Repository: open-mmlab/mmdeploy

Issue:
Could not export Python function call 'RoIAlignFunction'
While trying to export mask-rcnn-r50-c4, I'm getting following error:
```
Process Process-2:
Traceback (most recent call last):
  File "/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/pytorch2torchscript.py", line 54, in torch2torchscript
    trace(
  File "/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 356, in _wrap
    return self.call_function(func_name_, *args, **kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 326, in call_function
    return self.call_function_local(func_name, *args, **kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 275, in call_function_local
    return pipe_caller(*args, **kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py", line 107, in __call__
    ret = func(*args, **kwargs)
  File "/home/sort/mmdeploy/mmdeploy/apis/torch_jit/trace.py", line 121, in trace
    torch.jit.save(ts_model, output_path)
  File "/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/jit/_serialization.py", line 81, in save
    m.save(f, _extra_files=_extra_files)
  File "/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/jit/_script.py", line 693, in save
    return self._c.save(str(f), **kwargs)
RuntimeError: 
Could not export Python function call 'RoIAlignFunction'. Remove calls to Python functions before export. Did you forget to add @script or @script_method annotation? If this is a nn.ModuleList, add it to __constants__:
/home/sort/ved/mmcv/mmcv/ops/roi_align.py(215): forward
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1098): _slow_forward
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1110): _call_impl
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py(130): single_roi_extractor__forward
/home/sort/mmdeploy/mmdeploy/core/optimizers/function_marker.py(261): g
/home/sort/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py(371): wrapper
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1098): _slow_forward
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1110): _call_impl
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/mmdet/models/roi_heads/standard_roi_head.py(121): _bbox_forward
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/roi_heads/test_mixins.py(51): bbox_test_mixin__simple_test_bboxes
/home/sort/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py(371): wrapper
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/roi_heads/standard_roi_head.py(38): standard_roi_head__simple_test
/home/sort/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py(371): wrapper
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/detectors/two_stage.py(59): two_stage_detector__simple_test
/home/sort/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py(371): wrapper
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/detectors/base.py(26): __forward_impl
/home/sort/mmdeploy/mmdeploy/core/optimizers/function_marker.py(261): g
/home/sort/mmdeploy/mmdeploy/codebase/mmdet/models/detectors/base.py(70): base_detector__forward
/home/sort/mmdeploy/mmdeploy/core/rewriters/rewriter_utils.py(371): wrapper
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1098): _slow_forward
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/nn/modules/module.py(1110): _call_impl
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/jit/_trace.py(958): trace_module
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/site-packages/torch/jit/_trace.py(741): trace
/home/sort/mmdeploy/mmdeploy/apis/torch_jit/trace.py(97): trace
/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py(107): __call__
/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py(275): call_function_local
/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py(326): call_function
/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py(356): _wrap
/home/sort/mmdeploy/mmdeploy/apis/pytorch2torchscript.py(54): torch2torchscript
/home/sort/mmdeploy/mmdeploy/apis/core/pipeline_manager.py(107): __call__
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/process.py(108): run
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/process.py(315): _bootstrap
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/spawn.py(129): _main
/home/sort/miniconda3/envs/openmmlab/lib/python3.8/multiprocessing/spawn.py(116): spawn_main
<string>(1): <module>

2022-07-07 14:51:36,787 - mmdeploy - ERROR - `mmdeploy.apis.pytorch2torchscript.torch2torchscript` with Call id: 0 failed. exit.
```

The command I'm using is as follows:
```
python /home/sort/mmdeploy/tools/deploy.py \
    /home/sort/mmdeploy/configs/mmdet/detection/detection_torchscript.py \
    /home/sort/ved/sort/onion/mask_rcnn_r50_caffe_c4_onion_l1.py \
    /home/sort/ved/sort/onion/mask_rcnn_r50/latest.pth \
    /home/sort/ved/sort/onion/data_l1/val/images/0.jpg \
    --work-dir export_mask_rcnn_r50_onion_l1 \
    --show \
    --device cuda:1
```


Discussion:
@anshkumar Hi, you could post your env info here by running `python tools/check_env.py`.
I've reproduced the error above. Will fix it ASAP.

Resolving patch:
diff --git a/mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py b/mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py
--- a/mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py
+++ b/mmdeploy/codebase/mmdet/models/roi_heads/single_level_roi_extractor.py
@@ -127,6 +127,8 @@ def single_roi_extractor__forward(ctx,
     roi_feats = feats[0].new_zeros(rois.shape[0], self.out_channels, *out_size)
     if num_levels == 1:
         assert len(rois) > 0, 'The number of rois should be positive'
+        if backend == Backend.TORCHSCRIPT:
+            self.roi_layers[0].use_torchvision = True
         return self.roi_layers[0](feats[0], rois)
 
     target_lvls = self.map_roi_levels(rois, num_levels)


Test patch: