TensorRT segmentation probabilities scores

[user41pp]

Describe the feature

MMSegmentation TensorRT deployment: output float probabilities/scores instead of int class labels.

In segmentor.h we have
int* mask; ///< segmentation mask of the input image, in which mask[i * width + j] indicates the label id of pixel at (i, j)

Request: change mask to be a 3d matrix of probabilities of shape (width, height, num_classes) with float probabilities for all classes.

This would require modifications in the export pipeline as well.

Motivation

• Variable thresholds for different classes at runtime
• Decision making based on all scores, not just the top one, if available

Related resources

https://github.com/open-mmlab/mmdeploy/blob/master/csrc/mmdeploy/apis/c/mmdeploy/segmentor.h

[RunningLeon]

@user41pp Hi, exporting to onnx with score map(without argmax) is OK and easy to implement, but setting thresholds for each class at runtime is a little confused. Maybe you could provide sample code or PR to show your thoughts.

[user41pp]

@user41pp Hi, exporting to onnx with score map(without argmax) is OK and easy to implement, but setting thresholds for each class at runtime is a little confused. Maybe you could provide sample code or PR to show your thoughts.

Hi, sorry for the late reply. In the simplest form, I am asking for exactly that - adding the option (or making it the default) to remove the argmax layer and instead directly output the softmax values.

Here is a screenshot of the last layers of a UNET model (configs and commands listed further below) with that change:

This would solve our problem already and if that's easy, then this is the way to go. The "dynamic threshold at runtime" part does not have to be done by mmdeploy.

However, outputting the softmax layer would increase the output tensor by the number of classes the model has (e.g. current output shape is 1024 x 512, and if we would output the softmax values, output shape would be 1024 x 512 x num_classes). Because of the potentially huge amount of classes maybe it makes sense to not output all the scores but instead only the topk scores and correspondig class indices. For the most common case k=1 we would output the top softmax score along with its class index. The returned tensor then would be a tuple containing (argmax, softmax[argmax]).

MMdeploy config:
https://github.com/open-mmlab/mmdeploy/blob/master/configs/mmseg/segmentation_tensorrt-fp16_dynamic-512x1024-2048x2048.py

MMseg config:
https://github.com/open-mmlab/mmsegmentation/blob/master/configs/unet/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes.py

MMseg model checkpoint:
https://download.openmmlab.com/mmsegmentation/v0.5/unet/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes_20211210_145204-6860854e.pth

Deploy command:

python %MMDEPLOY_DIR%/tools/deploy.py^  
 %MMDEPLOY_DIR%/configs/mmseg/segmentation_tensorrt-fp16_dynamic-512x1024-2048x2048.py^  
 %MMSEG_DIR%/configs/unet/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes.py^  
 %CHECKPOINT_DIR%/fcn_unet_s5-d16_4x4_512x1024_160k_cityscapes_20211210_145204-6860854e.pth^  
 %MMSEG_DIR%/demo/demo.jpg^  
 --work-dir %WORK_DIR%^  
 --device cuda:0^  
 --dump-info

[RunningLeon]

@user41pp Hi, if you only need to operate on onnx, it's a customized exporting. You could refer to here for how to get a sub-graph onnx file.

[user41pp]

@user41pp Hi, if you only need to operate on onnx, it's a customized exporting. You could refer to here for how to get a sub-graph onnx file.

We use TensorRT like shown here https://github.com/open-mmlab/mmdeploy/blob/master/demo/csrc/image_segmentation.cpp

Specifically, we use an export config similar to the TensorRT export config linked above (segmentation_tensorrt-fp16_dynamic-512x1024-2048x2048.py; the end2end.onnx file is an intermediate artifact which is used while creating the TensorRT engine)

[RunningLeon]

@user41pp Hi, do you have time to PR us this feature? Seems not so complicated. It includes exporting onnx without argmax and add argmax as post-process in mmdeploy-sdk.

[yuanyuangoo]

Hi, Is this problem solved?

[RunningLeon]

Hi, Is this problem solved?

@yuanyuangoo Hi, have given the solution, not sure if it's solved the problem.

[kacifer999]

@RunningLeon
Exporting onnx with out argmax is quit easy. Could you point out where to add armax in post process.
I have tried the first part. The second part raise error as:

[2022-11-21 18:20:59.885] [mmdeploy] [info] [inference.cpp:50] ["img"] <- ["img"]
[2022-11-21 18:20:59.885] [mmdeploy] [info] [inference.cpp:61] ["post_output"] -> ["mask"]
[2022-11-21 18:21:00.877] [mmdeploy] [error] [segment.cpp:35] unsupported output tensor, shape: [1, 2, 512, 512]

Which is a check you guys implement here:
https://github.com/open-mmlab/mmdeploy/blob/99040d5655aff9b39f331db925dce0c065a8176a/csrc/mmdeploy/codebase/mmseg/segment.cpp

[RunningLeon]

@kacifer999 hi, maybe this #1379 is what you want.

[github-actions]

This issue is closed because it has been stale for 5 days. Please open a new issue if you have similar issues or you have any new updates now.