Added trt_compile configs for vista2d and vista3d

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)