def test_mean_teacher_model() -> None: """ Test training and weight updates of the mean teacher model computation. """ def _get_parameters_of_model( model: Union[torch.nn.Module, DataParallelModel]) -> Any: """ Returns the iterator of model parameters """ if isinstance(model, DataParallelModel): return model.module.parameters() else: return model.parameters() config = DummyClassification() config.num_epochs = 1 # Set train batch size to be arbitrary big to ensure we have only one training step # i.e. one mean teacher update. config.train_batch_size = 100 # Train without mean teacher model_train(config) # Retrieve the weight after one epoch model = create_model_with_temperature_scaling(config) print(config.get_path_to_checkpoint(1)) _ = model_util.load_checkpoint(model, config.get_path_to_checkpoint(1)) model_weight = next(_get_parameters_of_model(model)) # Get the starting weight of the mean teacher model ml_util.set_random_seed(config.get_effective_random_seed()) _ = create_model_with_temperature_scaling(config) mean_teach_model = create_model_with_temperature_scaling(config) initial_weight_mean_teacher_model = next( _get_parameters_of_model(mean_teach_model)) # Now train with mean teacher and check the update of the weight alpha = 0.999 config.mean_teacher_alpha = alpha model_train(config) # Retrieve weight of mean teacher model saved in the checkpoint mean_teacher_model = create_model_with_temperature_scaling(config) _ = model_util.load_checkpoint( mean_teacher_model, config.get_path_to_checkpoint(1, for_mean_teacher_model=True)) result_weight = next(_get_parameters_of_model(mean_teacher_model)) # Retrieve the associated student weight _ = model_util.load_checkpoint(model, config.get_path_to_checkpoint(1)) student_model_weight = next(_get_parameters_of_model(model)) # Assert that the student weight corresponds to the weight of a simple training without mean teacher # computation assert student_model_weight.allclose(model_weight) # Check the update of the parameters assert torch.all(alpha * initial_weight_mean_teacher_model + (1 - alpha) * student_model_weight == result_weight)
def test_visualization_for_different_target_weeks(test_output_dirs: TestOutputDirectories) -> None: """ Tests that the visualizations are differentiated depending on the target week for which we visualize it. """ config = ToyMultiLabelSequenceModel(should_validate=False) config.set_output_to(test_output_dirs.root_dir) config.dataset_data_frame = _get_multi_label_sequence_dataframe() config.pre_process_dataset_dataframe() model = create_model_with_temperature_scaling(config) dataloader = SequenceDataset(config, data_frame=config.dataset_data_frame).as_data_loader(shuffle=False, batch_size=2) batch = next(iter(dataloader)) model_inputs_and_labels = get_scalar_model_inputs_and_labels(config, model, batch) # type: ignore visualizer = VisualizationMaps(model, config) # Pseudo-grad cam explaining the prediction at target sequence 2 _, _, pseudo_cam_non_img_3, probas_3 = visualizer.generate(model_inputs_and_labels.model_inputs, target_position=2, target_label_index=2) # Pseudo-grad cam explaining the prediction at target sequence 0 _, _, pseudo_cam_non_img_1, probas_1 = visualizer.generate(model_inputs_and_labels.model_inputs, target_position=0, target_label_index=0) assert pseudo_cam_non_img_1.shape[1] == 1 assert pseudo_cam_non_img_3.shape[1] == 3 # Both visualizations should not be equal assert np.any(pseudo_cam_non_img_1 != pseudo_cam_non_img_3) assert np.any(probas_3 != probas_1)
def test_visualization_with_sequence_model(use_combined_model: bool, imaging_feature_type: ImagingFeatureType, test_output_dirs: TestOutputDirectories) -> None: config = ToySequenceModel(use_combined_model, imaging_feature_type, should_validate=False) config.set_output_to(test_output_dirs.root_dir) config.dataset_data_frame = _get_mock_sequence_dataset() config.num_epochs = 1 model = create_model_with_temperature_scaling(config) update_model_for_multiple_gpus(ModelAndInfo(model), config) dataloader = SequenceDataset(config, data_frame=config.dataset_data_frame).as_data_loader(shuffle=False, batch_size=2) # Patch the load_images function that will be called once we access a dataset item image_and_seg = ImageAndSegmentations[np.ndarray](images=np.random.uniform(0, 1, SCAN_SIZE), segmentations=np.random.randint(0, 2, SCAN_SIZE)) with mock.patch('InnerEye.ML.utils.io_util.load_image_in_known_formats', return_value=image_and_seg): batch = next(iter(dataloader)) model_inputs_and_labels = get_scalar_model_inputs_and_labels(config, model, batch) # type: ignore number_sequences = model_inputs_and_labels.model_inputs[0].shape[1] number_subjects = len(model_inputs_and_labels.subject_ids) visualizer = VisualizationMaps(model, config) guided_grad_cams, grad_cams, pseudo_cam_non_img, probas = visualizer.generate( model_inputs_and_labels.model_inputs) if use_combined_model: if imaging_feature_type == ImagingFeatureType.ImageAndSegmentation: assert guided_grad_cams.shape[:2] == (number_subjects, number_sequences * 2) assert grad_cams.shape[:2] == (number_subjects, number_sequences * 2) else: assert guided_grad_cams.shape[:2] == (number_subjects, number_sequences) assert grad_cams.shape[:2] == (number_subjects, number_sequences) else: assert guided_grad_cams is None assert grad_cams is None assert pseudo_cam_non_img.shape[:2] == (number_subjects, number_sequences) assert probas.shape[0] == number_subjects non_image_features = config.numerical_columns + config.categorical_columns non_imaging_plot_labels = visualizer._get_non_imaging_plot_labels(model_inputs_and_labels.data_item, non_image_features, index=0, target_position=3) assert non_imaging_plot_labels == ['numerical1_0', 'numerical2_0', 'cat1_0', 'numerical1_1', 'numerical2_1', 'cat1_1', 'numerical1_2', 'numerical2_2', 'cat1_2', 'numerical1_3', 'numerical2_3', 'cat1_3']
def test_load_all_configs(model_name: str) -> None: """ Loads all model configurations that are present in the ML/src/configs folder, and carries out basic validations of the configuration. """ logger = logging.getLogger() logger.setLevel(logging.INFO) config = ModelConfigLoader().create_model_config_from_name(model_name) assert config.model_name == model_name, "Mismatch between definition .py file and model name" if config.is_segmentation_model: # Reduce the feature channels to a minimum, to make tests run fast on CPU. minimal_feature_channels = 1 config.feature_channels = [minimal_feature_channels] * len(config.feature_channels) print("Model architecture after restricting to 2 feature channels only:") model = create_model_with_temperature_scaling(config) generate_and_print_model_summary(config, model) # type: ignore else: # For classification models, we can't always print a model summary: The model could require arbitrary # numbers of input tensors, and we'd only know once we load the training data. # Hence, only try to create the model, but don't attempt to print the summary. create_model_with_temperature_scaling(config)
def extract_activation_maps(args: ModelConfigBase) -> None: """ Extracts and saves activation maps of a specific layer of a trained network :param args: :return: """ model = create_model_with_temperature_scaling(args) if args.use_gpu: model = torch.nn.DataParallel(model, device_ids=list( range(torch.cuda.device_count()))) model = model.cuda() checkpoint_path = args.get_path_to_checkpoint() if checkpoint_path.is_file(): checkpoint = torch.load(checkpoint_path) # type: ignore model.load_state_dict(checkpoint['state_dict']) else: raise FileNotFoundError("Could not find checkpoint") model.eval() val_loader = args.create_data_loaders()[ModelExecutionMode.VAL] feature_extractor = model_hooks.HookBasedFeatureExtractor( model, layer_name=args.activation_map_layers) for batch, sample in enumerate(val_loader): sample = CroppedSample.from_dict(sample=sample) input_image = sample.image.cuda().float() feature_extractor(input_image) # access first image of batch of feature maps activation_map = feature_extractor.outputs[0][0].cpu().numpy() if len(activation_map.shape) == 4: visualize_3d_activation_map(activation_map, args) elif len(activation_map.shape) == 3: visualize_2d_activation_map(activation_map, args) else: raise NotImplementedError( 'cannot visualize activation map of shape', activation_map.shape) # Only visualize the first validation example break
def run_inference_on_unet(size: TupleInt3) -> None: """ Runs a model forward pass on a freshly created model, with an input image of the given size. Asserts that the model prediction has the same size as the input image. """ fg_classes = ["tumour_mass", "subtract"] number_of_classes = len(fg_classes) + 1 config = SegmentationModelBase( architecture="UNet3D", local_dataset=Path("dummy"), feature_channels=[1], kernel_size=3, largest_connected_component_foreground_classes=fg_classes, posterior_smoothing_mm=(2, 2, 2), crop_size=(64, 64, 64), # test_crop_size must be larger than 'size for the bug to trigger test_crop_size=(80, 80, 80), image_channels=["mr"], ground_truth_ids=fg_classes, ground_truth_ids_display_names=fg_classes, colours=[(255, 0, 0)] * len(fg_classes), fill_holes=[False] * len(fg_classes), mask_id=None, class_weights=[1.0 / number_of_classes] * number_of_classes, train_batch_size=8, inference_batch_size=1, inference_stride_size=(40, 40, 40), use_mixed_precision=True) model = create_model_with_temperature_scaling(config) pipeline = InferencePipeline(model=model, model_config=config, epoch=1) image = np.random.uniform(-1, 1, (1, ) + size) result = pipeline.predict_and_post_process_whole_image( image, mask=np.ones(size), voxel_spacing_mm=(1, 1, 1)) # All posteriors and segmentations must have the size of the input image for p in [*result.posteriors, result.segmentation]: assert p.shape == size # Check that all results are not NaN. In particular, if stride size is not adjusted # correctly, the results would be partially NaN. image_util.check_array_range(p)
def test_model_summary_on_classification1() -> None: model = create_model_with_temperature_scaling(GlaucomaPublic()) ModelSummary(model).generate_summary(input_sizes=[(1, 6, 64, 60)])
def test_visualization_with_scalar_model(use_non_imaging_features: bool, imaging_feature_type: ImagingFeatureType, encode_channels_jointly: bool, test_output_dirs: OutputFolderForTests) -> None: dataset_contents = """subject,channel,path,label,numerical1,numerical2,categorical1,categorical2 S1,week0,scan1.npy,,1,10,Male,Val1 S1,week1,scan2.npy,True,2,20,Female,Val2 S2,week0,scan3.npy,,3,30,Female,Val3 S2,week1,scan4.npy,False,4,40,Female,Val1 S3,week0,scan1.npy,,5,50,Male,Val2 S3,week1,scan3.npy,True,6,60,Male,Val2 """ dataset_dataframe = pd.read_csv(StringIO(dataset_contents), dtype=str) numerical_columns = ["numerical1", "numerical2"] if use_non_imaging_features else [] categorical_columns = ["categorical1", "categorical2"] if use_non_imaging_features else [] non_image_feature_channels = get_non_image_features_dict(default_channels=["week1", "week0"], specific_channels={"categorical2": ["week1"]}) \ if use_non_imaging_features else {} config = ImageEncoder( local_dataset=Path(), encode_channels_jointly=encode_channels_jointly, should_validate=False, numerical_columns=numerical_columns, categorical_columns=categorical_columns, imaging_feature_type=imaging_feature_type, non_image_feature_channels=non_image_feature_channels, categorical_feature_encoder=CategoricalToOneHotEncoder.create_from_dataframe( dataframe=dataset_dataframe, columns=categorical_columns) ) dataloader = ScalarDataset(config, data_frame=dataset_dataframe) \ .as_data_loader(shuffle=False, batch_size=2) config.set_output_to(test_output_dirs.root_dir) config.num_epochs = 1 model = create_model_with_temperature_scaling(config) visualizer = VisualizationMaps(model, config) # Patch the load_images function that will be called once we access a dataset item image_and_seg = ImageAndSegmentations[np.ndarray](images=np.random.uniform(0, 1, (6, 64, 60)), segmentations=np.random.randint(0, 2, (6, 64, 60))) with mock.patch('InnerEye.ML.utils.io_util.load_image_in_known_formats', return_value=image_and_seg): batch = next(iter(dataloader)) if config.use_gpu: device = visualizer.grad_cam.device batch = transfer_batch_to_device(batch, device) visualizer.grad_cam.model = visualizer.grad_cam.model.to(device) model_inputs_and_labels = get_scalar_model_inputs_and_labels(model, target_indices=[], sample=batch) number_channels = len(config.image_channels) number_subjects = len(model_inputs_and_labels.subject_ids) guided_grad_cams, grad_cams, pseudo_cam_non_img, probas = visualizer.generate( model_inputs_and_labels.model_inputs) if imaging_feature_type == ImagingFeatureType.ImageAndSegmentation: assert guided_grad_cams.shape[:2] == (number_subjects, number_channels * 2) else: assert guided_grad_cams.shape[:2] == (number_subjects, number_channels) assert grad_cams.shape[:2] == (number_subjects, 1) if encode_channels_jointly \ else (number_subjects, number_channels) if use_non_imaging_features: non_image_features = config.numerical_columns + config.categorical_columns non_imaging_plot_labels = visualizer._get_non_imaging_plot_labels(model_inputs_and_labels.data_item, non_image_features, index=0) assert non_imaging_plot_labels == ['numerical1_week1', 'numerical1_week0', 'numerical2_week1', 'numerical2_week0', 'categorical1_week1', 'categorical1_week0', 'categorical2_week1'] assert pseudo_cam_non_img.shape == (number_subjects, 1, len(non_imaging_plot_labels))
def test_image_encoder(test_output_dirs: OutputFolderForTests, encode_channels_jointly: bool, use_non_imaging_features: bool, kernel_size_per_encoding_block: Optional[Union[TupleInt3, List[TupleInt3]]], stride_size_per_encoding_block: Optional[Union[TupleInt3, List[TupleInt3]]], reduction_factor: float, expected_num_reduced_features: int, aggregation_type: AggregationType) -> None: """ Test if the image encoder networks can be trained without errors (including GradCam computation and data augmentation). """ logging_to_stdout() set_random_seed(0) dataset_folder = Path(test_output_dirs.make_sub_dir("dataset")) scan_size = (6, 64, 60) scan_files: List[str] = [] for s in range(4): random_scan = np.random.uniform(0, 1, scan_size) scan_file_name = f"scan{s + 1}{NumpyFile.NUMPY.value}" np.save(str(dataset_folder / scan_file_name), random_scan) scan_files.append(scan_file_name) dataset_contents = """subject,channel,path,label,numerical1,numerical2,categorical1,categorical2 S1,week0,scan1.npy,,1,10,Male,Val1 S1,week1,scan2.npy,True,2,20,Female,Val2 S2,week0,scan3.npy,,3,30,Female,Val3 S2,week1,scan4.npy,False,4,40,Female,Val1 S3,week0,scan1.npy,,5,50,Male,Val2 S3,week1,scan3.npy,True,6,60,Male,Val2 """ (dataset_folder / "dataset.csv").write_text(dataset_contents) numerical_columns = ["numerical1", "numerical2"] if use_non_imaging_features else [] categorical_columns = ["categorical1", "categorical2"] if use_non_imaging_features else [] non_image_feature_channels = get_non_image_features_dict(default_channels=["week1", "week0"], specific_channels={"categorical2": ["week1"]}) \ if use_non_imaging_features else {} config_for_dataset = ScalarModelBase( local_dataset=dataset_folder, image_channels=["week0", "week1"], image_file_column="path", label_channels=["week1"], label_value_column="label", non_image_feature_channels=non_image_feature_channels, numerical_columns=numerical_columns, categorical_columns=categorical_columns, should_validate=False ) config_for_dataset.read_dataset_into_dataframe_and_pre_process() dataset = ScalarDataset(config_for_dataset, sample_transforms=ScalarItemAugmentation( RandAugmentSlice(is_transformation_for_segmentation_maps=False))) assert len(dataset) == 3 config = ImageEncoder( encode_channels_jointly=encode_channels_jointly, should_validate=False, numerical_columns=numerical_columns, categorical_columns=categorical_columns, non_image_feature_channels=non_image_feature_channels, categorical_feature_encoder=config_for_dataset.categorical_feature_encoder, encoder_dimensionality_reduction_factor=reduction_factor, aggregation_type=aggregation_type, scan_size=(6, 64, 60) ) if kernel_size_per_encoding_block: config.kernel_size_per_encoding_block = kernel_size_per_encoding_block if stride_size_per_encoding_block: config.stride_size_per_encoding_block = stride_size_per_encoding_block config.set_output_to(test_output_dirs.root_dir) config.max_batch_grad_cam = 1 model = create_model_with_temperature_scaling(config) input_size: List[Tuple] = [(len(config.image_channels), *scan_size)] if use_non_imaging_features: input_size.append((config.get_total_number_of_non_imaging_features(),)) # Original number output channels (unreduced) is # num initial channel * (num encoder block - 1) = 4 * (3-1) = 8 if encode_channels_jointly: # reduced_num_channels + num_non_img_features assert model.final_num_feature_channels == expected_num_reduced_features + \ config.get_total_number_of_non_imaging_features() else: # num_img_channels * reduced_num_channels + num_non_img_features assert model.final_num_feature_channels == len(config.image_channels) * expected_num_reduced_features + \ config.get_total_number_of_non_imaging_features() summarizer = ModelSummary(model) summarizer.generate_summary(input_sizes=input_size) config.local_dataset = dataset_folder config.validate() model_train(config, checkpoint_handler=get_default_checkpoint_handler(model_config=config, project_root=Path(test_output_dirs.root_dir)))