[WIP] Support generative bundles and pathology hovernet bundles with TensorRT

models/brats_mri_axial_slices_generative_diffusion/configs/inference_trt.json

@@ -0,0 +1,13 @@
+{
+    "imports": [
+        "$from datetime import datetime",
+        "$from pathlib import Path",
+        "$from PIL import Image",
+        "$from scripts.utils import visualize_2d_image",
+        "$import torch_tensorrt"
+    ],
+    "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder_trt.ts'",
+    "load_autoencoder": "$torch.jit.load(@load_autoencoder_path)",
+    "load_diffusion_path": "$@bundle_root + '/models/model_trt.ts'",
+    "load_diffusion": "$torch.jit.load(@load_diffusion_path)"
+}

models/brats_mri_axial_slices_generative_diffusion/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20230507.json",
-    "version": "1.0.6",
+    "version": "1.0.7",
     "changelog": {
+        "1.0.7": "add support for TensorRT conversion and inference",
         "1.0.6": "update with new lr scheduler api in inference",
         "1.0.5": "fix the wrong GPU index issue of multi-node",
         "1.0.4": "update with new lr scheduler api",

models/brats_mri_axial_slices_generative_diffusion/docs/README.md

@@ -143,6 +143,26 @@ The following code generates a synthetic image from a random sampled noise.
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
+#### Export checkpoint to TensorRT based models with fp32 or fp16 precision:
+
+The Autoencoder can be exported into a TensorRT model.
+
+```
+python -m monai.bundle trt_export --net_id autoencoder_def --filepath models/model_autoencoder_trt.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json --use_trace "True"  --input_shape "[1, 1, 240, 240]" --precision <fp32/fp16> --converter_kwargs "{'truncate_long_and_double': True}"
+```
+
+The Diffusion model can be exported into a TensorRT model.
+
+```
+ 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"  --input_shape "[[1, 1, 64, 64], [1,]]" --converter_kwargs "{'truncate_long_and_double': True}"
+```
+
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
+
 # References
 [1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf
 

models/brats_mri_generative_diffusion/configs/inference_trt.json

@@ -0,0 +1,11 @@
+{
+    "imports": [
+        "$from datetime import datetime",
+        "$from pathlib import Path",
+        "$import torch_tensorrt"
+    ],
+    "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder_trt.ts'",
+    "load_autoencoder": "$torch.jit.load(@load_autoencoder_path)",
+    "load_diffusion_path": "$@bundle_root + '/models/model_trt.ts'",
+    "load_diffusion": "$torch.jit.load(@load_diffusion_path)"
+}

models/brats_mri_generative_diffusion/configs/metadata.json

@@ -1,7 +1,8 @@
 {
     "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_generator_ldm_20230507.json",
-    "version": "1.0.6",
+    "version": "1.0.7",
     "changelog": {
+        "1.0.7": "add support for TensorRT conversion and inference",
         "1.0.6": "update with new lr scheduler api in inference",
         "1.0.5": "fix the wrong GPU index issue of multi-node",
         "1.0.4": "update with new lr scheduler api",

models/brats_mri_generative_diffusion/docs/README.md

@@ -151,6 +151,26 @@ The Autoencoder can be exported into a TorchScript file.
 python -m monai.bundle ckpt_export autoencoder_def --filepath models/model_autoencoder.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json
 ```
 
+#### Export checkpoint to TensorRT based models with fp32 or fp16 precision:
+
+The Autoencoder can be exported into a TensorRT model.
+
+```
+python -m monai.bundle trt_export --net_id autoencoder_def --filepath models/model_autoencoder_trt.ts --ckpt_file models/model_autoencoder.pt --meta_file configs/metadata.json --config_file configs/inference.json --use_trace "True"  --input_shape "[1, 1, 112, 128, 80]" --precision <fp32/fp16> --converter_kwargs "{'truncate_long_and_double': True}"
+```
+
+The Diffusion model can be exported into a TensorRT model.
+
+```
+ 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"  --input_shape "[[1, 8, 36, 44, 28], [1,]]" --converter_kwargs "{'truncate_long_and_double': True}"
+```
+
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
+
 # References
 [1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022. https://openaccess.thecvf.com/content/CVPR2022/papers/Rombach_High-Resolution_Image_Synthesis_With_Latent_Diffusion_Models_CVPR_2022_paper.pdf
 

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",

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
+}

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