Hi @apivovarov @kevinthesun @anijain2305:
There is a conv2d_transpose_legalize function that forced convert the conv2d_transpose input of NHWC layout to NCHW by called the relay.transpose before and after the relay.nn.conv2d_transpose in python/tvm/topi/nn/conv2d_transpose.py.
we have a problem: Because our target only support nhwc layout, so we convert the nchw layout network’s operators to nhwc through ConvertLayout pass and Specify the desired_layouts(transform.ConvertLayout(desired_layouts)), but after convert succuss and called the
if (targets.size() == 1) {
pass_seqs.push_back(transform::Legalize());
}
the layout convert to nchw again, which is not our expected.
%497 = nn.relu(%496) /* ty=Tensor[(1000, 7, 7, 2048), float16] */;
%498 = transpose(%497, axes=[0, 3, 1, 2]) /* ty=Tensor[(1000, 2048, 7, 7), float16] */;
%499 = nn.conv2d_transpose(%498, meta[relay.Constant][282] /* ty=Tensor[(2048, 256, 2, 2), float16] */ /* ty=Tensor[(2048, 256, 2, 2), float16] */, channels=256, kernel_size=[2, 2], strides=[2, 2], padding=[0, 0, 0, 0]) /* ty=Tensor[(1000, 256, 14, 14), float16] */;
%500 = transpose(%499, axes=[0, 2, 3, 1]) /* ty=Tensor[(1000, 14, 14, 256), float16] */;
So, I think the forced conversion in conv2d_transpose_legalize is not suitable to those targets need supported nhwc layout and called ConvertLayout to attempt to adjust the layout.
we should make a change in conv2d_transpose_legalize that
# Set new attrs for conv2d_transpose.
new_attrs = {k: attrs[k] for k in attrs.keys()}
new_attrs["data_layout"] = "NCHW"
# layout of kernel should be IOHW, but kernel_layout should be swapped - OIHW
new_attrs["kernel_layout"] = "OIHW"
# Convert data to NCHW.
data = relay.transpose(data, axes=(0, 3, 1, 2))
deconv = relay.nn.conv2d_transpose(data, kernel, **new_attrs)
# Convert back to original NHWC layout.
out = relay.transpose(deconv, axes=(0, 2, 3, 1))
return out
to
# # Set new attrs for conv2d_transpose.
new_attrs = {k: attrs[k] for k in attrs.keys()}
new_attrs["kernel_layout"] = "OIHW"
out = relay.nn.conv2d_transpose(data, kernel, **new_attrs)
return out
and move the check for NCHW from relay type infer into op strategy for corresponding targets which don’t have NHWC implementation.(already exists cuda/x86/hls assert layout == "NCHW", "only support nchw for now")
What do you think? Can you give some suggestions to support the nhwc layout conversion of conv2d_transpose?