monai-weekly 1.5.dev2506__py3-none-any.whl → 1.5.dev2507__py3-none-any.whl

tests/networks/blocks/test_crossattention.py

@@ -0,0 +1,185 @@
+# 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 unittest
+from unittest import skipUnless
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.networks import eval_mode
+from monai.networks.blocks.crossattention import CrossAttentionBlock
+from monai.networks.layers.factories import RelPosEmbedding
+from monai.utils import optional_import
+from tests.test_utils import SkipIfBeforePyTorchVersion, assert_allclose
+
+einops, has_einops = optional_import("einops")
+
+TEST_CASE_CABLOCK = []
+for dropout_rate in np.linspace(0, 1, 4):
+    for hidden_size in [360, 480, 600, 768]:
+        for num_heads in [4, 6, 8, 12]:
+            for rel_pos_embedding in [None, RelPosEmbedding.DECOMPOSED]:
+                for input_size in [(16, 32), (8, 8, 8)]:
+                    for flash_attn in [True, False]:
+                        test_case = [
+                            {
+                                "hidden_size": hidden_size,
+                                "num_heads": num_heads,
+                                "dropout_rate": dropout_rate,
+                                "rel_pos_embedding": rel_pos_embedding if not flash_attn else None,
+                                "input_size": input_size,
+                                "use_flash_attention": flash_attn,
+                            },
+                            (2, 512, hidden_size),
+                            (2, 512, hidden_size),
+                        ]
+                        TEST_CASE_CABLOCK.append(test_case)
+
+
+class TestResBlock(unittest.TestCase):
+    @parameterized.expand(TEST_CASE_CABLOCK)
+    @skipUnless(has_einops, "Requires einops")
+    @SkipIfBeforePyTorchVersion((2, 0))
+    def test_shape(self, input_param, input_shape, expected_shape):
+        # Without flash attention
+        net = CrossAttentionBlock(**input_param)
+        with eval_mode(net):
+            result = net(torch.randn(input_shape), context=torch.randn(2, 512, input_param["hidden_size"]))
+            self.assertEqual(result.shape, expected_shape)
+
+    def test_ill_arg(self):
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(hidden_size=128, num_heads=12, dropout_rate=6.0)
+
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(hidden_size=620, num_heads=8, dropout_rate=0.4)
+
+    @SkipIfBeforePyTorchVersion((2, 0))
+    def test_save_attn_with_flash_attention(self):
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(
+                hidden_size=128, num_heads=3, dropout_rate=0.1, use_flash_attention=True, save_attn=True
+            )
+
+    @SkipIfBeforePyTorchVersion((2, 0))
+    def test_rel_pos_embedding_with_flash_attention(self):
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(
+                hidden_size=128,
+                num_heads=3,
+                dropout_rate=0.1,
+                use_flash_attention=True,
+                save_attn=False,
+                rel_pos_embedding=RelPosEmbedding.DECOMPOSED,
+            )
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_attention_dim_not_multiple_of_heads(self):
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(hidden_size=128, num_heads=3, dropout_rate=0.1)
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_inner_dim_different(self):
+        CrossAttentionBlock(hidden_size=128, num_heads=4, dropout_rate=0.1, dim_head=30)
+
+    def test_causal_no_sequence_length(self):
+        with self.assertRaises(ValueError):
+            CrossAttentionBlock(hidden_size=128, num_heads=4, dropout_rate=0.1, causal=True)
+
+    @skipUnless(has_einops, "Requires einops")
+    @SkipIfBeforePyTorchVersion((2, 0))
+    def test_causal_flash_attention(self):
+        block = CrossAttentionBlock(
+            hidden_size=128,
+            num_heads=1,
+            dropout_rate=0.1,
+            causal=True,
+            sequence_length=16,
+            save_attn=False,
+            use_flash_attention=True,
+        )
+        input_shape = (1, 16, 128)
+        # Check it runs correctly
+        block(torch.randn(input_shape))
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_causal(self):
+        block = CrossAttentionBlock(
+            hidden_size=128, num_heads=1, dropout_rate=0.1, causal=True, sequence_length=16, save_attn=True
+        )
+        input_shape = (1, 16, 128)
+        block(torch.randn(input_shape))
+        # check upper triangular part of the attention matrix is zero
+        assert torch.triu(block.att_mat, diagonal=1).sum() == 0
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_context_input(self):
+        block = CrossAttentionBlock(
+            hidden_size=128, num_heads=1, dropout_rate=0.1, causal=True, sequence_length=16, context_input_size=12
+        )
+        input_shape = (1, 16, 128)
+        block(torch.randn(input_shape), context=torch.randn(1, 3, 12))
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_context_wrong_input_size(self):
+        block = CrossAttentionBlock(
+            hidden_size=128, num_heads=1, dropout_rate=0.1, causal=True, sequence_length=16, context_input_size=12
+        )
+        input_shape = (1, 16, 128)
+        with self.assertRaises(RuntimeError):
+            block(torch.randn(input_shape), context=torch.randn(1, 3, 24))
+
+    @skipUnless(has_einops, "Requires einops")
+    def test_access_attn_matrix(self):
+        # input format
+        hidden_size = 128
+        num_heads = 2
+        dropout_rate = 0
+        input_shape = (2, 256, hidden_size)
+
+        # be  not able to access the matrix
+        no_matrix_acess_blk = CrossAttentionBlock(
+            hidden_size=hidden_size, num_heads=num_heads, dropout_rate=dropout_rate
+        )
+        no_matrix_acess_blk(torch.randn(input_shape))
+        assert isinstance(no_matrix_acess_blk.att_mat, torch.Tensor)
+        # no of elements is zero
+        assert no_matrix_acess_blk.att_mat.nelement() == 0
+
+        # be able to acess the attention matrix.
+        matrix_acess_blk = CrossAttentionBlock(
+            hidden_size=hidden_size, num_heads=num_heads, dropout_rate=dropout_rate, save_attn=True
+        )
+        matrix_acess_blk(torch.randn(input_shape))
+        assert matrix_acess_blk.att_mat.shape == (input_shape[0], input_shape[0], input_shape[1], input_shape[1])
+
+    @parameterized.expand([[True], [False]])
+    @skipUnless(has_einops, "Requires einops")
+    @SkipIfBeforePyTorchVersion((2, 0))
+    def test_flash_attention(self, causal):
+        input_param = {"hidden_size": 128, "num_heads": 1, "causal": causal, "sequence_length": 16 if causal else None}
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        block_w_flash_attention = CrossAttentionBlock(**input_param, use_flash_attention=True).to(device)
+        block_wo_flash_attention = CrossAttentionBlock(**input_param, use_flash_attention=False).to(device)
+        block_wo_flash_attention.load_state_dict(block_w_flash_attention.state_dict())
+        test_data = torch.randn(1, 16, 128).to(device)
+
+        out_1 = block_w_flash_attention(test_data)
+        out_2 = block_wo_flash_attention(test_data)
+        assert_allclose(out_1, out_2, atol=1e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/blocks/test_denseblock.py

@@ -0,0 +1,105 @@
+# 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 unittest
+
+import torch.nn as nn
+
+from monai.networks.blocks import ConvDenseBlock, DenseBlock
+from tests.test_utils import TorchImageTestCase2D, TorchImageTestCase3D
+
+
+class TestDenseBlock2D(TorchImageTestCase2D):
+    def test_block_empty(self):
+        block = DenseBlock([])
+        out = block(self.imt)
+        expected_shape = self.imt.shape
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_block_conv(self):
+        conv1 = nn.Conv2d(self.input_channels, self.output_channels, 3, padding=1)
+        conv2 = nn.Conv2d(self.input_channels + self.output_channels, self.input_channels, 3, padding=1)
+        block = DenseBlock([conv1, conv2])
+        out = block(self.imt)
+        expected_shape = (1, self.output_channels + self.input_channels * 2, self.im_shape[0], self.im_shape[1])
+        self.assertEqual(out.shape, expected_shape)
+
+
+class TestDenseBlock3D(TorchImageTestCase3D):
+    def test_block_conv(self):
+        conv1 = nn.Conv3d(self.input_channels, self.output_channels, 3, padding=1)
+        conv2 = nn.Conv3d(self.input_channels + self.output_channels, self.input_channels, 3, padding=1)
+        block = DenseBlock([conv1, conv2])
+        out = block(self.imt)
+        expected_shape = (
+            1,
+            self.output_channels + self.input_channels * 2,
+            self.im_shape[0],
+            self.im_shape[1],
+            self.im_shape[2],
+        )
+        self.assertEqual(out.shape, expected_shape)
+
+
+class TestConvDenseBlock2D(TorchImageTestCase2D):
+    def test_block_empty(self):
+        conv = ConvDenseBlock(spatial_dims=2, in_channels=self.input_channels, channels=[])
+        out = conv(self.imt)
+        expected_shape = self.imt.shape
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_except(self):
+        with self.assertRaises(ValueError):
+            _ = ConvDenseBlock(spatial_dims=2, in_channels=self.input_channels, channels=[1, 2], dilations=[1, 2, 3])
+
+    def test_block1(self):
+        channels = [2, 4]
+        conv = ConvDenseBlock(spatial_dims=2, in_channels=self.input_channels, channels=channels)
+        out = conv(self.imt)
+        expected_shape = (1, self.input_channels + sum(channels), self.im_shape[0], self.im_shape[1])
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_block2(self):
+        channels = [2, 4]
+        dilations = [1, 2]
+        conv = ConvDenseBlock(spatial_dims=2, in_channels=self.input_channels, channels=channels, dilations=dilations)
+        out = conv(self.imt)
+        expected_shape = (1, self.input_channels + sum(channels), self.im_shape[0], self.im_shape[1])
+        self.assertEqual(out.shape, expected_shape)
+
+
+class TestConvDenseBlock3D(TorchImageTestCase3D):
+    def test_block_empty(self):
+        conv = ConvDenseBlock(spatial_dims=3, in_channels=self.input_channels, channels=[])
+        out = conv(self.imt)
+        expected_shape = self.imt.shape
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_block1(self):
+        channels = [2, 4]
+        conv = ConvDenseBlock(spatial_dims=3, in_channels=self.input_channels, channels=channels)
+        out = conv(self.imt)
+        expected_shape = (1, self.input_channels + sum(channels), self.im_shape[0], self.im_shape[1], self.im_shape[2])
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_block2(self):
+        channels = [2, 4]
+        dilations = [1, 2]
+        conv = ConvDenseBlock(spatial_dims=3, in_channels=self.input_channels, channels=channels, dilations=dilations)
+        out = conv(self.imt)
+        expected_shape = (1, self.input_channels + sum(channels), self.im_shape[0], self.im_shape[1], self.im_shape[2])
+        self.assertEqual(out.shape, expected_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/blocks/test_dynunet_block.py

@@ -0,0 +1,116 @@
+# 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 unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks import eval_mode
+from monai.networks.blocks.dynunet_block import UnetBasicBlock, UnetResBlock, UnetUpBlock, get_padding
+from tests.test_utils import test_script_save
+
+TEST_CASE_RES_BASIC_BLOCK = []
+for spatial_dims in range(2, 4):
+    for kernel_size in [1, 3]:
+        for stride in [1, 2]:
+            for norm_name in [("GROUP", {"num_groups": 16}), ("batch", {"track_running_stats": False}), "instance"]:
+                for in_size in [15, 16]:
+                    padding = get_padding(kernel_size, stride)
+                    if not isinstance(padding, int):
+                        padding = padding[0]
+                    out_size = int((in_size + 2 * padding - kernel_size) / stride) + 1
+                    test_case = [
+                        {
+                            "spatial_dims": spatial_dims,
+                            "in_channels": 16,
+                            "out_channels": 16,
+                            "kernel_size": kernel_size,
+                            "norm_name": norm_name,
+                            "act_name": ("leakyrelu", {"inplace": True, "negative_slope": 0.1}),
+                            "stride": stride,
+                        },
+                        (1, 16, *([in_size] * spatial_dims)),
+                        (1, 16, *([out_size] * spatial_dims)),
+                    ]
+                    TEST_CASE_RES_BASIC_BLOCK.append(test_case)
+
+TEST_UP_BLOCK = []
+in_channels, out_channels = 4, 2
+for spatial_dims in range(2, 4):
+    for kernel_size in [1, 3]:
+        for stride in [1, 2]:
+            for norm_name in ["batch", "instance"]:
+                for in_size in [15, 16]:
+                    for trans_bias in [True, False]:
+                        out_size = in_size * stride
+                        test_case = [
+                            {
+                                "spatial_dims": spatial_dims,
+                                "in_channels": in_channels,
+                                "out_channels": out_channels,
+                                "kernel_size": kernel_size,
+                                "norm_name": norm_name,
+                                "stride": stride,
+                                "upsample_kernel_size": stride,
+                                "trans_bias": trans_bias,
+                            },
+                            (1, in_channels, *([in_size] * spatial_dims)),
+                            (1, out_channels, *([out_size] * spatial_dims)),
+                            (1, out_channels, *([in_size * stride] * spatial_dims)),
+                        ]
+                        TEST_UP_BLOCK.append(test_case)
+
+
+class TestResBasicBlock(unittest.TestCase):
+    @parameterized.expand(TEST_CASE_RES_BASIC_BLOCK)
+    def test_shape(self, input_param, input_shape, expected_shape):
+        for net in [UnetResBlock(**input_param), UnetBasicBlock(**input_param)]:
+            with eval_mode(net):
+                result = net(torch.randn(input_shape))
+                self.assertEqual(result.shape, expected_shape)
+
+    def test_ill_arg(self):
+        with self.assertRaises(ValueError):
+            UnetBasicBlock(3, 4, 2, kernel_size=3, stride=1, norm_name="norm")
+        with self.assertRaises(AssertionError):
+            UnetResBlock(3, 4, 2, kernel_size=1, stride=4, norm_name="batch")
+
+    def test_script(self):
+        input_param, input_shape, _ = TEST_CASE_RES_BASIC_BLOCK[0]
+
+        for net_type in (UnetResBlock, UnetBasicBlock):
+            net = net_type(**input_param)
+            test_data = torch.randn(input_shape)
+            test_script_save(net, test_data)
+
+
+class TestUpBlock(unittest.TestCase):
+    @parameterized.expand(TEST_UP_BLOCK)
+    def test_shape(self, input_param, input_shape, expected_shape, skip_shape):
+        net = UnetUpBlock(**input_param)
+        with eval_mode(net):
+            result = net(torch.randn(input_shape), torch.randn(skip_shape))
+            self.assertEqual(result.shape, expected_shape)
+
+    def test_script(self):
+        input_param, input_shape, _, skip_shape = TEST_UP_BLOCK[0]
+
+        net = UnetUpBlock(**input_param)
+        test_data = torch.randn(input_shape)
+        skip_data = torch.randn(skip_shape)
+        test_script_save(net, test_data, skip_data)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/blocks/test_fpn_block.py

@@ -0,0 +1,88 @@
+# 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 unittest
+from collections import OrderedDict
+
+import torch
+from parameterized import parameterized
+
+from monai.networks.blocks.backbone_fpn_utils import _resnet_fpn_extractor
+from monai.networks.blocks.feature_pyramid_network import FeaturePyramidNetwork
+from monai.networks.nets.resnet import resnet50
+from monai.utils import optional_import
+from tests.test_utils import SkipIfBeforePyTorchVersion, test_script_save
+
+_, has_torchvision = optional_import("torchvision")
+
+TEST_CASES = [
+    [
+        {"spatial_dims": 3, "in_channels_list": [32, 64], "out_channels": 6},
+        ((7, 32, 16, 32, 64), (7, 64, 8, 16, 32)),
+        ((7, 6, 16, 32, 64), (7, 6, 8, 16, 32)),
+    ],
+    [
+        {"spatial_dims": 2, "in_channels_list": [32, 64], "out_channels": 6},
+        ((7, 32, 16, 32), (7, 64, 8, 16)),
+        ((7, 6, 16, 32), (7, 6, 8, 16)),
+    ],
+]
+
+TEST_CASES2 = [
+    [{"spatial_dims": 3, "returned_layers": [1]}, (7, 3, 32, 64, 32), ((7, 256, 16, 32, 16), (7, 256, 8, 16, 8))]
+]
+
+
+class TestFPNBlock(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_fpn_block(self, input_param, input_shape, expected_shape):
+        net = FeaturePyramidNetwork(**input_param)
+        data = OrderedDict()
+        data["feat0"] = torch.rand(input_shape[0])
+        data["feat1"] = torch.rand(input_shape[1])
+        result = net(data)
+        self.assertEqual(result["feat0"].shape, expected_shape[0])
+        self.assertEqual(result["feat1"].shape, expected_shape[1])
+
+    @parameterized.expand(TEST_CASES)
+    @SkipIfBeforePyTorchVersion((1, 9, 1))
+    def test_script(self, input_param, input_shape, expected_shape):
+        # test whether support torchscript
+        net = FeaturePyramidNetwork(**input_param)
+        data = OrderedDict()
+        data["feat0"] = torch.rand(input_shape[0])
+        data["feat1"] = torch.rand(input_shape[1])
+        test_script_save(net, data)
+
+
+@unittest.skipUnless(has_torchvision, "Requires torchvision")
+class TestFPN(unittest.TestCase):
+    @parameterized.expand(TEST_CASES2)
+    def test_fpn(self, input_param, input_shape, expected_shape):
+        net = _resnet_fpn_extractor(backbone=resnet50(), spatial_dims=input_param["spatial_dims"], returned_layers=[1])
+        data = torch.rand(input_shape)
+        result = net(data)
+        self.assertEqual(result["0"].shape, expected_shape[0])
+        self.assertEqual(result["pool"].shape, expected_shape[1])
+
+    @parameterized.expand(TEST_CASES2)
+    @SkipIfBeforePyTorchVersion((1, 9, 1))
+    def test_script(self, input_param, input_shape, expected_shape):
+        # test whether support torchscript
+        net = _resnet_fpn_extractor(backbone=resnet50(), spatial_dims=input_param["spatial_dims"], returned_layers=[1])
+        data = torch.rand(input_shape)
+        test_script_save(net, data)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/blocks/test_localnet_block.py

@@ -0,0 +1,121 @@
+# 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 unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks import eval_mode
+from monai.networks.blocks.localnet_block import (
+    LocalNetDownSampleBlock,
+    LocalNetFeatureExtractorBlock,
+    LocalNetUpSampleBlock,
+)
+
+TEST_CASE_DOWN_SAMPLE = [
+    [{"spatial_dims": spatial_dims, "in_channels": 2, "out_channels": 4, "kernel_size": 3}] for spatial_dims in [2, 3]
+]
+
+TEST_CASE_UP_SAMPLE = [
+    [
+        {
+            "spatial_dims": spatial_dims,
+            "in_channels": 4,
+            "out_channels": 2,
+            "mode": "bilinear" if spatial_dims == 2 else "trilinear",
+        }
+    ]
+    for spatial_dims in [2, 3]
+]
+
+TEST_CASE_EXTRACT = [
+    [{"spatial_dims": spatial_dims, "in_channels": 2, "out_channels": 3, "act": act, "initializer": initializer}]
+    for spatial_dims, act, initializer in zip([2, 3], ["sigmoid", None], ["kaiming_uniform", "zeros"])
+]
+
+in_size = 4
+
+
+class TestLocalNetDownSampleBlock(unittest.TestCase):
+
+    @parameterized.expand(TEST_CASE_DOWN_SAMPLE)
+    def test_shape(self, input_param):
+        net = LocalNetDownSampleBlock(**input_param)
+        input_shape = (1, input_param["in_channels"], *([in_size] * input_param["spatial_dims"]))
+        expect_mid_shape = (1, input_param["out_channels"], *([in_size] * input_param["spatial_dims"]))
+        expect_x_shape = (1, input_param["out_channels"], *([in_size / 2] * input_param["spatial_dims"]))
+        with eval_mode(net):
+            x, mid = net(torch.randn(input_shape))
+            self.assertEqual(x.shape, expect_x_shape)
+            self.assertEqual(mid.shape, expect_mid_shape)
+
+    def test_ill_arg(self):
+        # even kernel_size
+        with self.assertRaises(NotImplementedError):
+            LocalNetDownSampleBlock(spatial_dims=2, in_channels=2, out_channels=4, kernel_size=4)
+
+    @parameterized.expand(TEST_CASE_DOWN_SAMPLE)
+    def test_ill_shape(self, input_param):
+        net = LocalNetDownSampleBlock(**input_param)
+        input_shape = (1, input_param["in_channels"], *([5] * input_param["spatial_dims"]))
+        with self.assertRaises(ValueError):
+            with eval_mode(net):
+                net(torch.randn(input_shape))
+
+
+class TestLocalNetUpSampleBlock(unittest.TestCase):
+
+    @parameterized.expand(TEST_CASE_UP_SAMPLE)
+    def test_shape(self, input_param):
+        net = LocalNetUpSampleBlock(**input_param)
+        input_shape = (1, input_param["in_channels"], *([in_size] * input_param["spatial_dims"]))
+        mid_shape = (1, input_param["out_channels"], *([in_size * 2] * input_param["spatial_dims"]))
+        expected_shape = mid_shape
+        with eval_mode(net):
+            result = net(torch.randn(input_shape), torch.randn(mid_shape))
+            self.assertEqual(result.shape, expected_shape)
+
+    def test_ill_arg(self):
+        # channel unmatch
+        with self.assertRaises(ValueError):
+            LocalNetUpSampleBlock(spatial_dims=2, in_channels=2, out_channels=2)
+
+    @parameterized.expand(TEST_CASE_UP_SAMPLE)
+    def test_ill_shape(self, input_param):
+        net = LocalNetUpSampleBlock(**input_param)
+        input_shape = (1, input_param["in_channels"], *([in_size] * input_param["spatial_dims"]))
+        mid_shape = (1, input_param["out_channels"], *([in_size] * input_param["spatial_dims"]))
+        with self.assertRaises(ValueError):
+            with eval_mode(net):
+                net(torch.randn(input_shape), torch.randn(mid_shape))
+
+
+class TestExtractBlock(unittest.TestCase):
+
+    @parameterized.expand(TEST_CASE_EXTRACT)
+    def test_shape(self, input_param):
+        net = LocalNetFeatureExtractorBlock(**input_param)
+        input_shape = (1, input_param["in_channels"], *([in_size] * input_param["spatial_dims"]))
+        expected_shape = (1, input_param["out_channels"], *([in_size] * input_param["spatial_dims"]))
+        with eval_mode(net):
+            result = net(torch.randn(input_shape))
+            self.assertEqual(result.shape, expected_shape)
+
+    def test_ill_arg(self):
+        with self.assertRaises(ValueError):
+            LocalNetFeatureExtractorBlock(spatial_dims=2, in_channels=2, out_channels=2, initializer="none")
+
+
+if __name__ == "__main__":
+    unittest.main()