Generating post-processed prediction masks with C++ TorchScript

[SammyProgrammer]

Hi,

I'm using TorchScript tracing to generate a Mask R-CNN model that I'd like to make use of in a C++ application. It would be wonderful if I could access the post-processed prediction masks that Detectron2 is capable of generating in my application, but I'm not sure what the best way to do that is so I'm looking for advice.

As per the deployment docs:

The converted models do not contain post-processing operations that transform raw layer outputs into formatted predictions. For example, the C++ examples only produce raw outputs (28x28 masks) from the final layers that are not post-processed, because in actual deployment, an application often needs its custom lightweight post-processing, so this step is left for users.

If I use the example deployment code:

./export_model.py --config-file ../../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \
    --output ./output --export-method tracing --format torchscript \
    MODEL.WEIGHTS detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl \
    MODEL.DEVICE cuda

./build/torchscript_traced_mask_rcnn output/model.ts input.jpg tracing

I get the expected output which includes the bounding box and the 28x28 mask:

Number of output tensors: 5
bbox: CUDAFloatType [1, 4]
scores: CUDAFloatType [1]
pred_classes: CUDALongType [1]
pred_masks: CUDAFloatType [1, 1, 28, 28]
 246.0907   14.2438  416.8796  189.9806

I would like to extend this to include the full size post-processed prediction masks. I was hoping it might be as simple as setting do_postprocess=True at this line:

detectron2/tools/deploy/export_model.py

Line 86 in 1ba62f1

inst = model.inference(inputs, do_postprocess=False)[0]

This generates a model.ts file, but it errors when running the C++ torchscript_traced_mask_rcnn program. Note that I set CUDA_LAUNCH_BLOCKING=1 to try and get a more complete trace:

/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:142: operator(): block: [28,0,0], thread: [96,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:142: operator(): block: [28,0,0], thread: [97,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:142: operator(): block: [28,0,0], thread: [98,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:142: operator(): block: [28,0,0], thread: [99,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
/pytorch/aten/src/ATen/native/cuda/IndexKernel.cu:142: operator(): block: [28,0,0], thread: [100,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
terminate called after throwing an instance of 'std::runtime_error'
  what():  The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/detectron2/export/flatten.py", line 101, in forward
    inds0 = torch.to(inds, dtype=4, layout=0, device=torch.device("cuda:0"), pin_memory=None, non_blocking=False, copy=False, memory_format=None)
    _77 = annotate(List[Optional[Tensor]], [inds0])
    masks0 = torch.index(_76, _77)
             ~~~~~~~~~~~ <--- HERE
    inds1 = torch.to(inds0, dtype=4, layout=0, device=torch.device("cuda:0"), pin_memory=None, non_blocking=False, copy=False, memory_format=None)
    _78 = annotate(List[Optional[Tensor]], [inds1])

Traceback of TorchScript, original code (most recent call last):
detectron2/detectron2/layers/mask_ops.py(134): paste_masks_in_image
detectron2/detectron2/utils/memory.py(70): wrapped
detectron2/detectron2/modeling/postprocessing.py(68): detector_postprocess
detectron2/detectron2/modeling/meta_arch/rcnn.py(246): _postprocess
detectron2/detectron2/modeling/meta_arch/rcnn.py(219): inference
export_model.py(90): inference
detectron2/detectron2/export/flatten.py(249): forward
detectron_venv/lib64/python3.6/site-packages/torch/nn/modules/module.py(860): _slow_forward
detectron_venv/lib64/python3.6/site-packages/torch/nn/modules/module.py(887): _call_impl
detectron_venv/lib64/python3.6/site-packages/torch/jit/_trace.py(957): trace_module
detectron_venv/lib64/python3.6/site-packages/torch/jit/_trace.py(742): trace
export_model.py(107): export_tracing
export_model.py(185): <module>
RuntimeError: CUDA error: device-side assert triggered

Abort

I'm not sure what the best way is to proceed from here, so any advice would be very welcome.

Many thanks,
Sammy

[ppwwyyxx]

This is not yet supported.

The easiest way to get there is for you to call cv2.resize after getting 28x28 mask to resize it to the size of the box

[SammyProgrammer]

Thanks for the info! I'll use OpenCV as you suggest.