jwfromm commented on code in PR #16654:
URL: https://github.com/apache/tvm/pull/16654#discussion_r1506386423
##########
python/tvm/relax/frontend/nn/modules.py:
##########
@@ -270,7 +272,88 @@ def forward(self, x: Tensor) -> Tensor: # pylint:
disable=invalid-name
The output tensor for the conv2d layer.
"""
return op.conv2d(
- x, self.weight, self.bias, self.stride, self.padding,
self.dilation, self.groups
+ x,
+ self.weight,
+ self.bias,
+ self.stride,
+ self.padding,
+ self.dilation,
+ self.groups,
+ self.data_layout,
+ )
+
+
+class Conv3D(Module):
+ """
+ Module for conv3d layer.
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[List[int], int],
+ stride: Union[List[int], int] = 1,
+ padding: Union[List[int], int] = 0,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ dtype: Optional[str] = None,
+ data_layout: str = "NCDHW",
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.stride = stride
+ self.padding = padding
+ self.dilation = dilation
+ self.groups = groups
+ self.data_layout = data_layout
+
+ # Allow dynamic input channels.
+ if isinstance(self.in_channels, int):
+ in_channels = int(self.in_channels / self.groups)
+ else:
+ in_channels = tir.floordiv(self.in_channels, self.groups)
+
+ # Expand kernel size if given an integer.
+ if isinstance(kernel_size, int):
+ self.kernel_size = [kernel_size] * 3
+ else:
+ self.kernel_size = kernel_size
+
+ kernel_shape = [self.out_channels, self.in_channels] +
list(self.kernel_size)
+
+ self.weight = Parameter(kernel_shape, dtype)
+
+ if bias:
+ self.bias = Parameter((self.out_channels,), dtype)
+ else:
+ self.bias = None
+
+ def forward(self, x: Tensor) -> Tensor: # pylint: disable=invalid-name
+ """
+ Forward method for conv3d layer.
+
+ Parameters
+ ----------
+ x : Tensor
+ The input tensor.
+
+ Returns
+ -------
+ ret : Tensor
+ The output tensor for the conv2d layer.
Review Comment:
busted for copy pasting haha
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
