From 798a88dffba97de2bf12cc20757614aa94f78a2e Mon Sep 17 00:00:00 2001 From: kyinhub Date: Sat, 25 Jul 2026 11:19:30 -0700 Subject: [PATCH 1/2] fix(networks): support QuickNAT without optional SE blocks Signed-off-by: kyinhub --- monai/networks/nets/quicknat.py | 27 ++++++++++++++++-- tests/networks/nets/test_quicknat.py | 42 ++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/monai/networks/nets/quicknat.py b/monai/networks/nets/quicknat.py index f82e3d7e78..01381af289 100644 --- a/monai/networks/nets/quicknat.py +++ b/monai/networks/nets/quicknat.py @@ -43,7 +43,28 @@ class SkipConnectionWithIdx(SkipConnection): """ def forward(self, input, indices): # type: ignore[override] - return super().forward(input), indices + """Combine the input with the indexed submodule output. + + Args: + input: input tensor for the skip connection and submodule. + indices: pooling indices to preserve for the following decoder. + + Returns: + A tuple containing the combined tensor and the original indices. + + Raises: + NotImplementedError: if the configured skip mode is unsupported. + """ + submodule_output, _ = self.submodule(input, None) + if self.mode == "cat": + output = torch.cat([input, submodule_output], dim=self.dim) + elif self.mode == "add": + output = torch.add(input, submodule_output) + elif self.mode == "mul": + output = torch.mul(input, submodule_output) + else: + raise NotImplementedError(f"Unsupported mode {self.mode}.") + return output, indices class SequentialWithIdx(nn.Sequential): @@ -325,7 +346,7 @@ class Quicknat(nn.Module): stride_convolution: convolution stride. Defaults to 1. pool: kernel size of the pooling layer, stride_pool: stride for the pooling layer. - se_block: Squeeze and Excite block type to be included, defaults to None. Valid options : NONE, CSE, SSE, CSSE, + se_block: Squeeze and Excite block type to include. Use ``None`` or ``"None"`` to disable it; valid enabled values are ``"CSE"``, ``"SSE"``, and ``"CSSE"``. droup_out: dropout ratio. Defaults to no dropout. act: activation type and arguments. Defaults to PReLU. norm: feature normalization type and arguments. Defaults to instance norm. @@ -358,7 +379,7 @@ def __init__( pool: int = 2, stride_pool: int = 2, # Valid options : NONE, CSE, SSE, CSSE - se_block: str = "None", + se_block: str | None = "None", drop_out: float = 0, act: tuple | str = Act.PRELU, norm: tuple | str = Norm.INSTANCE, diff --git a/tests/networks/nets/test_quicknat.py b/tests/networks/nets/test_quicknat.py index 6653965c08..01ed623a41 100644 --- a/tests/networks/nets/test_quicknat.py +++ b/tests/networks/nets/test_quicknat.py @@ -18,11 +18,21 @@ from monai.networks import eval_mode from monai.networks.nets import Quicknat +from monai.networks.nets.quicknat import SkipConnectionWithIdx from monai.utils import optional_import from tests.test_utils import test_script_save _, has_se = optional_import("squeeze_and_excitation") + +class _DoubleWithIdx(torch.nn.Module): + """Double tensor values using QuickNAT's indexed-module signature.""" + + def forward(self, input, indices): + """Return the doubled tensor and unchanged indices.""" + return input * 2, indices + + TEST_CASES = [ # params, input_shape, expected_shape [{"num_classes": 1, "num_channels": 1, "num_filters": 1, "se_block": None}, (1, 1, 32, 32), (1, 1, 32, 32)], @@ -36,6 +46,38 @@ ] +class TestQuicknatCore(unittest.TestCase): + """Test QuickNAT paths that do not require optional dependencies.""" + + @parameterized.expand(["cat", "add", "mul"]) + def test_skip_connection_modes_preserve_indices(self, mode): + """Verify each fusion mode and preservation of pooling indices.""" + skip = SkipConnectionWithIdx(_DoubleWithIdx(), mode=mode) + input = torch.tensor([[[[2.0]]]]) + indices = torch.tensor([[[[3]]]]) + + output, returned_indices = skip(input, indices) + + expected = {"cat": torch.cat([input, input * 2], dim=1), "add": input * 3, "mul": input * input * 2}[mode] + self.assertTrue(torch.equal(output, expected)) + self.assertIs(returned_indices, indices) + + def test_skip_connection_rejects_unsupported_mode(self): + """Verify unsupported fusion modes fail explicitly.""" + skip = SkipConnectionWithIdx(_DoubleWithIdx(), mode="cat") + skip.mode = "unsupported" + + with self.assertRaisesRegex(NotImplementedError, "Unsupported mode"): + skip(torch.ones((1, 1, 1, 1)), torch.ones((1, 1, 1, 1))) + + def test_forward_without_optional_se_dependency(self): + """Verify QuickNAT runs when squeeze-and-excitation is disabled.""" + net = Quicknat(num_classes=2, num_channels=1, num_filters=4, se_block=None) + with eval_mode(net): + result = net(torch.randn(1, 1, 32, 32)) + self.assertEqual(result.shape, (1, 2, 32, 32)) + + @unittest.skipUnless(has_se, "squeeze_and_excitation not installed") class TestQuicknat(unittest.TestCase): @parameterized.expand(TEST_CASES) From fa90c9f55c6f65b8bd047954217d552522544338 Mon Sep 17 00:00:00 2001 From: kyinhub Date: Sat, 25 Jul 2026 14:04:37 -0700 Subject: [PATCH 2/2] fix(networks): forward QuickNAT skip indices Signed-off-by: kyinhub --- monai/networks/nets/quicknat.py | 2 +- tests/networks/nets/test_quicknat.py | 29 +++++++++++++++++++++------- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/monai/networks/nets/quicknat.py b/monai/networks/nets/quicknat.py index 01381af289..a934b26c31 100644 --- a/monai/networks/nets/quicknat.py +++ b/monai/networks/nets/quicknat.py @@ -55,7 +55,7 @@ def forward(self, input, indices): # type: ignore[override] Raises: NotImplementedError: if the configured skip mode is unsupported. """ - submodule_output, _ = self.submodule(input, None) + submodule_output, _ = self.submodule(input, indices) if self.mode == "cat": output = torch.cat([input, submodule_output], dim=self.dim) elif self.mode == "add": diff --git a/tests/networks/nets/test_quicknat.py b/tests/networks/nets/test_quicknat.py index 01ed623a41..913f30b661 100644 --- a/tests/networks/nets/test_quicknat.py +++ b/tests/networks/nets/test_quicknat.py @@ -28,9 +28,22 @@ class _DoubleWithIdx(torch.nn.Module): """Double tensor values using QuickNAT's indexed-module signature.""" - def forward(self, input, indices): - """Return the doubled tensor and unchanged indices.""" - return input * 2, indices + def __init__(self): + super().__init__() + self.received_indices = None + + def forward(self, tensor, indices): + """Return doubled values and the unchanged index object. + + Args: + tensor: input tensor to double. + indices: pooling indices passed through the indexed module. + + Returns: + A tuple containing the doubled tensor and the original indices. + """ + self.received_indices = indices + return tensor * 2, indices TEST_CASES = [ @@ -52,14 +65,16 @@ class TestQuicknatCore(unittest.TestCase): @parameterized.expand(["cat", "add", "mul"]) def test_skip_connection_modes_preserve_indices(self, mode): """Verify each fusion mode and preservation of pooling indices.""" - skip = SkipConnectionWithIdx(_DoubleWithIdx(), mode=mode) - input = torch.tensor([[[[2.0]]]]) + submodule = _DoubleWithIdx() + skip = SkipConnectionWithIdx(submodule, mode=mode) + tensor = torch.tensor([[[[2.0]]]]) indices = torch.tensor([[[[3]]]]) - output, returned_indices = skip(input, indices) + output, returned_indices = skip(tensor, indices) - expected = {"cat": torch.cat([input, input * 2], dim=1), "add": input * 3, "mul": input * input * 2}[mode] + expected = {"cat": torch.cat([tensor, tensor * 2], dim=1), "add": tensor * 3, "mul": tensor * tensor * 2}[mode] self.assertTrue(torch.equal(output, expected)) + self.assertIs(submodule.received_indices, indices) self.assertIs(returned_indices, indices) def test_skip_connection_rejects_unsupported_mode(self):