def test_predict_non_ensemble(batch_size: int, empty_labels: bool) -> None: config = ConstantScalarConfig(1.) model = create_lightning_model(config, set_optimizer_and_scheduler=False) assert isinstance(model, ScalarLightning) pipeline = ScalarInferencePipeline(model, config, 0) actual_labels = torch.zeros( (batch_size, 1)) * np.nan if empty_labels else torch.zeros( (batch_size, 1)) data = { "metadata": [GeneralSampleMetadata(id='2')] * batch_size, "label": actual_labels, "images": torch.zeros( ((batch_size, 1) + config.expected_image_size_zyx)), "numerical_non_image_features": torch.tensor([]), "categorical_non_image_features": torch.tensor([]), "segmentations": torch.tensor([]) } results = pipeline.predict(data) ids, labels, predicted = results.subject_ids, results.labels, results.posteriors assert ids == ['2'] * batch_size assert torch.allclose(labels, actual_labels, equal_nan=True) # The model always returns 1, so predicted should be sigmoid(1) assert torch.allclose(predicted, torch.full((batch_size, 1), 0.731058578))
def test_predict_ensemble(batch_size: int) -> None: config_returns_0 = ConstantScalarConfig(0.) model_returns_0 = create_lightning_model(config_returns_0, set_optimizer_and_scheduler=False) assert isinstance(model_returns_0, ScalarLightning) config_returns_1 = ConstantScalarConfig(1.) model_returns_1 = create_lightning_model(config_returns_1, set_optimizer_and_scheduler=False) assert isinstance(model_returns_1, ScalarLightning) pipeline_0 = ScalarInferencePipeline(model_returns_0, config_returns_0, 0) pipeline_1 = ScalarInferencePipeline(model_returns_0, config_returns_0, 1) pipeline_2 = ScalarInferencePipeline(model_returns_0, config_returns_0, 2) pipeline_3 = ScalarInferencePipeline(model_returns_1, config_returns_1, 3) pipeline_4 = ScalarInferencePipeline(model_returns_1, config_returns_1, 4) ensemble_pipeline = ScalarEnsemblePipeline( [pipeline_0, pipeline_1, pipeline_2, pipeline_3, pipeline_4], config_returns_0, EnsembleAggregationType.Average) data = { "metadata": [GeneralSampleMetadata(id='2')] * batch_size, "label": torch.zeros((batch_size, 1)), "images": torch.zeros( ((batch_size, 1) + config_returns_0.expected_image_size_zyx)), "numerical_non_image_features": torch.tensor([]), "categorical_non_image_features": torch.tensor([]), "segmentations": torch.tensor([]) } results = ensemble_pipeline.predict(data) ids, labels, predicted = results.subject_ids, results.labels, results.posteriors assert ids == ['2'] * batch_size assert torch.equal(labels, torch.zeros((batch_size, 1))) # 3 models return 0, 2 return 1, so predicted should be ((sigmoid(0)*3)+(sigmoid(1)*2))/5 assert torch.allclose(predicted, torch.full((batch_size, 1), 0.592423431))
def load_from_lightning_checkpoint(config: ModelConfigBase, checkpoint_path: Path) -> InnerEyeLightning: """ Reads a PyTorch model from a checkpoint. First, a PyTorch Lightning model is created matching the InnerEye model configuration, its parameter tensors are then populated from the given checkpoint. :param config: An InnerEye model configuration object :param checkpoint_path: The location of the checkpoint file. :return: A PyTorch Lightning model object. """ # Create a Lighting model that matches the configuration, but keep only the type of it lightning_model_type = type(create_lightning_model(config)) # For model debugging, allow loading a GPU trained model onto the CPU. This will clearly only work # if the model is small. map_location = None if config.use_gpu else 'cpu' lightning_model = lightning_model_type.load_from_checkpoint(checkpoint_path=str(checkpoint_path), map_location=map_location, config=config) return lightning_model
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) lightning_model = create_lightning_model(config) assert isinstance(lightning_model, SegmentationLightning) pipeline = InferencePipeline(model=lightning_model, model_config=config) 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 create_model_and_store_checkpoint(config: ModelConfigBase, checkpoint_path: Path, weights_only: bool = True) -> None: """ Creates a Lightning model for the given model configuration, and stores it as a checkpoint file. If a GPU is available, the model is moved to the GPU before storing. The trainer properties `current_epoch` and `global_step` are set to fixed non-default values. :param config: The model configuration. :param checkpoint_path: The path and filename of the checkpoint file. """ container = InnerEyeContainer(config) trainer, _ = create_lightning_trainer(container) model = create_lightning_model(config) if machine_has_gpu: model = model.cuda() # type: ignore trainer.model = model # Before saving, the values for epoch and step are incremented. Save them here in such a way that we can assert # easily later. We can't mock that because otherwise the mock object would be written to disk (that fails) trainer.fit_loop.current_epoch = FIXED_EPOCH - 1 # type: ignore trainer.fit_loop.global_step = FIXED_GLOBAL_STEP - 1 # type: ignore # In PL, it is the Trainer's responsibility to save the model. Checkpoint handling refers back to the trainer # to get a save_func. Mimicking that here. trainer.save_checkpoint(checkpoint_path, weights_only=weights_only)
def create_model(self) -> LightningModule: # type: ignore from InnerEye.ML.lightning_models import create_lightning_model return create_lightning_model(self.config)