Repository: Project-MONAI/MONAI
Issue category: Performance Issue

Issue:
Optimize the VISTA3D latency
**Is your feature request related to a problem? Please describe.**
The VISTA3D is a very brilliant model for everything segmentation in a medical image. However it suffers a slowdown [issue](https://github.com/Project-MONAI/model-zoo/tree/dev/models/vista3d/docs#inference-gpu-benchmarks) when running on a GPU platform, which takes minutes to run a single inference. Therefore, a lower latency version of the VISTA3D model should be a user-friendly model and gain more usage.

**Describe the solution you'd like**
1. Update the VISTA3D bundle parameters to achieve a better latency.
2. Update the VISTA3D network to solve some latency bottlenecks.
3. Try to accelerate the bundle with acceleration tools like TensorRT.

**Additional context**
Some analyzation has already done in [this PR](https://github.com/Project-MONAI/model-zoo/pull/671).
The corresponding MONAI core code change is included [here](https://github.com/binliunls/MONAI/pull/1)



Hints:

Patch:
diff --git a/monai/networks/nets/segresnet_ds.py b/monai/networks/nets/segresnet_ds.py
index 1ac5a79ee3..098e490511 100644
--- a/monai/networks/nets/segresnet_ds.py
+++ b/monai/networks/nets/segresnet_ds.py
@@ -508,8 +508,10 @@ def forward(  # type: ignore
 
         outputs: list[torch.Tensor] = []
         outputs_auto: list[torch.Tensor] = []
-        x_ = x.clone()
+        x_ = x
         if with_point:
+            if with_label:
+                x_ = x.clone()
             i = 0
             for level in self.up_layers:
                 x = level["upsample"](x)
diff --git a/monai/networks/nets/vista3d.py b/monai/networks/nets/vista3d.py
index 4215a9a594..6313b7812d 100644
--- a/monai/networks/nets/vista3d.py
+++ b/monai/networks/nets/vista3d.py
@@ -639,12 +639,10 @@ def forward(self, src: torch.Tensor, class_vector: torch.Tensor):
         if self.use_mlp:
             class_embedding = self.mlp(class_embedding)
         # [b,1,feat] @ [1,feat,dim], batch dimension become class_embedding batch dimension.
-        masks = []
-        for i in range(b):
-            mask = class_embedding @ src[[i]].view(1, c, h * w * d)
-            masks.append(mask.view(-1, 1, h, w, d))
+        masks_embedding = class_embedding.squeeze() @ src.view(b, c, h * w * d)
+        masks_embedding = masks_embedding.view(b, -1, h, w, d).transpose(0, 1)
 
-        return torch.cat(masks, 1), class_embedding
+        return masks_embedding, class_embedding
 
 
 class TwoWayTransformer(nn.Module):


Test patch: