Implementation of TRT wrapping via inference.json

models/brats_mri_axial_slices_generative_diffusion/configs/inference.json

@@ -77,8 +77,7 @@
     "load_autoencoder": "$@autoencoder_def.load_old_state_dict(torch.load(@load_autoencoder_path))",
     "autoencoder": "$@autoencoder_def.to(@device)",
     "load_diffusion_path": "$@model_dir + '/model.pt'",
-    "load_diffusion": "$@network_def.load_old_state_dict(torch.load(@load_diffusion_path))",
-    "diffusion": "$@network_def.to(@device)",
+    "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))",
     "noise_scheduler": {
         "_target_": "monai.networks.schedulers.ddim.DDIMScheduler",
         "_requires_": [

models/brats_mri_axial_slices_generative_diffusion/docs/README.md

@@ -143,6 +143,12 @@ The following code generates a synthetic image from a random sampled noise.
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
+#### 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.json

@@ -64,11 +64,10 @@
         "num_res_blocks": 2
     },
     "load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'",
-    "load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))",
+    "load_autoencoder": "$@autoencoder_def.to(@device).load_state_dict(torch.load(@load_autoencoder_path, weights_only=True))",
     "autoencoder": "$@autoencoder_def.to(@device)",
     "load_diffusion_path": "$@model_dir + '/model.pt'",
-    "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))",
-    "diffusion": "$@network_def.to(@device)",
+    "load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path, weights_only=True))",
     "noise_scheduler": {
         "_target_": "generative.networks.schedulers.DDIMScheduler",
         "_requires_": [

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.9",
+    "version": "1.0.10",
     "changelog": {
+        "1.0.10": "update to use monai 1.3.3rc1 + trt_wrap",
         "1.0.9": "update to use monai 1.3.1",
         "1.0.8": "update run section",
         "1.0.7": "update with EnsureChannelFirstd",
@@ -13,7 +14,7 @@
         "1.0.1": "update dependency, update trained model weights",
         "1.0.0": "Initial release"
     },
-    "monai_version": "1.3.1",
+    "monai_version": "1.3.3rc1",
     "pytorch_version": "2.2.2",
     "numpy_version": "1.24.4",
     "optional_packages_version": {

models/brats_mri_generative_diffusion/docs/README.md

@@ -143,14 +143,13 @@ 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 TorchScript file
-
-The Autoencoder can be exported into a TorchScript file.
+#### Execute inference with the TensorRT model:
 
 ```
-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
+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_classification/configs/inference.json

@@ -6,6 +6,7 @@
         "$import os"
     ],
     "bundle_root": ".",
+    "checkpoint": "$@bundle_root + '/models/model.pt'",
     "output_dir": "$@bundle_root + '/eval'",
     "dataset_dir": "/workspace/data/CoNSePNuclei",
     "images": "$list(sorted(glob.glob(@dataset_dir + '/Test/Images/*.png')))[:1]",
@@ -88,7 +89,7 @@
     "handlers": [
         {
             "_target_": "CheckpointLoader",
-            "load_path": "$@bundle_root + '/models/model.pt'",
+            "load_path": "$@checkpoint",
             "load_dict": {
                 "model": "@network"
             }

models/pathology_nuclei_classification/configs/inference_trt.json

@@ -1,12 +1,6 @@
 {
-    "imports": [
-        "$import glob",
-        "$import os",
-        "$import pathlib",
-        "$import json",
-        "$import torch_tensorrt"
+    "+imports": [
+        "$from monai.networks import trt_compile"
     ],
-    "handlers#0#_disabled_": true,
-    "network_def": "$torch.jit.load(@bundle_root + '/models/model_trt.ts')",
-    "evaluator#amp": false
+    "network": "$trt_compile(@network_def.to(@device), @checkpoint)"
 }

models/pathology_nuclei_classification/docs/README.md

@@ -207,12 +207,6 @@ torchrun --standalone --nnodes=1 --nproc_per_node=2 -m monai.bundle run --config
 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>
-```
-
 #### Execute inference with the TensorRT model:
 
 ```

models/pathology_nuclei_segmentation_classification/configs/inference.json

@@ -74,17 +74,18 @@
         "progress": true,
         "extra_input_padding": "$((@patch_size - @out_size) // 2,) * 4"
     },
+    "sub_keys": [
+        "horizontal_vertical",
+        "nucleus_prediction",
+        "type_prediction"
+    ],
     "postprocessing": {
         "_target_": "Compose",
         "transforms": [
             {
                 "_target_": "FlattenSubKeysd",
                 "keys": "pred",
-                "sub_keys": [
-                    "horizontal_vertical",
-                    "nucleus_prediction",
-                    "type_prediction"
-                ],
+                "sub_keys": "$@sub_keys",
                 "delete_keys": true
             },
             {

models/pathology_nuclei_segmentation_classification/configs/inference_trt.json

@@ -0,0 +1,10 @@
+{
+    "+imports": [
+        "$from monai.networks import trt_compile"
+    ],
+    "trt_args": {
+        "output_names": "$@sub_keys",
+        "dynamic_batchsize": "$[1, @sw_batch_size, @sw_batch_size]"
+    },
+    "network": "$trt_compile(@network_def.to(@device), @bundle_root + '/models/model.pt', args=@trt_args)"
+}

models/pathology_nuclei_segmentation_classification/docs/README.md

@@ -135,6 +135,12 @@ python -m monai.bundle run --config_file "['configs/train.json','configs/evaluat
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
+#### 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
 

models/swin_unetr_btcv_segmentation/configs/inference.json

@@ -5,6 +5,7 @@
         "$import os"
     ],
     "bundle_root": ".",
+    "checkpoint": "$@bundle_root + '/models/model.pt'",
     "image_key": "image",
     "output_dir": "$@bundle_root + '/eval'",
     "output_ext": ".nii.gz",
@@ -139,7 +140,7 @@
     },
     "checkpointloader": {
         "_target_": "CheckpointLoader",
-        "load_path": "$@bundle_root + '/models/model.pt'",
+        "load_path": "$@checkpoint",
         "load_dict": {
             "model": "@network"
         }

models/swin_unetr_btcv_segmentation/configs/inference_trt.json

@@ -0,0 +1,9 @@
+{
+    "+imports": [
+        "$from monai.networks import trt_compile"
+    ],
+    "trt_args": {
+        "dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]"
+    },
+    "network": "$trt_compile(@network_def.to(@device), @checkpoint, args=@trt_args)"
+}

models/swin_unetr_btcv_segmentation/docs/README.md

@@ -108,9 +108,12 @@ 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 TorchScript file:
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
 
-TorchScript conversion is currently not supported.
 
 # References
 [1] Hatamizadeh, Ali, et al. "Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images." arXiv preprint arXiv:2201.01266 (2022). https://arxiv.org/abs/2201.01266.

models/vista2d/configs/inference.json

@@ -11,7 +11,7 @@
         256,
         256
     ],
-    "input_dict": "${'image': '/home/venn/Desktop/data/medical/cellpose_dataset/test/001_img.png'}",
+    "input_dict": "${'image': '/cellpose_dataset/test/001_img.png'}",
     "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
     "sam_ckpt_path": "$@ckpt_dir + '/sam_vit_b_01ec64.pth'",
     "pretrained_ckpt_path": "$@ckpt_dir + '/model.pt'",

models/vista2d/configs/inference_trt.json

@@ -0,0 +1,10 @@
+{
+    "imports": [
+        "$import numpy",
+        "$from monai.networks import trt_compile"
+    ],
+    "trt_args": {
+        "dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]"
+    },
+    "network": "$trt_compile(@network_def.to(@device), @pretrained_ckpt_path, args=@trt_args)"
+}

models/vista2d/docs/README.md

@@ -66,7 +66,13 @@ torchrun --nproc_per_node=gpu -m monai.bundle run_workflow "scripts.workflow.Vis
 python -m monai.bundle run --config_file configs/inference.json
 ```
 
-Please note that the data used in the config file is: "/cellpose_dataset/test/001_img.png", if the dataset path is different or you want to do inference on another file, please modify in `configs/inference.json` accordingly.
+Please note that the data used in this config file is: "/cellpose_dataset/test/001_img.png", if the dataset path is different or you want to do inference on another file, please modify in `configs/inference.json` accordingly.
+
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
 
 ### Execute multi-GPU inference
 ```bash

models/vista3d/configs/inference_trt.json

@@ -0,0 +1,9 @@
+{
+    "+imports": [
+        "$from monai.networks import trt_compile"
+    ],
+    "trt_args": {
+        "dynamic_batchsize": "$[1, @inferer#sw_batch_size, @inferer#sw_batch_size]"
+    },
+    "network": "$trt_compile(@network_def.to(@device), @bundle_root + '/models/model.pt', args=@trt_args, submodule=['image_encoder.encoder', 'class_head'])"
+}

models/vista3d/docs/README.md

@@ -172,6 +172,13 @@ This default is overridable by changing the input folder `input_dir`, or the inp
 
 Set `"postprocessing#transforms#0#_disabled_": false` to move the postprocessing to cpu to reduce the GPU memory footprint.
 
+#### Execute inference with the TensorRT model:
+
+```
+python -m monai.bundle run --config_file "['configs/inference.json', 'configs/inference_trt.json']"
+```
+
+
 ## Automatic segmentation label prompts :
 The mapping between organ name and label prompt is in the [json file](labels.json)