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

tests/networks/schedulers/test_scheduler_ddpm.py

@@ -0,0 +1,104 @@
+# 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.schedulers import DDPMScheduler
+from tests.test_utils import assert_allclose
+
+TEST_2D_CASE = []
+for beta_schedule in ["linear_beta", "scaled_linear_beta"]:
+    for variance_type in ["fixed_small", "fixed_large"]:
+        TEST_2D_CASE.append(
+            [{"schedule": beta_schedule, "variance_type": variance_type}, (2, 6, 16, 16), (2, 6, 16, 16)]
+        )
+
+TEST_3D_CASE = []
+for beta_schedule in ["linear_beta", "scaled_linear_beta"]:
+    for variance_type in ["fixed_small", "fixed_large"]:
+        TEST_3D_CASE.append(
+            [{"schedule": beta_schedule, "variance_type": variance_type}, (2, 6, 16, 16, 16), (2, 6, 16, 16, 16)]
+        )
+
+TEST_CASES = TEST_2D_CASE + TEST_3D_CASE
+
+TEST_FULl_LOOP = [
+    [{"schedule": "linear_beta"}, (1, 1, 2, 2), torch.Tensor([[[[-1.0153, -0.3218], [0.8454, -0.7870]]]])]
+]
+
+
+class TestDDPMScheduler(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_add_noise(self, input_param, input_shape, expected_shape):
+        scheduler = DDPMScheduler(**input_param)
+        original_sample = torch.zeros(input_shape)
+        noise = torch.randn_like(original_sample)
+        timesteps = torch.randint(0, scheduler.num_train_timesteps, (original_sample.shape[0],)).long()
+
+        noisy = scheduler.add_noise(original_samples=original_sample, noise=noise, timesteps=timesteps)
+        self.assertEqual(noisy.shape, expected_shape)
+
+    @parameterized.expand(TEST_CASES)
+    def test_step_shape(self, input_param, input_shape, expected_shape):
+        scheduler = DDPMScheduler(**input_param)
+        model_output = torch.randn(input_shape)
+        sample = torch.randn(input_shape)
+        output_step = scheduler.step(model_output=model_output, timestep=500, sample=sample)
+        self.assertEqual(output_step[0].shape, expected_shape)
+        self.assertEqual(output_step[1].shape, expected_shape)
+
+    @parameterized.expand(TEST_FULl_LOOP)
+    def test_full_timestep_loop(self, input_param, input_shape, expected_output):
+        scheduler = DDPMScheduler(**input_param)
+        scheduler.set_timesteps(50)
+        torch.manual_seed(42)
+        model_output = torch.randn(input_shape)
+        sample = torch.randn(input_shape)
+        for t in range(50):
+            sample, _ = scheduler.step(model_output=model_output, timestep=t, sample=sample)
+        assert_allclose(sample, expected_output, rtol=1e-3, atol=1e-3)
+
+    @parameterized.expand(TEST_CASES)
+    def test_get_velocity_shape(self, input_param, input_shape, expected_shape):
+        scheduler = DDPMScheduler(**input_param)
+        sample = torch.randn(input_shape)
+        timesteps = torch.randint(0, scheduler.num_train_timesteps, (input_shape[0],)).long()
+        velocity = scheduler.get_velocity(sample=sample, noise=sample, timesteps=timesteps)
+        self.assertEqual(velocity.shape, expected_shape)
+
+    def test_step_learned(self):
+        for variance_type in ["learned", "learned_range"]:
+            scheduler = DDPMScheduler(variance_type=variance_type)
+        model_output = torch.randn(2, 6, 16, 16)
+        sample = torch.randn(2, 3, 16, 16)
+        output_step = scheduler.step(model_output=model_output, timestep=500, sample=sample)
+        self.assertEqual(output_step[0].shape, sample.shape)
+        self.assertEqual(output_step[1].shape, sample.shape)
+
+    def test_set_timesteps(self):
+        scheduler = DDPMScheduler(num_train_timesteps=1000)
+        scheduler.set_timesteps(num_inference_steps=100)
+        self.assertEqual(scheduler.num_inference_steps, 100)
+        self.assertEqual(len(scheduler.timesteps), 100)
+
+    def test_set_timesteps_with_num_inference_steps_bigger_than_num_train_timesteps(self):
+        scheduler = DDPMScheduler(num_train_timesteps=1000)
+        with self.assertRaises(ValueError):
+            scheduler.set_timesteps(num_inference_steps=2000)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/schedulers/test_scheduler_pndm.py

@@ -0,0 +1,108 @@
+# 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.schedulers import PNDMScheduler
+from tests.test_utils import assert_allclose
+
+TEST_2D_CASE = []
+for beta_schedule in ["linear_beta", "scaled_linear_beta"]:
+    TEST_2D_CASE.append([{"schedule": beta_schedule}, (2, 6, 16, 16), (2, 6, 16, 16)])
+
+TEST_3D_CASE = []
+for beta_schedule in ["linear_beta", "scaled_linear_beta"]:
+    TEST_3D_CASE.append([{"schedule": beta_schedule}, (2, 6, 16, 16, 16), (2, 6, 16, 16, 16)])
+
+TEST_CASES = TEST_2D_CASE + TEST_3D_CASE
+
+TEST_FULl_LOOP = [
+    [
+        {"schedule": "linear_beta"},
+        (1, 1, 2, 2),
+        torch.Tensor([[[[-2123055.2500, -459014.2812], [2863438.0000, -1263401.7500]]]]),
+    ]
+]
+
+
+class TestDDPMScheduler(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_add_noise(self, input_param, input_shape, expected_shape):
+        scheduler = PNDMScheduler(**input_param)
+        original_sample = torch.zeros(input_shape)
+        noise = torch.randn_like(original_sample)
+        timesteps = torch.randint(0, scheduler.num_train_timesteps, (original_sample.shape[0],)).long()
+        noisy = scheduler.add_noise(original_samples=original_sample, noise=noise, timesteps=timesteps)
+        self.assertEqual(noisy.shape, expected_shape)
+
+    @parameterized.expand(TEST_CASES)
+    def test_step_shape(self, input_param, input_shape, expected_shape):
+        scheduler = PNDMScheduler(**input_param)
+        scheduler.set_timesteps(600)
+        model_output = torch.randn(input_shape)
+        sample = torch.randn(input_shape)
+        output_step = scheduler.step(model_output=model_output, timestep=500, sample=sample)
+        self.assertEqual(output_step[0].shape, expected_shape)
+        self.assertEqual(output_step[1], None)
+
+    @parameterized.expand(TEST_FULl_LOOP)
+    def test_full_timestep_loop(self, input_param, input_shape, expected_output):
+        scheduler = PNDMScheduler(**input_param)
+        scheduler.set_timesteps(50)
+        torch.manual_seed(42)
+        model_output = torch.randn(input_shape)
+        sample = torch.randn(input_shape)
+        for t in range(50):
+            sample, _ = scheduler.step(model_output=model_output, timestep=t, sample=sample)
+        assert_allclose(sample, expected_output, rtol=1e-3, atol=1e-3)
+
+    @parameterized.expand(TEST_FULl_LOOP)
+    def test_timestep_two_loops(self, input_param, input_shape, expected_output):
+        scheduler = PNDMScheduler(**input_param)
+        scheduler.set_timesteps(50)
+        torch.manual_seed(42)
+        model_output = torch.randn(input_shape)
+        sample = torch.randn(input_shape)
+        for t in range(50):
+            sample, _ = scheduler.step(model_output=model_output, timestep=t, sample=sample)
+        torch.manual_seed(42)
+        model_output2 = torch.randn(input_shape)
+        sample2 = torch.randn(input_shape)
+        scheduler.set_timesteps(50)
+        for t in range(50):
+            sample2, _ = scheduler.step(model_output=model_output2, timestep=t, sample=sample2)
+        assert_allclose(sample, sample2, rtol=1e-3, atol=1e-3)
+
+    def test_set_timesteps(self):
+        scheduler = PNDMScheduler(num_train_timesteps=1000, skip_prk_steps=True)
+        scheduler.set_timesteps(num_inference_steps=100)
+        self.assertEqual(scheduler.num_inference_steps, 100)
+        self.assertEqual(len(scheduler.timesteps), 100)
+
+    def test_set_timesteps_prk(self):
+        scheduler = PNDMScheduler(num_train_timesteps=1000, skip_prk_steps=False)
+        scheduler.set_timesteps(num_inference_steps=100)
+        self.assertEqual(scheduler.num_inference_steps, 109)
+        self.assertEqual(len(scheduler.timesteps), 109)
+
+    def test_set_timesteps_with_num_inference_steps_bigger_than_num_train_timesteps(self):
+        scheduler = PNDMScheduler(num_train_timesteps=1000)
+        with self.assertRaises(ValueError):
+            scheduler.set_timesteps(num_inference_steps=2000)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_bundle_onnx_export.py

@@ -0,0 +1,71 @@
+# 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 os
+import tempfile
+import unittest
+from pathlib import Path
+
+from parameterized import parameterized
+
+from monai.bundle import ConfigParser
+from monai.networks import save_state
+from tests.test_utils import SkipIfBeforePyTorchVersion, SkipIfNoModule, command_line_tests, skip_if_windows
+
+TEST_CASE_1 = ["True"]
+TEST_CASE_2 = ["False"]
+
+
+@skip_if_windows
+@SkipIfNoModule("onnx")
+@SkipIfBeforePyTorchVersion((1, 10))
+class TestONNXExport(unittest.TestCase):
+    def setUp(self):
+        self.device = os.environ.get("CUDA_VISIBLE_DEVICES")
+        if not self.device:
+            os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # default
+
+    def tearDown(self):
+        if self.device is not None:
+            os.environ["CUDA_VISIBLE_DEVICES"] = self.device
+        else:
+            del os.environ["CUDA_VISIBLE_DEVICES"]  # previously unset
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_onnx_export(self, use_trace):
+        tests_path = Path(__file__).parents[1]
+        meta_file = os.path.join(tests_path, "testing_data", "metadata.json")
+        config_file = os.path.join(tests_path, "testing_data", "inference.json")
+        with tempfile.TemporaryDirectory() as tempdir:
+            def_args = {"meta_file": "will be replaced by `meta_file` arg"}
+            def_args_file = os.path.join(tempdir, "def_args.yaml")
+
+            ckpt_file = os.path.join(tempdir, "model.pt")
+            onnx_file = os.path.join(tempdir, "model.onnx")
+
+            parser = ConfigParser()
+            parser.export_config_file(config=def_args, filepath=def_args_file)
+            parser.read_config(config_file)
+            net = parser.get_parsed_content("network_def")
+            save_state(src=net, path=ckpt_file)
+
+            cmd = ["python", "-m", "monai.bundle", "onnx_export", "network_def", "--filepath", onnx_file]
+            cmd += ["--meta_file", meta_file, "--config_file", f"['{config_file}','{def_args_file}']"]
+            cmd += ["--ckpt_file", ckpt_file, "--args_file", def_args_file, "--input_shape", "[1, 1, 96, 96, 96]"]
+            cmd += ["--use_trace", use_trace]
+            command_line_tests(cmd)
+            self.assertTrue(os.path.exists(onnx_file))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_convert_to_onnx.py

@@ -0,0 +1,106 @@
+# 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 itertools
+import platform
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks import convert_to_onnx
+from monai.networks.nets import SegResNet, UNet
+from tests.test_utils import SkipIfBeforePyTorchVersion, SkipIfNoModule, optional_import, skip_if_quick
+
+if torch.cuda.is_available():
+    TORCH_DEVICE_OPTIONS = ["cpu", "cuda"]
+else:
+    TORCH_DEVICE_OPTIONS = ["cpu"]
+TESTS = list(itertools.product(TORCH_DEVICE_OPTIONS, [True, False], [True, False]))
+TESTS_ORT = list(itertools.product(TORCH_DEVICE_OPTIONS, [True]))
+
+ON_AARCH64 = platform.machine() == "aarch64"
+if ON_AARCH64:
+    rtol, atol = 1e-1, 1e-2
+else:
+    rtol, atol = 1e-3, 1e-4
+
+onnx, _ = optional_import("onnx")
+
+
+@SkipIfNoModule("onnx")
+@SkipIfBeforePyTorchVersion((1, 9))
+@skip_if_quick
+class TestConvertToOnnx(unittest.TestCase):
+    @parameterized.expand(TESTS)
+    def test_unet(self, device, use_trace, use_ort):
+        if use_ort:
+            _, has_onnxruntime = optional_import("onnxruntime")
+            if not has_onnxruntime:
+                self.skipTest("onnxruntime is not installed probably due to python version >= 3.11.")
+        model = UNet(
+            spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0
+        )
+        if use_trace:
+            onnx_model = convert_to_onnx(
+                model=model,
+                inputs=[torch.randn((16, 1, 32, 32), requires_grad=False)],
+                input_names=["x"],
+                output_names=["y"],
+                verify=True,
+                device=device,
+                use_ort=use_ort,
+                use_trace=use_trace,
+                rtol=rtol,
+                atol=atol,
+            )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+    @parameterized.expand(TESTS_ORT)
+    @SkipIfBeforePyTorchVersion((1, 12))
+    def test_seg_res_net(self, device, use_ort):
+        if use_ort:
+            _, has_onnxruntime = optional_import("onnxruntime")
+            if not has_onnxruntime:
+                self.skipTest("onnxruntime is not installed probably due to python version >= 3.11.")
+        model = SegResNet(
+            spatial_dims=3,
+            init_filters=32,
+            in_channels=1,
+            out_channels=105,
+            dropout_prob=0.2,
+            act=("RELU", {"inplace": True}),
+            norm=("GROUP", {"num_groups": 8}),
+            norm_name="",
+            num_groups=8,
+            use_conv_final=True,
+            blocks_down=[1, 2, 2, 4],
+            blocks_up=[1, 1, 1],
+        )
+        onnx_model = convert_to_onnx(
+            model=model,
+            inputs=[torch.randn((1, 1, 24, 24, 24), requires_grad=False)],
+            input_names=["x"],
+            output_names=["y"],
+            verify=True,
+            device=device,
+            use_ort=use_ort,
+            use_trace=True,
+            rtol=rtol,
+            atol=atol,
+        )
+        self.assertTrue(isinstance(onnx_model, onnx.ModelProto))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_convert_to_torchscript.py

@@ -0,0 +1,46 @@
+# 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 os
+import tempfile
+import unittest
+
+import torch
+
+from monai.networks import convert_to_torchscript
+from monai.networks.nets import UNet
+
+
+class TestConvertToTorchScript(unittest.TestCase):
+
+    def test_value(self):
+        model = UNet(
+            spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0
+        )
+        with tempfile.TemporaryDirectory() as tempdir:
+            torchscript_model = convert_to_torchscript(
+                model=model,
+                filename_or_obj=os.path.join(tempdir, "model.ts"),
+                extra_files={"foo.txt": b"bar"},
+                verify=True,
+                inputs=[torch.randn((16, 1, 32, 32), requires_grad=False)],
+                device="cuda" if torch.cuda.is_available() else "cpu",
+                rtol=1e-3,
+                atol=1e-4,
+                optimize=None,
+            )
+            self.assertTrue(isinstance(torchscript_model, torch.nn.Module))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_convert_to_trt.py

@@ -0,0 +1,79 @@
+# 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 tempfile
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.networks import convert_to_trt
+from monai.networks.nets import UNet
+from monai.utils import optional_import
+from tests.test_utils import SkipIfBeforeComputeCapabilityVersion, skip_if_no_cuda, skip_if_quick, skip_if_windows
+
+_, has_torchtrt = optional_import(
+    "torch_tensorrt",
+    version="1.4.0",
+    descriptor="Torch-TRT is not installed. Are you sure you have a Torch-TensorRT compilation?",
+)
+_, has_tensorrt = optional_import(
+    "tensorrt", descriptor="TensorRT is not installed. Are you sure you have a TensorRT compilation?"
+)
+
+TEST_CASE_1 = ["fp32"]
+TEST_CASE_2 = ["fp16"]
+
+
+@skip_if_windows
+@skip_if_no_cuda
+@skip_if_quick
+@SkipIfBeforeComputeCapabilityVersion((7, 5))
+class TestConvertToTRT(unittest.TestCase):
+    def setUp(self):
+        self.gpu_device = torch.cuda.current_device()
+
+    def tearDown(self):
+        current_device = torch.cuda.current_device()
+        if current_device != self.gpu_device:
+            torch.cuda.set_device(self.gpu_device)
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    @unittest.skipUnless(has_torchtrt and has_tensorrt, "Torch-TensorRT is required for convert!")
+    def test_value(self, precision):
+        model = UNet(
+            spatial_dims=3,
+            in_channels=1,
+            out_channels=2,
+            channels=(2, 2, 4, 8, 4),
+            strides=(2, 2, 2, 2),
+            num_res_units=2,
+            norm="batch",
+        )
+        with tempfile.TemporaryDirectory() as _:
+            torchscript_model = convert_to_trt(
+                model=model,
+                precision=precision,
+                input_shape=[1, 1, 96, 96, 96],
+                dynamic_batchsize=[1, 4, 8],
+                use_trace=False,
+                verify=True,
+                device=0,
+                rtol=1e-2,
+                atol=1e-2,
+            )
+            self.assertTrue(isinstance(torchscript_model, torch.nn.Module))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_save_state.py

@@ -0,0 +1,73 @@
+# 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 os
+import tempfile
+import unittest
+
+import torch
+import torch.optim as optim
+from parameterized import parameterized
+
+from monai.networks import save_state
+
+TEST_CASE_1 = [torch.nn.PReLU(), ["weight"]]
+
+TEST_CASE_2 = [{"net": torch.nn.PReLU()}, ["net"]]
+
+TEST_CASE_3 = [{"net": torch.nn.PReLU(), "opt": optim.SGD(torch.nn.PReLU().parameters(), lr=0.02)}, ["net", "opt"]]
+
+TEST_CASE_4 = [torch.nn.DataParallel(torch.nn.PReLU()), ["weight"]]
+
+TEST_CASE_5 = [{"net": torch.nn.DataParallel(torch.nn.PReLU())}, ["net"]]
+
+TEST_CASE_6 = [torch.nn.PReLU(), ["weight"], True, True, None, {"pickle_protocol": 2}]
+
+TEST_CASE_7 = [torch.nn.PReLU().state_dict(), ["weight"]]
+
+TEST_CASE_8 = [torch.nn.PReLU(), ["weight"], False]
+
+TEST_CASE_9 = [torch.nn.PReLU(), ["weight"], True, False]
+
+TEST_CASE_10 = [torch.nn.PReLU(), ["weight"], True, True, torch.save]
+
+
+class TestSaveState(unittest.TestCase):
+
+    @parameterized.expand(
+        [
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+            TEST_CASE_8,
+            TEST_CASE_9,
+            TEST_CASE_10,
+        ]
+    )
+    def test_file(self, src, expected_keys, create_dir=True, atomic=True, func=None, kwargs=None):
+        with tempfile.TemporaryDirectory() as tempdir:
+            path = os.path.join(tempdir, "test_ckpt.pt")
+            if kwargs is None:
+                kwargs = {}
+            save_state(src=src, path=path, create_dir=create_dir, atomic=atomic, func=func, **kwargs)
+            ckpt = dict(torch.load(path))
+            for k in ckpt.keys():
+                self.assertIn(k, expected_keys)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/networks/test_to_onehot.py

@@ -0,0 +1,63 @@
+# 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 numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.networks import one_hot
+
+TEST_CASE_1 = [  # single channel 2D, batch 3, shape (2, 1, 2, 2)
+    {"labels": torch.tensor([[[[0, 1], [1, 2]]], [[[2, 1], [1, 0]]]]), "num_classes": 3},
+    (2, 3, 2, 2),
+]
+
+TEST_CASE_2 = [  # single channel 1D, batch 2, shape (2, 1, 4)
+    {"labels": torch.tensor([[[1, 2, 2, 0]], [[2, 1, 0, 1]]]), "num_classes": 3},
+    (2, 3, 4),
+    np.array([[[0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 1, 0]], [[0, 0, 1, 0], [0, 1, 0, 1], [1, 0, 0, 0]]]),
+]
+
+TEST_CASE_3 = [  # single channel 0D, batch 2, shape (2, 1)
+    {"labels": torch.tensor([[1.0], [2.0]]), "num_classes": 3},
+    (2, 3),
+    np.array([[0, 1, 0], [0, 0, 1]]),
+]
+
+TEST_CASE_4 = [  # no channel 0D, batch 3, shape (3)
+    {"labels": torch.tensor([1, 2, 0]), "num_classes": 3, "dtype": torch.long},
+    (3, 3),
+    np.array([[0, 1, 0], [0, 0, 1], [1, 0, 0]]),
+]
+
+
+class TestToOneHot(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_shape(self, input_data, expected_shape, expected_result=None):
+        result = one_hot(**input_data)
+        self.assertEqual(result.shape, expected_shape)
+        if expected_result is not None:
+            self.assertTrue(np.allclose(expected_result, result.numpy()))
+
+        if "dtype" in input_data:
+            self.assertEqual(result.dtype, input_data["dtype"])
+        else:
+            # by default, expecting float type
+            self.assertEqual(result.dtype, torch.float)
+
+
+if __name__ == "__main__":
+    unittest.main()