Mask-rcnn torchscript doesn't load in Android: Unknown builtin op: torchvision::_new_empty_tensor_op

[shehroz010]

I have trained a faster rcnn model on a custom dataset but, when I try to take trace of the model using jit.trace() it throws some errors. So I tried jit.script(). It works but when I try to load it into an android module it fails and throws and error
"could not find a similar ops to torchvision: :_new_empty_tensor_op .This op may not exist or may not be currently supported in TorchScript".

Below is the image of the jit.script() issue while loading it into android module

For the second issue when I use jit.trace() is throws an runtime error even I change the input to cuda type

Code:

example = torch.rand(1, 3, 224, 224)
examplex = example.to('cuda')
model_trace = torch.jit.trace(modelx, examplex)

Error

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Any help will be appreciated.

[fmassa]

cc @dreiss for loading custom ops in mobile

@shehroz010 for the second issue, maskrcnn models only support torch.jit.script, and not torch.jit.trace.

[shehroz010]

@fmassa does the above point is also true for faster rcnn model too? As I am using a pretrained faster rcnn model which is trained on custom a dataset.

[fmassa]

@shehroz010 yes, this is true for all R-CNN models in torchvision (Faster R-CNN, Mask R-CNN and Keypoint R-CNN), you should use torch.jit.script instead of torch.jit.trace.

[shehroz010]

Ok, I switched to jit.script() but in that I got another issue regarding "new_empty_tensor_op" which you have already seen. And thanks for clearing this as I can focus on jit.script() issue only. One more thing, is there will be any support going to be added to RCNN models for jit.trace() ?

[fmassa]

@shehroz010 we will not be supporting torch.jit.trace for Mask R-CNN, as it's less robust / flexible than torch.jit.script. If you want to export your model to ONNX you can, via torch.onnx.export.

The issue with new_empty_tensor_op is a known problem, and @dreiss is aware of it and will be working on a solution in the future.

[shehroz010]

@fmassa thankyou so much for the info.

[datumbox]

The new_empty_tensor operator has been removed by #3156