ConvTranspose has the wrong bias shape

[yuanyao-nv]

In some exported model with ConvTranspose I'm seeing the wrong shape for bias:


where B should have shape (128) instead of (256) according to the ONNX spec.

Looking at the onnxscript code, the problem seems to be that since
_aten_convolution_onnx() is used for both Conv and ConvTranspose, the logic for filling in bias should check the 0th dimension of weight if it's Conv and the 1st dimension of weight if it's ConvTranspose:

    if bias is None:
        weight_dim_0 = op.Shape(weight, start=0, end=1)
        bias_shape = op.Expand(weight_dim_0, op.Constant(value_ints=[1]))
        zero = op.CastLike(0.0, input)
        bias = op.Expand(zero, bias_shape)

I think the value of groups is also needed since the shape of weight is (M x C/group x kH x kW) for Conv and (C x M/group x kH x kW) for ConvTranspose?

[justinchuby]

@fatcat-z would you be able to help take a look? Thanks!

[gramalingam]

Hi, wondering about this issue ... any progress? Thanks!

[fatcat-z]

Hi, wondering about this issue ... any progress? Thanks!

Just finished another issue and start working on it now.

[fatcat-z]

@yuanyao-nv ,
The function _aten_convolution_onnx was designed to serve those convolution ops, like aten::convolution, aten::conv1d, aten::2d, aten::3d. It is not used by any conv_transpose op(conv_transpose1d/conv_transpose2d/conv_transpose3d).

Based on this, the logic of filling bias will be as same as Conv op. We call op.ConvTranspose op in _aten_convolution_onnx() because the parameter 'transposed' of aten::convolution op. If transposed it True, then we can call op.ConvTranspose() op to get the same results as aten::convolution op. And it has been covered by our tests and proved it is true.

    if transposed:
        result = op.ConvTranspose(
            input,
            weight,
            bias,
            strides=strides,
            pads=pads,
            group=groups,
            dilations=dilations,
            output_padding=output_padding,
        )
    else:
        result = op.Conv(
            input,
            weight,
            bias,
            strides=strides,
            pads=pads,
            group=groups,
            dilations=dilations,
        )

Could you please share a code snippet in this issue so that I can repro the problem you met and dig into the root cause?

[yuanyao-nv]

@fatcat-z Sorry I missed your message.
Here's an example export script:

import torch
from monai.networks.nets import DynUNet

def UNetDataloader(dimensions=3, height=128, width=128, depth=128):
    if dimensions == 2:
        yield torch.randn(1,1,height,width).cuda()
    else:
        yield torch.randn(1,1,height,width,depth).cuda()

model = lambda : DynUNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=3,
            kernel_size=[[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
            strides=[[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
            upsample_kernel_size=[[2, 2, 2], [2, 2, 2], [2, 2, 2]],
            norm_name="instance",
            deep_supervision=True,
            deep_supr_num=1
        )
data = next(UNetDataloader())
model = model().eval().to('cuda')
export_output = torch.onnx.dynamo_export(
    model,
    data,
)
export_output.save('Clara_DynUNet_dynamo.onnx')

Some relevant env info:
monai 1.3.0
onnx 1.16.0
onnxscript 0.1.0.dev20240510
torch 2.4.0.dev20240509+cu121

Let me know if you need more info. Thanks!

[fatcat-z]

@fatcat-z Sorry I missed your message. Here's an example export script:

import torch
from monai.networks.nets import DynUNet

def UNetDataloader(dimensions=3, height=128, width=128, depth=128):
    if dimensions == 2:
        yield torch.randn(1,1,height,width).cuda()
    else:
        yield torch.randn(1,1,height,width,depth).cuda()

model = lambda : DynUNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=3,
            kernel_size=[[3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]],
            strides=[[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2]],
            upsample_kernel_size=[[2, 2, 2], [2, 2, 2], [2, 2, 2]],
            norm_name="instance",
            deep_supervision=True,
            deep_supr_num=1
        )
data = next(UNetDataloader())
model = model().eval().to('cuda')
export_output = torch.onnx.dynamo_export(
    model,
    data,
)
export_output.save('Clara_DynUNet_dynamo.onnx')

Some relevant env info: monai 1.3.0 onnx 1.16.0 onnxscript 0.1.0.dev20240510 torch 2.4.0.dev20240509+cu121

Let me know if you need more info. Thanks!

I couldn't find such problematic ConvTranspose node in this example ONNX file. And this ONNX file exactly shows my above comments about the usage of _aten_convolution_onnx() method.

[yuanyao-nv]

I think the problem still exists. I've uploaded the model here: https://drive.google.com/drive/folders/1WcUNY-O1mVWRC26EP4h5k94ptLcnUuzs?usp=sharing

It's easier to see the problem in the ConvTranspose nodes in the inlined version:

According to the ONNX spec, The size of B should be the number of feature maps 128, not the number of channels 256, right?

[fatcat-z]

I think the problem still exists. I've uploaded the model here: https://drive.google.com/drive/folders/1WcUNY-O1mVWRC26EP4h5k94ptLcnUuzs?usp=sharing

It's easier to see the problem in the ConvTranspose nodes in the inlined version:

According to the ONNX spec, The size of B should be the number of feature maps 128, not the number of channels 256, right?

May I know how you get the Clara_DynUNet_dynamo_inlined.onnx file?

[yuanyao-nv]

I just used the ONNX inliner onnx.inliner.inline_local_functions()