yongwww commented on code in PR #16654:
URL: https://github.com/apache/tvm/pull/16654#discussion_r1506342194
##########
src/relax/op/image/resize.cc:
##########
@@ -105,14 +105,26 @@ StructInfo InferStructInfoResize2D(const Call& call,
const BlockBuilder& ctx) {
InferLayoutOutput InferLayoutResize2d(const Call& call,
const Map<String, Array<String>>&
desired_layouts,
const VarLayoutMap& var_layout_map) {
- ICHECK(NoDesiredLayout(call, desired_layouts));
+ const auto& it = desired_layouts.find("relax.image.resize2d");
const auto* attrs = call->attrs.as<Resize2DAttrs>();
ICHECK(attrs) << "Invalid Call";
- LayoutDecision layout = GetLayoutDecision(var_layout_map, call->args[0]);
+ LayoutDecision data_layout;
ObjectPtr<Resize2DAttrs> new_attrs = make_object<Resize2DAttrs>(*attrs);
- new_attrs->layout = TransposeLike(attrs->layout, InitialLayout(4),
layout->layout).name();
- return InferLayoutOutput({layout, InitialNLayout(call->args[1])}, {layout},
Attrs(new_attrs));
+
+ if (it != desired_layouts.end()) {
Review Comment:
it would be good to add test cases for this change in `test_op_image.py`
##########
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:
typo: conv3d
##########
python/tvm/relax/frontend/nn/op.py:
##########
@@ -413,10 +417,84 @@ def conv2d(
strides=stride,
padding=padding,
dilation=dilation,
+ data_layout=data_layout,
groups=groups,
)
if bias is not None:
- conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, -1, 1, 1]))
+ if data_layout == "NCHW":
+ conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, -1, 1,
1]))
+ elif data_layout == "NHWC":
+ conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, 1, 1,
-1]))
+ else:
+ raise NotImplementedError(f"Dont know how to handle layout
{data_layout}.")
+
+ return wrap_nested(conv_out, name)
+
+
+def conv3d(
+ x: Tensor,
+ weight: Tensor,
+ bias: Optional[Tensor] = None,
+ stride: Optional[Union[int, Tuple]] = 1,
+ padding: Optional[Union[int, Tuple, str]] = 0,
+ dilation: Optional[Union[int, Tuple]] = 1,
+ groups: Optional[int] = 1,
+ data_layout: Optional[str] = "NCDHW",
+ name: str = "conv3d",
+) -> Tensor:
+ """Applies a 3D convolution over an input image composed of sevaral input
planes
+
+ Parameters
+ ----------
+ x : Tensor
+ Input tensor of shape [B, N, D, H, W]
+
+ weight : Tensor
+ Filters of shape [O, N/groups, kD, kH, kW]
+
+ bias : Optional[Tensor]
+ Optional bias tensor of shape [O].
+
+ stride : Optional[Union[int, Tuple]]
+ The stride of the convolving kernel. Can be a single number
+ or tuple of (sD, sH, sW).
+
+ padding : Optional[[Union[int, Tuple]]]
+ Implicit paddings on both sides of the input.
+
+ dilation : Optional[Union[int, Tuple]]
+ The spacing between kernel elements. Can be a single number of tuple
(dD, dH, dW).
+
+ groups : Optional[int]
+ Split input into a number of groups.
+
+ data_layout : Optional[str]
+ Optional layout of the input and output data.
+
+ name : str
+ Name hint.
+
+ Returns
+ -------
+ result : Tensor
+ The computed result with shape [B, O, oD, oH, oW].
Review Comment:
how about we use [B, oC, oD, oH, oW]? I initially mistook the 'O' for the '0'
--
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]
