Repository: Project-MONAI/MONAI

Issue:
Enhance `monai.bundle.verify_net_in_out` to enable multiple input in forward
Currently, the function only supports to input a tensor into the network:
https://github.com/Project-MONAI/MONAI/blob/5fae0c1c7525bebee1e43a625c7777f594cc97a4/monai/bundle/scripts.py#L936

However, to make it more generic, we should support kwargs for networks that have multiple input args, for example:

https://github.com/Project-MONAI/GenerativeModels/blob/fd04ec6f98a1aec7b6886dff1cfb4d0fa72fe4fe/generative/networks/nets/diffusion_model_unet.py#L1848

This ticket is also related to:
https://github.com/Project-MONAI/model-zoo/pull/390


Discussion:
cc @Nic-Ma @Can-Zhao 
Thanks for raising the request, please submit a PR for it.

Resolving patch:
diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py
--- a/monai/bundle/scripts.py
+++ b/monai/bundle/scripts.py
@@ -862,6 +862,7 @@ def verify_net_in_out(
     p: int | None = None,
     n: int | None = None,
     any: int | None = None,
+    extra_forward_args: dict | None = None,
     args_file: str | None = None,
     **override: Any,
 ) -> None:
@@ -885,6 +886,8 @@ def verify_net_in_out(
         p: power factor to generate fake data shape if dim of expected shape is "x**p", default to 1.
         n: multiply factor to generate fake data shape if dim of expected shape is "x*n", default to 1.
         any: specified size to generate fake data shape if dim of expected shape is "*", default to 1.
+        extra_forward_args: a dictionary that contains other args for the forward function of the network.
+            Default to an empty dictionary.
         args_file: a JSON or YAML file to provide default values for `net_id`, `meta_file`, `config_file`,
             `device`, `p`, `n`, `any`, and override pairs. so that the command line inputs can be simplified.
         override: id-value pairs to override or add the corresponding config content.
@@ -901,11 +904,20 @@ def verify_net_in_out(
         p=p,
         n=n,
         any=any,
+        extra_forward_args=extra_forward_args,
         **override,
     )
     _log_input_summary(tag="verify_net_in_out", args=_args)
-    config_file_, meta_file_, net_id_, device_, p_, n_, any_ = _pop_args(
-        _args, "config_file", "meta_file", net_id="", device="cuda:0" if is_available() else "cpu", p=1, n=1, any=1
+    config_file_, meta_file_, net_id_, device_, p_, n_, any_, extra_forward_args_ = _pop_args(
+        _args,
+        "config_file",
+        "meta_file",
+        net_id="",
+        device="cuda:0" if is_available() else "cpu",
+        p=1,
+        n=1,
+        any=1,
+        extra_forward_args={},
     )
 
     parser = ConfigParser()
@@ -933,10 +945,10 @@ def verify_net_in_out(
             from torch.cuda.amp import autocast
 
             with autocast():
-                output = net(test_data.cuda())
+                output = net(test_data.cuda(), **extra_forward_args_)
             net.to(device_)
         else:
-            output = net(test_data)
+            output = net(test_data, **extra_forward_args_)
         if output.shape[1] != output_channels:
             raise ValueError(f"output channel number `{output.shape[1]}` doesn't match: `{output_channels}`.")
         if output.dtype != output_dtype:


Test patch:
diff --git a/tests/test_bundle_verify_net.py b/tests/test_bundle_verify_net.py
--- a/tests/test_bundle_verify_net.py
+++ b/tests/test_bundle_verify_net.py
@@ -17,7 +17,7 @@
 
 from parameterized import parameterized
 
-from monai.bundle import ConfigParser
+from monai.bundle import ConfigParser, verify_net_in_out
 from tests.utils import command_line_tests, skip_if_no_cuda, skip_if_windows
 
 TEST_CASE_1 = [
@@ -55,6 +55,19 @@ def test_verify_fp16(self, meta_file, config_file):
             cmd += ["--_meta_#network_data_format#outputs#pred#dtype", "float16"]
             command_line_tests(cmd)
 
+    @parameterized.expand([TEST_CASE_1])
+    @skip_if_no_cuda
+    def test_verify_fp16_extra_forward_args(self, meta_file, config_file):
+        verify_net_in_out(
+            net_id="network_def",
+            meta_file=meta_file,
+            config_file=config_file,
+            n=4,
+            any=16,
+            extra_forward_args={"extra_arg1": 1, "extra_arg2": 2},
+            **{"network_def#_target_": "tests.testing_data.bundle_test_network.TestMultiInputUNet"},
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/testing_data/bundle_test_network.py b/tests/testing_data/bundle_test_network.py
new file mode 100644
--- /dev/null
+++ b/tests/testing_data/bundle_test_network.py
@@ -0,0 +1,29 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import torch
+
+from monai.networks.nets import UNet
+
+
+class TestMultiInputUNet(UNet):
+    """
+    This class is used for "tests/test_bundle_verify_net.py" to show that the monai.bundle.verify_net_in_out
+    function supports to verify networks that have multiple args as the input in the forward function.
+    """
+
+    def forward(self, x: torch.Tensor, extra_arg1: int, extra_arg2: int) -> torch.Tensor:  # type: ignore
+        x = self.model(x)
+        x += extra_arg1
+        x += extra_arg2
+        return x