Export to ONNX Advice

johnorford · July 31, 2023, 12:58pm

I am trying to export some code to ONNX:

Atten4Vis/ConditionalDETR/blob/main/models/backbone.py#L77C10-L77C10


      
          self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)

Some layers (‘IntermediateLayer’) are outputting dictionaries:

github.com

pytorch/vision/blob/main/torchvision/models/_utils.py#L37-L42


      
          >>> new_m = torchvision.models._utils.IntermediateLayerGetter(m,
          >>>     {'layer1': 'feat1', 'layer3': 'feat2'})
          >>> out = new_m(torch.rand(1, 3, 224, 224))
          >>> print([(k, v.shape) for k, v in out.items()])
          >>>     [('feat1', torch.Size([1, 64, 56, 56])),
          >>>      ('feat2', torch.Size([1, 256, 14, 14]))]

…which ONNX export is complaining about (no aten::Dict support).

Any general advice on best next steps?

I will try replacing the dictionary with a list of tuples.

But I’d appreciate any other advice.