Repository: Project-MONAI/MONAI

Issue:
Cannot run the lung_nodule_ct_detection TensorRT model end-to-end
**Is your feature request related to a problem? Please describe.**
When converting a PyTorch model to a TensorRT model, there are two ways, which are Torch-TensorRT way and ONNX-TensorRT way.  When using the ONNX-TensorRT way, it could be set to either use a TorchScript module or PyTorch  module as the input.

The Torch-TensorRT way of conversion suffers a slowdown issue currently.

The output of the lung_nodule_ct_detection model is a python dict, which is not  fully supported by the ONNX-TensorRT way. This way exports the pytorch model to an ONNX model first, then converts the ONNX model to a TensorRT engine and  wraps the TensorRT engine back to a torchscript at the end. 

For the first step which exports pytorch models to ONNX models, there are two ways. The first way directly exports the pytorch model with the `torch.onnx.export` API. The second way is firstly export the pytorch model to a torchscript model using `torch.jit.script` or `torch.jit.trace` and then pass the torchscript model to the `torch.onnx.export` API to get an onnx model, which is used by the [trt_export](https://github.com/Project-MONAI/MONAI/blob/c30a5b9e9a33bb0670c2b10821aecd3a2ceb9f85/monai/bundle/scripts.py#L1224).
 
If using the first way to exporting the onnx model, the output of the converted model would be a list of Tensor. If using the second way, this model would fail at `torch.jit.trace` phase or the `torch.jit.onnx` API would report an error if using `torch.jit.script` to trace the model. Neither way could make the model run an end-to-end inference, since it needs to feed the output to a detector head which takes a dict as input. 

**Describe the solution you'd like**
1. Add a control parameter  like `use_tuple_output` to the [__init__](https://github.com/Project-MONAI/MONAI/blob/79723cace5dbf3a7e9db55997e51629d08bed0ed/monai/apps/detection/networks/retinanet_network.py#L262) function of `RetinaNet` to control whether to use a dict as ouput or a list of tensor as output. Please note to add it to the `__init__` function and use it in the `foward` function so that it would be convenient for bundle.
2. The [outputs](https://github.com/Project-MONAI/MONAI/blob/79723cace5dbf3a7e9db55997e51629d08bed0ed/monai/apps/detection/networks/retinanet_network.py#L326)  of `foward` function should be a tuple of tensors or a dict decided by the control parameter above.
3. Remove [this](https://github.com/Project-MONAI/MONAI/blob/79723cace5dbf3a7e9db55997e51629d08bed0ed/monai/apps/detection/networks/retinanet_detector.py#L191) check and add `"spatial_dims", "num_classes", "cls_key", "box_reg_key", "num_anchors", "size_divisible"` as input parameters to the [__init__](https://github.com/Project-MONAI/MONAI/blob/79723cace5dbf3a7e9db55997e51629d08bed0ed/monai/apps/detection/networks/retinanet_detector.py#L182) function of RetinaNetDetector. And assign the attributes of the class with the input values if the `self.network`  doesn't have these attribute.
4. Rewrite [this logic](https://github.com/Project-MONAI/MONAI/blob/79723cace5dbf3a7e9db55997e51629d08bed0ed/monai/apps/detection/utils/predict_utils.py#L76) to parse the output list like the code below. (Just for reference)
```
head_outputs = network(images)
    if isinstance(head_outputs, (tuple, list)):
        tmp_dict = {}
        tmp_dict["classification"] = head_outputs[:3]
        tmp_dict["box_regression"] = head_outputs[3:]
        head_outputs = tmp_dict
    else:
        ensure_dict_value_to_list_(head_outputs, keys)
```

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.



Discussion:
Hi @Can-Zhao ,

Could you please help take a look at this issue?

Thanks in advance.


Resolving patch:
diff --git a/monai/apps/detection/networks/retinanet_detector.py b/monai/apps/detection/networks/retinanet_detector.py
--- a/monai/apps/detection/networks/retinanet_detector.py
+++ b/monai/apps/detection/networks/retinanet_detector.py
@@ -122,19 +122,23 @@ class RetinaNetDetector(nn.Module):
             - num_anchors (int) is the number of anchor shapes at each location. it should equal to
               ``self.anchor_generator.num_anchors_per_location()[0]``.
 
+            If network does not have these attributes, user needs to provide them for the detector.
+
         2. Its input should be an image Tensor sized (B, C, H, W) or (B, C, H, W, D).
 
-        3. About its output ``head_outputs``:
+        3. About its output ``head_outputs``, it should be either a list of tensors or a dictionary of str: List[Tensor]:
 
-            - It should be a dictionary with at least two keys:
-              ``network.cls_key`` and ``network.box_reg_key``.
-            - ``head_outputs[network.cls_key]`` should be List[Tensor] or Tensor. Each Tensor represents
+            - If it is a dictionary, it needs to have at least two keys:
+              ``network.cls_key`` and ``network.box_reg_key``, representing predicted classification maps and box regression maps.
+              ``head_outputs[network.cls_key]`` should be List[Tensor] or Tensor. Each Tensor represents
               classification logits map at one resolution level,
               sized (B, num_classes*num_anchors, H_i, W_i) or (B, num_classes*num_anchors, H_i, W_i, D_i).
-            - ``head_outputs[network.box_reg_key]`` should be List[Tensor] or Tensor. Each Tensor represents
+              ``head_outputs[network.box_reg_key]`` should be List[Tensor] or Tensor. Each Tensor represents
               box regression map at one resolution level,
               sized (B, 2*spatial_dims*num_anchors, H_i, W_i)or (B, 2*spatial_dims*num_anchors, H_i, W_i, D_i).
-            - ``len(head_outputs[network.cls_key]) == len(head_outputs[network.box_reg_key])``.
+              ``len(head_outputs[network.cls_key]) == len(head_outputs[network.box_reg_key])``.
+            - If it is a list of 2N tensors, the first N tensors should be the predicted classification maps,
+              and the second N tensors should be the predicted box regression maps.
 
     Example:
 
@@ -184,34 +188,33 @@ def __init__(
         network: nn.Module,
         anchor_generator: AnchorGenerator,
         box_overlap_metric: Callable = box_iou,
+        spatial_dims: int | None = None,  # used only when network.spatial_dims does not exist
+        num_classes: int | None = None,  # used only when network.num_classes does not exist
+        size_divisible: Sequence[int] | int = 1,  # used only when network.size_divisible does not exist
+        cls_key: str = "classification",  # used only when network.cls_key does not exist
+        box_reg_key: str = "box_regression",  # used only when network.box_reg_key does not exist
         debug: bool = False,
     ):
         super().__init__()
 
-        if not all(
-            hasattr(network, attr)
-            for attr in ["spatial_dims", "num_classes", "cls_key", "box_reg_key", "num_anchors", "size_divisible"]
-        ):
-            raise AttributeError(
-                "network should have attributes, including: "
-                "'spatial_dims', 'num_classes', 'cls_key', 'box_reg_key', 'num_anchors', 'size_divisible'."
-            )
-
         self.network = network
-        self.spatial_dims: int = self.network.spatial_dims  # type: ignore[assignment]
-        self.num_classes = self.network.num_classes
-        self.size_divisible = ensure_tuple_rep(self.network.size_divisible, self.spatial_dims)
+        # network attribute
+        self.spatial_dims = self.get_attribute_from_network("spatial_dims", default_value=spatial_dims)
+        self.num_classes = self.get_attribute_from_network("num_classes", default_value=num_classes)
+
+        self.size_divisible = self.get_attribute_from_network("size_divisible", default_value=size_divisible)
+        self.size_divisible = ensure_tuple_rep(self.size_divisible, self.spatial_dims)
         # keys for the network output
-        self.cls_key: str = self.network.cls_key  # type: ignore[assignment]
-        self.box_reg_key: str = self.network.box_reg_key  # type: ignore[assignment]
+        self.cls_key = self.get_attribute_from_network("cls_key", default_value=cls_key)
+        self.box_reg_key = self.get_attribute_from_network("box_reg_key", default_value=box_reg_key)
 
         # check if anchor_generator matches with network
         self.anchor_generator = anchor_generator
-
         self.num_anchors_per_loc = self.anchor_generator.num_anchors_per_location()[0]
-        if self.num_anchors_per_loc != self.network.num_anchors:
+        network_num_anchors = self.get_attribute_from_network("num_anchors", default_value=self.num_anchors_per_loc)
+        if self.num_anchors_per_loc != network_num_anchors:
             raise ValueError(
-                f"Number of feature map channels ({self.network.num_anchors}) "
+                f"Number of feature map channels ({network_num_anchors}) "
                 f"should match with number of anchors at each location ({self.num_anchors_per_loc})."
             )
         # if new coming input images has same shape with
@@ -252,6 +255,14 @@ def __init__(
             apply_sigmoid=True,
         )
 
+    def get_attribute_from_network(self, attr_name, default_value=None):
+        if hasattr(self.network, attr_name):
+            return getattr(self.network, attr_name)
+        elif default_value is not None:
+            return default_value
+        else:
+            raise ValueError(f"network does not have attribute {attr_name}, please provide it in the detector.")
+
     def set_box_coder_weights(self, weights: tuple[float]) -> None:
         """
         Set the weights for box coder.
@@ -497,7 +508,14 @@ def forward(
         # 3. Generate network outputs. Use inferer only in evaluation mode.
         if self.training or (not use_inferer):
             head_outputs = self.network(images)
-            ensure_dict_value_to_list_(head_outputs)  # ensure head_outputs is Dict[str, List[Tensor]]
+            if isinstance(head_outputs, (tuple, list)):
+                tmp_dict = {}
+                tmp_dict[self.cls_key] = head_outputs[: len(head_outputs) // 2]
+                tmp_dict[self.box_reg_key] = head_outputs[len(head_outputs) // 2 :]
+                head_outputs = tmp_dict
+            else:
+                # ensure head_outputs is Dict[str, List[Tensor]]
+                ensure_dict_value_to_list_(head_outputs)
         else:
             if self.inferer is None:
                 raise ValueError(
diff --git a/monai/apps/detection/networks/retinanet_network.py b/monai/apps/detection/networks/retinanet_network.py
--- a/monai/apps/detection/networks/retinanet_network.py
+++ b/monai/apps/detection/networks/retinanet_network.py
@@ -41,7 +41,7 @@
 
 import math
 from collections.abc import Callable, Sequence
-from typing import Dict
+from typing import Any, Dict
 
 import torch
 from torch import Tensor, nn
@@ -203,9 +203,11 @@ class RetinaNet(nn.Module):
     """
     The network used in RetinaNet.
 
-    It takes an image tensor as inputs, and outputs a dictionary ``head_outputs``.
+    It takes an image tensor as inputs, and outputs either 1) a dictionary ``head_outputs``.
     ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
     ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
+    or 2) a list of 2N tensors ``head_outputs``, with first N tensors being the predicted
+    classification maps and second N tensors being the predicted box regression maps.
 
     Args:
         spatial_dims: number of spatial dimensions of the images. We support both 2D and 3D images.
@@ -217,6 +219,11 @@ class RetinaNet(nn.Module):
             It can be the output of ``resnet_fpn_feature_extractor(*args, **kwargs)``.
         size_divisible: the spatial size of the network input should be divisible by size_divisible,
             decided by the feature_extractor.
+        use_list_output: default False. If False, the network outputs a dictionary ``head_outputs``,
+            ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
+            ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
+            If True, the network outputs a list of 2N tensors ``head_outputs``, with first N tensors being
+            the predicted classification maps and second N tensors being the predicted box regression maps.
 
     Example:
 
@@ -255,7 +262,7 @@ class RetinaNet(nn.Module):
                 size_divisible = size_divisible,
             ).to(device)
             result = model(torch.rand(2, 1, 128,128,128))
-            cls_logits_maps = result["cls_logits"]  # a list of len(returned_layers)+1 Tensor
+            cls_logits_maps = result["classification"]  # a list of len(returned_layers)+1 Tensor
             box_regression_maps = result["box_regression"]  # a list of len(returned_layers)+1 Tensor
     """
 
@@ -266,12 +273,14 @@ def __init__(
         num_anchors: int,
         feature_extractor: nn.Module,
         size_divisible: Sequence[int] | int = 1,
+        use_list_output: bool = False,
     ):
         super().__init__()
 
         self.spatial_dims = look_up_option(spatial_dims, supported=[1, 2, 3])
         self.num_classes = num_classes
         self.size_divisible = ensure_tuple_rep(size_divisible, self.spatial_dims)
+        self.use_list_output = use_list_output
 
         if not hasattr(feature_extractor, "out_channels"):
             raise ValueError(
@@ -293,19 +302,21 @@ def __init__(
         self.cls_key: str = "classification"
         self.box_reg_key: str = "box_regression"
 
-    def forward(self, images: Tensor) -> dict[str, list[Tensor]]:
+    def forward(self, images: Tensor) -> Any:
         """
-        It takes an image tensor as inputs, and outputs a dictionary ``head_outputs``.
-        ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
-        ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
+        It takes an image tensor as inputs, and outputs predicted classification maps
+        and predicted box regression maps in ``head_outputs``.
 
         Args:
             images: input images, sized (B, img_channels, H, W) or (B, img_channels, H, W, D).
 
         Return:
-            a dictionary ``head_outputs`` with keys including self.cls_key and self.box_reg_key.
+            1) If self.use_list_output is False, output a dictionary ``head_outputs`` with
+            keys including self.cls_key and self.box_reg_key.
             ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
             ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
+            2) if self.use_list_output is True, outputs a list of 2N tensors ``head_outputs``, with first N tensors being
+            the predicted classification maps and second N tensors being the predicted box regression maps.
 
         """
         # compute features maps list from the input images.
@@ -323,10 +334,15 @@ def forward(self, images: Tensor) -> dict[str, list[Tensor]]:
         # compute classification and box regression maps from the feature maps
         # expandable for mask prediction in the future
 
-        head_outputs: dict[str, list[Tensor]] = {self.cls_key: self.classification_head(feature_maps)}
-        head_outputs[self.box_reg_key] = self.regression_head(feature_maps)
-
-        return head_outputs
+        if not self.use_list_output:
+            # output dict
+            head_outputs = {self.cls_key: self.classification_head(feature_maps)}
+            head_outputs[self.box_reg_key] = self.regression_head(feature_maps)
+            return head_outputs
+        else:
+            # output list of tensor, first half is classification, second half is box regression
+            head_outputs_sequence = self.classification_head(feature_maps) + self.regression_head(feature_maps)
+            return head_outputs_sequence
 
 
 def resnet_fpn_feature_extractor(
diff --git a/monai/apps/detection/utils/predict_utils.py b/monai/apps/detection/utils/predict_utils.py
--- a/monai/apps/detection/utils/predict_utils.py
+++ b/monai/apps/detection/utils/predict_utils.py
@@ -73,6 +73,12 @@ def _network_sequence_output(images: Tensor, network: nn.Module, keys: list[str]
         network output values concat to a single List[Tensor]
     """
     head_outputs = network(images)
+
+    # if head_outputs is already a sequence of tensors, directly output it
+    if isinstance(head_outputs, (tuple, list)):
+        return list(head_outputs)
+
+    # if head_outputs is a dict
     ensure_dict_value_to_list_(head_outputs, keys)
     if keys is None:
         keys = list(head_outputs.keys())


Test patch: