Convert fasterrcnn_resnet50_fpn to torchscript

I fine tuned a fasterrcnn_resnet50_fpn model and try to convert it to torchscript.

Load pre-trained model:

model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

After training:

in_size = 416
inp = torch.Tensor(np.random.uniform(0.0, 250.0, size=(1, 3, in_size, in_size)))
inp = inp.cuda()
with torch.no_grad():
    out = model(inp)
    model_trace = torch.jit.trace(model, inp)

I got this error:

RuntimeError: Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions

Thank you for any help.
I do not want use it on mobile devices.

You could try to use torch.jit.script instead, which should be more flexible.

Thank you. It works.
not sure about the differences?

This section of the docs explain the main differences.

1 Like

@ptrblck Thank you for your reply.
I can convert it to torchscript in script mode.
However I need to have it in traced mode. I am converting it by AWS neo which is based on Apache TVM.
So I need to solve this error.

Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions

Thank you for any clue or solution.

You are getting this error because your model will be giving an output which is not tensors, lists, tuples of tensors or dictionary of tensors.
This can be fixed by modifying the forward function of your model class. Just wrap the current output with torch.tensor($current_output$).