How to convert torchvision MaskRCNN model to TensorRT?

114 Views Asked by At

I've tried to use Torch-TensorRT to make the conversion, but I'm getting an error.

Here is a minimal example reproducing the problem:

model = torchvision.models.detection.maskrcnn_resnet50_fpn_v2(weights_backbone=None)
model.cuda()
model.eval()

# Define inputs
input_shape = (3, 224, 224)
inputs = [
    torch_tensorrt.Input(
        min_shape=[1, *input_shape],
        opt_shape=[batch_size, *input_shape],
        max_shape=[batch_size, *input_shape],
        dtype=input_dtype,
    )
]

# Compile to TensorRT
trt_ts_module = torch_tensorrt.compile(model, inputs=inputs, enabled_precisions={torch.half})

Environment:

OS                Ubuntu 22.04
Python            3.8.16
torch             2.0.1
torch-tensorrt    1.4.0
torchvision       0.15.2

I get this error:

Traceback (most recent call last):
...
    trt_ts_module = torch_tensorrt.compile(model, inputs=inputs, enabled_precisions={torch.half})
  File "/home/ml/venvs/venv38_20231029/lib/python3.8/site-packages/torch_tensorrt/_compile.py", line 133, in compile
    return torch_tensorrt.ts.compile(
  File "/home/ml/venvs/venv38_20231029/lib/python3.8/site-packages/torch_tensorrt/ts/_compiler.py", line 139, in compile
    compiled_cpp_mod = _C.compile_graph(module._c, _parse_compile_spec(spec))
RuntimeError: 
temporary: the only valid use of a module is looking up an attribute but found  = prim::SetAttr[name="_has_warned"](%self, %77)
:

How can I make this conversion?

0

There are 0 best solutions below