add suppor t of trt inference of the pathology nuclei segmentation cl…

…assification bundle Signed-off-by: binliu <[email protected]>
Project-MONAI · Aug 20, 2023 · 934e839 · 934e839
1 parent 1ac569b
commit 934e839
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 1 deletion.
diff --git a/models/pathology_nuclei_segmentation_classification/configs/inference.json b/models/pathology_nuclei_segmentation_classification/configs/inference.json
@@ -12,13 +12,16 @@
     "hovernet_mode": "fast",
     "patch_size": 256,
     "out_size": 164,
+    "non_trt_inference": true,
+    "use_list_output": "$not @non_trt_inference",
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
     "network_def": {
         "_target_": "HoVerNet",
         "mode": "@hovernet_mode",
         "adapt_standard_resnet": true,
         "in_channels": 3,
-        "out_classes": 5
+        "out_classes": 5,
+        "use_list_output": "@use_list_output"
     },
     "network": "$@network_def.to(@device)",
     "preprocessing": {
@@ -77,6 +80,16 @@
     "postprocessing": {
         "_target_": "Compose",
         "transforms": [
+            {
+                "_target_": "ExtendSubKeysd",
+                "_disabled_": "@non_trt_inference",
+                "keys": "pred",
+                "map_names": [
+                    "horizontal_vertical",
+                    "nucleus_prediction",
+                    "type_prediction"
+                ]
+            },
             {
                 "_target_": "FlattenSubKeysd",
                 "keys": "pred",

diff --git a/models/pathology_nuclei_segmentation_classification/configs/inference_trt.json b/models/pathology_nuclei_segmentation_classification/configs/inference_trt.json
@@ -0,0 +1,11 @@
+{
+    "imports": [
+        "$import glob",
+        "$import os",
+        "$import torch_tensorrt"
+    ],
+    "non_trt_inference": false,
+    "handlers#0#_disabled_": true,
+    "network_def": "$torch.jit.load(@bundle_root + '/models/model_trt.ts')",
+    "evaluator#amp": false
+}
diff --git a/models/pathology_nuclei_segmentation_classification/docs/README.md b/models/pathology_nuclei_segmentation_classification/docs/README.md
@@ -135,6 +135,18 @@ python -m monai.bundle run --config_file "['configs/train.json','configs/evaluat
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
+#### Export checkpoint to TensorRT based models with fp32 or fp16 precision:
+
+```
+python -m monai.bundle trt_export --net_id network_def --filepath models/model_trt.ts --ckpt_file models/model.pt --meta_file configs/metadata.json --config_file configs/inference.json --precision <fp32/fp16> --use_trace "True" --use_onnx "True" --dynamic_batchsize "[1, 16, 16]" --onnx_output_names "['output_0', 'output_1', 'output_2']" --use_list_output "True"
+```
+
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
+
 # References
 [1] Simon Graham, Quoc Dang Vu, Shan E Ahmed Raza, Ayesha Azam, Yee Wah Tsang, Jin Tae Kwak, Nasir Rajpoot, Hover-Net: Simultaneous segmentation and classification of nuclei in multi-tissue histology images, Medical Image Analysis, 2019 https://doi.org/10.1016/j.media.2019.101563