Cast Operation in ONNX

1

Hi, I am trying to export the segmentation some segementation of the segmentation model, newly added to torchvision. They export sucessfully after applying this workaround. However the ONNX graph generates several cast operations as seen in the picture.

image
image1919×303 92.9 KB

It says its coming from the ASPPPooling layer but that layer is defined as:

image
image719×168 23.6 KB

And I dont think any of these layers do any casting. Any ideas where this Cast[to=1] comes from or how to avoid it? Any help will be greatly appreciated!! :slight_smile:

P.S I am using pytorch and torchvision both compiled from source

2

I guess the cast is related to this line of code in interpolate, which was added in this PR to fix tracing.

PS: I could run the code without modifying anything or building from source:

class AsppPooling(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer):
        super(AsppPooling, self).__init__()
        self.gap = nn.Sequential(nn.AdaptiveAvgPool2d(1),
                                 nn.Conv2d(in_channels, out_channels, 1, bias=False),
                                 norm_layer(out_channels),
                                 nn.ReLU(True))

    def forward(self, x):
        _, _, h, w = x.size()
        pool = self.gap(x)
        return F.interpolate(pool, (h,w))


model = AsppPooling(3, 3, nn.BatchNorm2d)
torch.onnx.export(model, torch.randn(2, 3, 2, 2), 'tmp.onnx', verbose=True)
3

Thank you so much! @ptrblck. What should I expect if I were to remove the .floor operation?. Would the results be the same if I wanted to use that model for segmentation?

4

I’m not sure, but I would assume you would get wrong results.
Also, I don’t think you should be too worried about casting a single value, as this should not change the time a lot.
That being said, these are just guesses, as I’m not very familiar with ONNX, so @eellison might have more insight as the author of the PR.

1 Like
5

I am not worried because of speed. I am trying to export some ONNX models to OpenCV as part of my GSoC project and the OpenCV ONNX importer does not support the Cast operation from ONNX yet. I am trying to see if I can get away with not implementing it :sweat_smile: