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)