Is it possible to convert this model to ONNX format?

Hello, I used your model to train it on custom dataset, everything works fine, but I want to convert it to ONNX format, so it could be used with ML.NET. To convert I used the following code:

from ssd.config import cfg from ssd.modeling.detector import build_detection_model import torch cfg.merge_from_file(‘configs/vgg_custom_ssd512.yaml’) model = build_detection_model(cfg) state_dict = torch.load(‘model_final.pth’, map_location=lambda storage, loc: storage)[‘model’] model.load_state_dict(state_dict) dummy_input = torch.randn(1, 3, 512, 512, device=‘cpu’) torch.onnx.export(model, dummy_input, ‘ssdpytorch.onnx’, verbose = True)

I got the following error: RuntimeError: Only tuples, lists and Variables supported as JIT inputs, but got Container

Though Container isn’t used either in model or in state_dict, this data type is used while converting model to ONNX. From your point of view, is it solvable? Exactly the same error with vgg_ssd512_voc0712, so this error is easily reproducible.

[lufficc]

Since the return values are not tuples or lists, it cannot be directly exported. But a model wrapper will solve this(You can ignore some warnings).

from ssd.config import cfg
from ssd.modeling.detector import build_detection_model
import torch
from torch import nn


class ONNXExportableModel(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, *args, **kwargs):
        detections = self.model(*args, **kwargs)
        detections = [(det['boxes'], det['labels'], det['scores']) for det in detections]
        return detections


if __name__ == '__main__':
    cfg.merge_from_file('configs/vgg_ssd300_voc0712.yaml')
    model = build_detection_model(cfg)
    state_dict = torch.load('weights/vgg_ssd300_voc0712.pth', map_location=lambda storage, loc: storage)['model']
    model.load_state_dict(state_dict)
    dummy_input = torch.randn(1, 3, 300, 300, device='cpu')
    torch.onnx.export(ONNXExportableModel(model), dummy_input, 'ssdpytorch.onnx', verbose=True)

[pawopawo]

I converted the onnx file as shown in the picture. Does anyone know what is wrong? Hope for help