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

Code before:
@skip_if_quick
class TestFLEXIBLEUNET(unittest.TestCase):

    @parameterized.expand(CASES_2D + CASES_3D + CASES_VARIATIONS)
    def test_shape(self, input_param, input_shape, expected_shape):
        device = "cuda" if torch.cuda.is_available() else "cpu"

        with skip_if_downloading_fails():
            net = FlexibleUNet(**input_param).to(device)

        # run inference with random tensor
        with eval_mode(net):
            result = net(torch.randn(input_shape).to(device))

        # check output shape
        self.assertEqual(result.shape, expected_shape)

    @parameterized.expand(CASES_PRETRAIN)
    def test_pretrain(self, flexunet_input_param, feature_extractor_class, feature_extractor_input_param, weight_list):
        device = "cuda" if torch.cuda.is_available() else "cpu"

        with skip_if_downloading_fails():
            net = FlexibleUNet(**flexunet_input_param).to(device)

        with skip_if_downloading_fails():
            eff_net = feature_extractor_class(**feature_extractor_input_param).to(device)

        for weight_name in weight_list:
            if weight_name in net.encoder.state_dict() and weight_name in eff_net.state_dict():
                net_weight = net.encoder.state_dict()[weight_name]
                download_weight = eff_net.state_dict()[weight_name]
                weight_diff = torch.abs(net_weight - download_weight)
                diff_sum = torch.sum(weight_diff)
                # check if a weight in weight_list equals to the downloaded weight.
                self.assertLess(abs(diff_sum.item() - 0), 1e-8)

    @parameterized.expand(CASE_ERRORS)
    def test_error_raise(self, input_param):
        with self.assertRaises((ValueError, NotImplementedError)):
            FlexibleUNet(**input_param)


class TestFlexUNetEncoderRegister(unittest.TestCase):

    @parameterized.expand(CASE_REGISTER_ENCODER)
    def test_regist(self, encoder):
        tmp_backbone = FlexUNetEncoderRegister()
        tmp_backbone.register_class(encoder)
        for backbone in tmp_backbone.register_dict:
            backbone_type = tmp_backbone.register_dict[backbone]["type"]
            feature_number = backbone_type.num_outputs()

Review comment (performance):
Could you please rename the variable `eff_net` to a more general one? It means  efficient net here. And since you are adding more backbones, it should be a more general one like `feature_extractor_net`.

Thanks,
Bin

Code after:
@skip_if_quick
class TestFLEXIBLEUNET(unittest.TestCase):

    @parameterized.expand(CASES_2D + CASES_3D + CASES_VARIATIONS)
    def test_shape(self, input_param, input_shape, expected_shape):
        device = "cuda" if torch.cuda.is_available() else "cpu"

        with skip_if_downloading_fails():
            net = FlexibleUNet(**input_param).to(device)

        # run inference with random tensor
        with eval_mode(net):
            result = net(torch.randn(input_shape).to(device))

        # check output shape
        self.assertEqual(result.shape, expected_shape)

    @parameterized.expand(CASES_PRETRAIN)
    def test_pretrain(self, flexunet_input_param, feature_extractor_class, feature_extractor_input_param, weight_list):
        device = "cuda" if torch.cuda.is_available() else "cpu"

        with skip_if_downloading_fails():
            net = FlexibleUNet(**flexunet_input_param).to(device)

        with skip_if_downloading_fails():
            feature_extractor_net = feature_extractor_class(**feature_extractor_input_param).to(device)

        for weight_name in weight_list:
            if weight_name in net.encoder.state_dict() and weight_name in feature_extractor_net.state_dict():
                net_weight = net.encoder.state_dict()[weight_name]
                download_weight = feature_extractor_net.state_dict()[weight_name]
                weight_diff = torch.abs(net_weight - download_weight)
                diff_sum = torch.sum(weight_diff)
                # check if a weight in weight_list equals to the downloaded weight.
                self.assertLess(abs(diff_sum.item() - 0), 1e-8)

    @parameterized.expand(CASE_ERRORS)
    def test_error_raise(self, input_param):
        with self.assertRaises((ValueError, NotImplementedError)):
            FlexibleUNet(**input_param)


class TestFlexUNetEncoderRegister(unittest.TestCase):

    @parameterized.expand(CASE_REGISTER_ENCODER)
    def test_regist(self, encoder):
        tmp_backbone = FlexUNetEncoderRegister()
        tmp_backbone.register_class(encoder)
        for backbone in tmp_backbone.register_dict:
            backbone_type = tmp_backbone.register_dict[backbone]["type"]
            feature_number = backbone_type.num_outputs()