Repository: Project-MONAI/MONAI
File: tests/test_attentionunet.py (Python)

Code before:

    @skip_if_quick
    def test_attentionunet(self):
        for dims in [2, 3]:
            shape = (3, 1) + (92,) * dims
            input = torch.rand(*shape)
            model = att.AttentionUnet(
                spatial_dims=dims, in_channels=1, out_channels=2, channels=(3, 4, 5), up_kernel_size=5, strides=(1, 2)
            )
            output = model(input)
            self.assertEqual(output.shape[2:], input.shape[2:])
            self.assertEqual(output.shape[0], input.shape[0])
            self.assertEqual(output.shape[1], 2)

    def test_attentionunet_kernel_size(self):
        args_dict = {
            "spatial_dims": 2,
            "in_channels": 1,
            "out_channels": 2,
            "channels": (3, 4, 5),
            "up_kernel_size": 5,
            "strides": (1, 2),
        }
        model_a = att.AttentionUnet(**args_dict, kernel_size=5)
        model_b = att.AttentionUnet(**args_dict, kernel_size=7)
        self.assertNotEqual(get_net_parameters(model_a), get_net_parameters(model_b))

    @skip_if_no_cuda
    def test_attentionunet_gpu(self):
        for dims in [2, 3]:
            shape = (3, 1) + (92,) * dims
            input = torch.rand(*shape).to("cuda:0")
            model = att.AttentionUnet(
                spatial_dims=dims, in_channels=1, out_channels=2, channels=(3, 4, 5), strides=(2, 2)
            ).to("cuda:0")
            with torch.no_grad():
                output = model(input)
                self.assertEqual(output.shape[2:], input.shape[2:])
                self.assertEqual(output.shape[0], input.shape[0])
                self.assertEqual(output.shape[1], 2)


if __name__ == "__main__":
    unittest.main()


Review comment (suggestion):
I would suggest as a more useful test to store the number of expected parameters for model_a and model_b and compare against those values. With that minor change we'd be ready to go with this PR I think.

Code after:

    @skip_if_quick
    def test_attentionunet(self):
        for dims in [2, 3]:
            shape = (3, 1) + (92,) * dims
            input = torch.rand(*shape)
            model = att.AttentionUnet(
                spatial_dims=dims, in_channels=1, out_channels=2, channels=(3, 4, 5), up_kernel_size=5, strides=(1, 2)
            )
            output = model(input)
            self.assertEqual(output.shape[2:], input.shape[2:])
            self.assertEqual(output.shape[0], input.shape[0])
            self.assertEqual(output.shape[1], 2)

    def test_attentionunet_kernel_size(self):
        args_dict = {
            "spatial_dims": 2,
            "in_channels": 1,
            "out_channels": 2,
            "channels": (3, 4, 5),
            "up_kernel_size": 5,
            "strides": (1, 2),
        }
        model_a = att.AttentionUnet(**args_dict, kernel_size=5)
        model_b = att.AttentionUnet(**args_dict, kernel_size=7)
        self.assertEqual(get_net_parameters(model_a), 3534)
        self.assertEqual(get_net_parameters(model_b), 5574)

    @skip_if_no_cuda
    def test_attentionunet_gpu(self):
        for dims in [2, 3]:
            shape = (3, 1) + (92,) * dims
            input = torch.rand(*shape).to("cuda:0")
            model = att.AttentionUnet(
                spatial_dims=dims, in_channels=1, out_channels=2, channels=(3, 4, 5), strides=(2, 2)
            ).to("cuda:0")
            with torch.no_grad():
                output = model(input)
                self.assertEqual(output.shape[2:], input.shape[2:])
                self.assertEqual(output.shape[0], input.shape[0])
                self.assertEqual(output.shape[1], 2)


if __name__ == "__main__":
    unittest.main()