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

tests/integration/test_seg_loss_integration.py

@@ -0,0 +1,159 @@
+# 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
+import torch.nn as nn
+import torch.optim as optim
+from parameterized import parameterized
+
+from monai.losses import DiceLoss, FocalLoss, GeneralizedDiceLoss, TverskyLoss
+from monai.networks import one_hot
+
+TEST_CASES = [
+    [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "smooth_nr": 1e-4, "smooth_dr": 1e-4}, {}],
+    [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "smooth_nr": 0, "smooth_dr": 1e-3}, {}],
+    [DiceLoss, {"to_onehot_y": False, "squared_pred": True, "smooth_nr": 0, "smooth_dr": 1e-3}, {}],
+    [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "batch": True}, {}],
+    [DiceLoss, {"to_onehot_y": True, "sigmoid": True}, {}],
+    [DiceLoss, {"to_onehot_y": True, "softmax": True}, {}],
+    [FocalLoss, {"to_onehot_y": True, "gamma": 1.5, "weight": torch.tensor([1, 2])}, {}],
+    [FocalLoss, {"to_onehot_y": False, "gamma": 1.5, "weight": [1, 2]}, {}],
+    [FocalLoss, {"to_onehot_y": False, "gamma": 1.5, "weight": 1.0}, {}],
+    [FocalLoss, {"to_onehot_y": True, "gamma": 1.5}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": True, "softmax": True}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "simple"}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "uniform"}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "uniform", "batch": True}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": False, "sigmoid": True, "w_type": "uniform", "batch": True}, {}],
+    [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 0.8, "beta": 0.2}, {}],
+    [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 0.8, "beta": 0.2, "batch": True}, {}],
+    [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 1.0, "beta": 0.0}, {}],
+    [TverskyLoss, {"to_onehot_y": False, "softmax": True, "alpha": 1.0, "beta": 0.0}, {}],
+]
+
+
+class TestSegLossIntegration(unittest.TestCase):
+
+    def setUp(self):
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        torch.manual_seed(0)
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
+
+    def tearDown(self):
+        torch.backends.cudnn.deterministic = False
+        torch.backends.cudnn.benchmark = True
+
+    @parameterized.expand(TEST_CASES)
+    def test_convergence(self, loss_type, loss_args, forward_args):
+        """
+        The goal of this test is to assess if the gradient of the loss function
+        is correct by testing if we can train a one layer neural network
+        to segment one image.
+        We verify that the loss is decreasing in almost all SGD steps.
+        """
+        learning_rate = 0.001
+        max_iter = 40
+
+        # define a simple 3d example
+        target_seg = torch.tensor(
+            [
+                [
+                    # raw 0
+                    [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                    # raw 1
+                    [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                    # raw 2
+                    [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                ]
+            ],
+            device=self.device,
+        )
+        target_seg = torch.unsqueeze(target_seg, dim=0)
+        image = 12 * target_seg + 27
+        image = image.float().to(self.device)
+        num_classes = 2
+        num_voxels = 3 * 4 * 4
+
+        target_onehot = one_hot(target_seg, num_classes=num_classes)
+
+        # define a one layer model
+        class OnelayerNet(nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.layer_1 = nn.Linear(num_voxels, 200)
+                self.acti = nn.ReLU()
+                self.layer_2 = nn.Linear(200, num_voxels * num_classes)
+
+            def forward(self, x):
+                x = x.view(-1, num_voxels)
+                x = self.layer_1(x)
+                x = self.acti(x)
+                x = self.layer_2(x)
+                x = x.view(-1, num_classes, 3, 4, 4)
+                return x
+
+        # initialise the network
+        net = OnelayerNet().to(self.device)
+
+        # initialize the loss
+        loss = loss_type(**loss_args)
+
+        # initialize a SGD optimizer
+        optimizer = optim.Adam(net.parameters(), lr=learning_rate)
+
+        loss_history = []
+        init_output = None
+
+        # train the network
+        for iter_i in range(max_iter):
+            # set the gradient to zero
+            optimizer.zero_grad()
+
+            # forward pass
+            output = net(image)
+            if init_output is None:
+                init_output = torch.argmax(output, 1).detach().cpu().numpy()
+
+            if loss_args["to_onehot_y"] is False:
+                loss_val = loss(output, target_onehot, **forward_args)
+            else:
+                loss_val = loss(output, target_seg, **forward_args)
+
+            if iter_i % 10 == 0:
+                pred = torch.argmax(output, 1).detach().cpu().numpy()
+                gt = target_seg.detach().cpu().numpy()[:, 0]
+                print(f"{loss_type.__name__} iter: {iter_i}, acc: {np.sum(pred == gt) / np.prod(pred.shape)}")
+
+            # backward pass
+            loss_val.backward()
+            optimizer.step()
+
+            # stats
+            loss_history.append(loss_val.item())
+
+        pred = torch.argmax(output, 1).detach().cpu().numpy()
+        target = target_seg.detach().cpu().numpy()[:, 0]
+        # initial predictions are bad
+        self.assertTrue(not np.allclose(init_output, target))
+        # final predictions are good
+        np.testing.assert_allclose(pred, target)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_spatial_combine_transforms.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
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+import monai.transforms as mt
+from monai.data import create_test_image_2d, create_test_image_3d
+from monai.data.meta_tensor import MetaTensor
+from monai.transforms.lazy.functional import apply_pending
+from monai.transforms.transform import MapTransform
+from monai.utils import set_determinism
+from tests.lazy_transforms_utils import get_apply_param
+from tests.test_utils import assert_allclose
+
+TEST_2D = [
+    [
+        (2, 62, 61),
+        [
+            (mt.Spacing, {"pixdim": (1.2, 1.5), "padding_mode": "zeros", "dtype": torch.float32}),
+            (mt.Orientation, {"axcodes": "RA"}),
+            (mt.Resize, {"spatial_size": (64, 48), "mode": "bilinear"}),
+            (mt.RandSpatialCrop, {"roi_size": (32, 32)}),
+            (
+                mt.RandAffine,
+                {
+                    "prob": 0.9,
+                    "rotate_range": (np.pi / 2,),
+                    "shear_range": [1, 2],
+                    "translate_range": [2, 1],
+                    "mode": "bilinear",
+                    "padding_mode": "reflection",
+                },
+            ),
+            (mt.RandFlip, {"prob": 0.9}),
+            (mt.RandRotate, {"prob": 0.9, "range_x": np.pi / 4, "mode": "bilinear", "padding_mode": "reflection"}),
+            (mt.CenterScaleCrop, {"roi_scale": (0.96, 0.8)}),
+            (mt.RandZoom, {"prob": 0.9, "mode": "bilinear", "keep_size": False, "align_corners": False}),
+        ],
+    ],
+    [
+        (2, 63, 64),
+        [
+            (mt.CenterScaleCropd, {"roi_scale": (0.96, 0.8), "keys": "img"}),
+            (mt.RandRotated, {"prob": 0.9, "range_x": np.pi / 4, "keys": "img"}),
+            (
+                mt.RandZoomd,
+                {"prob": 0.9, "mode": "bilinear", "keep_size": False, "keys": "img", "align_corners": False},
+            ),
+            (mt.Spacingd, {"pixdim": (1.2, 1.5), "padding_mode": "zeros", "dtype": torch.float32, "keys": "img"}),
+            (mt.RandFlipd, {"prob": 0.9, "keys": "img"}),
+            (
+                mt.RandAffined,
+                {
+                    "prob": 0.9,
+                    "rotate_range": (np.pi / 2,),
+                    "shear_range": [1, 2],
+                    "translate_range": [2, 1],
+                    "mode": "bilinear",
+                    "keys": "img",
+                },
+            ),
+            (mt.Orientationd, {"axcodes": "RA", "keys": "img"}),
+            (mt.Resized, {"spatial_size": (48, 48), "mode": "bilinear", "keys": "img"}),
+            (mt.RandScaleCropd, {"roi_scale": (0.4, 1.5), "random_size": False, "keys": "img"}),
+        ],
+    ],
+]
+
+TEST_3D = [
+    [
+        (2, 83, 100, 67),
+        [
+            (mt.Orientation, {"axcodes": "RAS"}),
+            (mt.CenterScaleCrop, {"roi_scale": (1.2, 0.8, 1.0)}),
+            (
+                mt.RandAffine,
+                {
+                    "prob": 0.9,
+                    "rotate_range": (np.pi / 2,),
+                    "shear_range": [1, 2],
+                    "translate_range": [2, 1],
+                    "mode": "bilinear",
+                },
+            ),
+            (mt.Spacing, {"pixdim": (0.9, 1.2, 1.0), "padding_mode": "zeros", "dtype": torch.float32}),
+            (mt.RandSpatialCrop, {"roi_size": (36, 36, 38), "random_size": False}),
+            (mt.RandZoom, {"prob": 0.9, "mode": "nearest", "keep_size": False}),
+            (mt.Resize, {"spatial_size": (32, 32, 32), "mode": "nearest"}),
+            (mt.RandFlip, {"prob": 0.9}),
+            (mt.RandRotate, {"prob": 0.9, "range_x": np.pi / 4}),
+        ],
+    ],
+    [
+        (2, 62, 64, 72),
+        [
+            (mt.RandScaleCropd, {"roi_scale": (0.9, 0.7, 1.1), "random_size": False, "keys": "img"}),
+            (mt.Spacingd, {"pixdim": (1.2, 1.5, 0.9), "padding_mode": "zeros", "dtype": torch.float32, "keys": "img"}),
+            (mt.Orientationd, {"axcodes": "RAS", "keys": "img"}),
+            (mt.Resized, {"spatial_size": (32, 32, 32), "mode": "nearest", "keys": "img"}),
+            (mt.RandFlipd, {"prob": 0.9, "keys": "img"}),
+            (mt.CenterScaleCropd, {"roi_scale": (0.96, 0.8, 1.25), "keys": "img"}),
+            (mt.RandZoomd, {"prob": 0.9, "mode": "nearest", "keep_size": False, "keys": "img"}),
+            (
+                mt.RandAffined,
+                {
+                    "prob": 0.9,
+                    "rotate_range": (np.pi / 2,),
+                    "shear_range": [1, 2],
+                    "translate_range": [2, 1],
+                    "mode": "bilinear",
+                    "keys": "img",
+                },
+            ),
+            (mt.RandRotated, {"prob": 0.9, "range_x": np.pi / 4, "keys": "img"}),
+        ],
+    ],
+]
+
+
+class CombineLazyTest(unittest.TestCase):
+    @parameterized.expand(TEST_2D + TEST_3D)
+    def test_combine_transforms(self, input_shape, funcs):
+        for device in ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]:
+            for seed in [10, 100, 1000, 10000]:
+                set_determinism(seed=seed)
+                _funcs = []
+                for _func, _params in funcs:
+                    _funcs.append(_func(**_params))
+                is_map = isinstance(_funcs[0], MapTransform)
+                chns, sp_size = input_shape[0], input_shape[1:]
+                imgs = []
+                for _ in range(chns):
+                    if len(sp_size) == 2:
+                        imgs.append(create_test_image_2d(sp_size[0], sp_size[1])[0])
+                    else:
+                        imgs.append(create_test_image_3d(sp_size[0], sp_size[1], sp_size[2])[0])
+                data = np.stack(imgs).astype(float)
+                im = MetaTensor(data, meta={"a": "b", "affine": np.eye(len(input_shape))}).to(device)
+                input_data = {"img": im} if is_map else im
+                # non lazy
+                non_lazy_result = input_data
+                for _func in _funcs:
+                    if isinstance(_func, mt.Randomizable):
+                        _func.set_random_state(seed=seed)
+                    non_lazy_result = _func(non_lazy_result)
+                expected = non_lazy_result["img"] if is_map else non_lazy_result
+
+                # lazy
+                pending_result = input_data
+                for _func in _funcs:
+                    _func.lazy = True
+                    if isinstance(_func, mt.Randomizable):
+                        _func.set_random_state(seed=seed)
+                    pending_result = _func(pending_result)
+                pending_result = pending_result["img"] if is_map else pending_result
+
+                assert_allclose(pending_result.peek_pending_affine(), expected.affine, atol=1e-7)
+                assert_allclose(pending_result.peek_pending_shape(), expected.shape[1:4])
+
+                # test final result
+                init_param = funcs[-1][1]
+                call_param = {}
+                apply_param = get_apply_param(init_param, call_param)
+                result = apply_pending(pending_result, overrides=apply_param)[0]
+
+                match_ratio = np.sum(np.isclose(result.array, expected.array, atol=5e-1)) / np.prod(result.shape)
+                self.assertGreater(match_ratio, 0.5)  # at least half of the images are very close
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_testtimeaugmentation.py

@@ -0,0 +1,186 @@
+# 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 functools import partial
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from monai.data import CacheDataset, DataLoader, create_test_image_2d
+from monai.data.test_time_augmentation import TestTimeAugmentation
+from monai.data.utils import pad_list_data_collate
+from monai.losses import DiceLoss
+from monai.networks.nets import UNet
+from monai.transforms import (
+    Activations,
+    AsDiscrete,
+    Compose,
+    CropForegroundd,
+    DivisiblePadd,
+    EnsureChannelFirstd,
+    RandAffined,
+    RandScaleIntensityd,
+)
+from monai.transforms.croppad.dictionary import SpatialPadd
+from monai.transforms.spatial.dictionary import RandFlipd
+from monai.utils import optional_import, set_determinism
+from monai.utils.enums import PostFix
+from tests.test_utils import TEST_NDARRAYS
+
+if TYPE_CHECKING:
+    import tqdm
+
+    has_tqdm = True
+    has_nib = True
+else:
+    tqdm, has_tqdm = optional_import("tqdm")
+    _, has_nib = optional_import("nibabel")
+
+trange = partial(tqdm.trange, desc="training") if has_tqdm else range
+
+
+class TestTestTimeAugmentation(unittest.TestCase):
+    @staticmethod
+    def get_data(num_examples, input_size, data_type=np.asarray, include_label=True):
+        custom_create_test_image_2d = partial(
+            create_test_image_2d, *input_size, rad_max=7, num_seg_classes=1, num_objs=1
+        )
+        data = []
+        for i in range(num_examples):
+            im, label = custom_create_test_image_2d()
+            d = {"image": data_type(im[:, i:])}
+            if include_label:
+                d["label"] = data_type(label[:, i:])
+                d[PostFix.meta("label")] = {"affine": np.eye(4)}
+            data.append(d)
+        return data[0] if num_examples == 1 else data
+
+    def setUp(self) -> None:
+        set_determinism(seed=0)
+
+    def tearDown(self) -> None:
+        set_determinism(None)
+
+    def test_test_time_augmentation(self):
+        input_size = (20, 40)  # test different input data shape to pad list collate
+        keys = ["image", "label"]
+        num_training_ims = 10
+
+        train_data = self.get_data(num_training_ims, input_size)
+        test_data = self.get_data(1, input_size)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        transforms = Compose(
+            [
+                EnsureChannelFirstd(keys, channel_dim="no_channel"),
+                RandAffined(
+                    keys,
+                    prob=1.0,
+                    spatial_size=(30, 30),
+                    rotate_range=(np.pi / 3, np.pi / 3),
+                    translate_range=(3, 3),
+                    scale_range=((0.8, 1), (0.8, 1)),
+                    padding_mode="zeros",
+                    mode=("bilinear", "nearest"),
+                ),
+                CropForegroundd(keys, source_key="image"),
+                DivisiblePadd(keys, 4),
+            ]
+        )
+
+        train_ds = CacheDataset(train_data, transforms)
+        # output might be different size, so pad so that they match
+        train_loader = DataLoader(train_ds, batch_size=2, collate_fn=pad_list_data_collate)
+
+        model = UNet(2, 1, 1, channels=(6, 6), strides=(2, 2)).to(device)
+        loss_function = DiceLoss(sigmoid=True)
+        optimizer = torch.optim.Adam(model.parameters(), 1e-3)
+
+        num_epochs = 10
+        for _ in trange(num_epochs):
+            epoch_loss = 0
+
+            for batch_data in train_loader:
+                inputs, labels = (batch_data["image"].to(device), batch_data["label"].to(device))
+                optimizer.zero_grad()
+                outputs = model(inputs)
+                loss = loss_function(outputs, labels)
+                loss.backward()
+                optimizer.step()
+                epoch_loss += loss.item()
+
+            epoch_loss /= len(train_loader)
+
+        post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+
+        tt_aug = TestTimeAugmentation(
+            transform=transforms,
+            batch_size=5,
+            num_workers=0,
+            inferrer_fn=model,
+            device=device,
+            to_tensor=True,
+            output_device="cpu",
+            post_func=post_trans,
+        )
+        mode, mean, std, vvc = tt_aug(test_data)
+        self.assertEqual(mode.shape, (1,) + input_size)
+        self.assertEqual(mean.shape, (1,) + input_size)
+        self.assertTrue(all(np.unique(mode) == (0, 1)))
+        self.assertGreaterEqual(mean.min(), 0.0)
+        self.assertLessEqual(mean.max(), 1.0)
+        self.assertEqual(std.shape, (1,) + input_size)
+        self.assertIsInstance(vvc, float)
+
+    def test_warn_non_random(self):
+        transforms = Compose([EnsureChannelFirstd("im", channel_dim="no_channel"), SpatialPadd("im", 1)])
+        with self.assertWarns(UserWarning):
+            TestTimeAugmentation(transforms, None, None, None)
+
+    def test_warn_random_but_has_no_invertible(self):
+        transforms = Compose(
+            [
+                EnsureChannelFirstd("image", channel_dim="no_channel"),
+                RandFlipd("image", prob=1.0),
+                RandScaleIntensityd("image", 0.1, prob=1.0),
+            ]
+        )
+        with self.assertWarns(UserWarning):
+            tta = TestTimeAugmentation(transforms, 5, 0, orig_key="image")
+            tta(self.get_data(1, (20, 20), data_type=np.float32))
+
+    def test_warn_random_but_all_not_invertible(self):
+        """test with no invertible stack"""
+        transforms = Compose(
+            [EnsureChannelFirstd("image", channel_dim="no_channel"), RandScaleIntensityd("image", 0.1, prob=1.0)]
+        )
+        with self.assertWarns(UserWarning):
+            tta = TestTimeAugmentation(transforms, 1, 0, orig_key="image")
+            tta(self.get_data(1, (20, 20), data_type=np.float32))
+
+    def test_single_transform(self):
+        for p in TEST_NDARRAYS:
+            transforms = RandFlipd(["image", "label"], prob=1.0)
+            tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x)
+            tta(self.get_data(1, (20, 20), data_type=p))
+
+    def test_image_no_label(self):
+        transforms = RandFlipd(["image"], prob=1.0)
+        tta = TestTimeAugmentation(transforms, batch_size=5, num_workers=0, inferrer_fn=lambda x: x, orig_key="image")
+        tta(self.get_data(1, (20, 20), include_label=False))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_vis_gradbased.py

@@ -0,0 +1,69 @@
+# 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.nets import DenseNet, DenseNet121, SEResNet50
+from monai.visualize import GuidedBackpropGrad, GuidedBackpropSmoothGrad, SmoothGrad, VanillaGrad
+
+
+class DenseNetAdjoint(DenseNet121):
+
+    def __call__(self, x, adjoint_info):
+        if adjoint_info != 42:
+            raise ValueError
+        return super().__call__(x)
+
+
+DENSENET2D = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3)
+DENSENET3D = DenseNet(spatial_dims=3, in_channels=1, out_channels=3, init_features=2, growth_rate=2, block_config=(6,))
+SENET2D = SEResNet50(spatial_dims=2, in_channels=3, num_classes=4)
+SENET3D = SEResNet50(spatial_dims=3, in_channels=3, num_classes=4)
+DENSENET2DADJOINT = DenseNetAdjoint(spatial_dims=2, in_channels=1, out_channels=3)
+
+TESTS = []
+for type in (VanillaGrad, SmoothGrad, GuidedBackpropGrad, GuidedBackpropSmoothGrad):
+    # 2D densenet
+    TESTS.append([type, DENSENET2D, (1, 1, 48, 64)])
+    # 3D densenet
+    TESTS.append([type, DENSENET3D, (1, 1, 6, 6, 6)])
+    # 2D senet
+    TESTS.append([type, SENET2D, (1, 3, 64, 64)])
+    # 3D senet
+    TESTS.append([type, SENET3D, (1, 3, 8, 8, 48)])
+    # 2D densenet - adjoint
+    TESTS.append([type, DENSENET2DADJOINT, (1, 1, 48, 64)])
+
+
+class TestGradientClassActivationMap(unittest.TestCase):
+
+    @parameterized.expand(TESTS)
+    def test_shape(self, vis_type, model, shape):
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+        # optionally test for adjoint info
+        kwargs = {"adjoint_info": 42} if isinstance(model, DenseNetAdjoint) else {}
+
+        model.to(device)
+        model.eval()
+        vis = vis_type(model)
+        x = torch.rand(shape, device=device)
+        result = vis(x, **kwargs)
+        self.assertTupleEqual(result.shape, x.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()