(self, input_data, expected_shape)
| 70 | |
| 71 | @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3]) |
| 72 | def test_shape(self, input_data, expected_shape): |
| 73 | model = None |
| 74 | |
| 75 | if input_data["model"] == "densenet2d": |
| 76 | model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=3) |
| 77 | if input_data["model"] == "densenet3d": |
| 78 | model = DenseNet( |
| 79 | spatial_dims=3, in_channels=1, out_channels=3, init_features=2, growth_rate=2, block_config=(6,) |
| 80 | ) |
| 81 | if input_data["model"] == "senet2d": |
| 82 | model = SEResNet50(spatial_dims=2, in_channels=3, num_classes=4) |
| 83 | if input_data["model"] == "senet3d": |
| 84 | model = SEResNet50(spatial_dims=3, in_channels=3, num_classes=4) |
| 85 | |
| 86 | device = "cuda:0" if torch.cuda.is_available() else "cpu" |
| 87 | model.to(device) |
| 88 | model.eval() |
| 89 | cam = CAM(nn_module=model, target_layers=input_data["target_layers"], fc_layers=input_data["fc_layers"]) |
| 90 | image = torch.rand(input_data["shape"], device=device) |
| 91 | result = cam(x=image, layer_idx=-1) |
| 92 | fea_shape = cam.feature_map_size(input_data["shape"], device=device) |
| 93 | self.assertTupleEqual(fea_shape, input_data["feature_shape"]) |
| 94 | self.assertTupleEqual(result.shape, expected_shape) |
| 95 | |
| 96 | |
| 97 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected