def test_metrics_dict_get_hues() -> None:
"""
Test to make sure metrics dict is configured properly with/without hues
"""
m = MetricsDict()
assert m.get_hue_names() == [MetricsDict.DEFAULT_HUE_KEY]
assert m.get_hue_names(include_default=False) == []
_hues = ["A", "B", "C"]
m = MetricsDict(hues=_hues)
assert m.get_hue_names() == _hues + [MetricsDict.DEFAULT_HUE_KEY]
assert m.get_hue_names(include_default=False) == _hues
def test_metrics_dict1() -> None:
"""
Test insertion of scalar values into a MetricsDict.
"""
m = MetricsDict()
assert m.get_hue_names() == [MetricsDict.DEFAULT_HUE_KEY]
name = "foo"
v1 = 2.7
v2 = 3.14
m.add_metric(name, v1)
m.add_metric(name, v2)
assert m.values()[name] == [v1, v2]
with pytest.raises(ValueError) as ex:
# noinspection PyTypeChecker
m.add_metric(name, [1.0]) # type: ignore
assert "Expected the metric to be a scalar" in str(ex)
assert m.skip_nan_when_averaging[name] is False
v3 = 3.0
name2 = "bar"
m.add_metric(name2, v3, skip_nan_when_averaging=True)
assert m.skip_nan_when_averaging[name2] is True
# Expected average: Metric "foo" averages over two values v1 and v2. For "bar", we only inserted one value anyhow
average = m.average()
mean_v1_v2 = mean([v1, v2])
assert average.values() == {name: [mean_v1_v2], name2: [v3]}
num_entries = m.num_entries()
assert num_entries == {name: 2, name2: 1}
def log_segmentation_epoch_metrics(self, metrics: MetricsDict,
learning_rates: List[float]) -> None:
"""
Logs segmentation metrics (e.g. loss, dice scores, learning rates) to an event file for TensorBoard
visualization and to the AzureML run context
:param learning_rates: The logged learning rates.
:param metrics: The metrics of the specified epoch, averaged along its batches.
"""
logging_fn = self.log_to_azure_and_tensorboard
logging_fn(MetricType.LOSS.value,
metrics.get_single_metric(MetricType.LOSS))
logging_fn("Dice/AverageExceptBackground",
metrics.get_single_metric(MetricType.DICE))
logging_fn(
"Voxels/ProportionForeground",
metrics.get_single_metric(MetricType.PROPORTION_FOREGROUND_VOXELS))
logging_fn("TimePerEpoch_Seconds",
metrics.get_single_metric(MetricType.SECONDS_PER_EPOCH))
if learning_rates is not None:
for i, lr in enumerate(learning_rates):
logging_fn("LearningRate/Index_{}".format(i), lr)
for class_name in metrics.get_hue_names(include_default=False):
# Tensorboard groups metrics by what is before the slash.
# With metrics Dice/Foo and Dice/Bar, it will create a section for "Dice",
# and inside of it, there are graphs for Foo and Bar
get_label = lambda x, y: "{}/{}".format(x, y)
logging_fn(
get_label("Dice", class_name),
metrics.get_single_metric(MetricType.DICE, hue=class_name))
logging_fn(
get_label("Voxels", class_name),
metrics.get_single_metric(
MetricType.PROPORTION_FOREGROUND_VOXELS, hue=class_name))
def aggregate_segmentation_metrics(metrics: MetricsDict) -> MetricsDict:
"""
Computes aggregate metrics for segmentation models, from a metrics dictionary that contains the results for
individual minibatches. Specifically, average Dice scores for only the foreground structures and proportions
of foreground voxels are computed. All metrics for the background class will be removed.
All other metrics that are already present in the input metrics will be averaged and available in the result.
Diagnostic values present in the input will be passed through unchanged.
:param metrics: A metrics dictionary that contains the per-minibatch results.
"""
class_names_with_background = metrics.get_hue_names(include_default=False)
has_background_class = class_names_with_background[0] == BACKGROUND_CLASS_NAME
foreground_classes = class_names_with_background[1:] if has_background_class else class_names_with_background
result = metrics.average(across_hues=False)
result.diagnostics = metrics.diagnostics.copy()
if has_background_class:
result.delete_hue(BACKGROUND_CLASS_NAME)
add_average_foreground_dice(result)
# Total number of voxels per class, including the background class
total_voxels = []
voxel_count = MetricType.VOXEL_COUNT.value
for g in class_names_with_background:
values = metrics.values(hue=g)
if voxel_count in values:
total_voxels.append(sum(values[voxel_count]))
if len(total_voxels) > 0:
# Proportion of voxels in foreground classes only
proportion_foreground = np.array(total_voxels[1:], dtype=float) / sum(total_voxels)
for i, foreground_class in enumerate(foreground_classes):
result.add_metric(MetricType.PROPORTION_FOREGROUND_VOXELS, proportion_foreground[i], hue=foreground_class)
result.add_metric(MetricType.PROPORTION_FOREGROUND_VOXELS, np.sum(proportion_foreground).item())
return result
def test_delete_hue() -> None:
h1 = "a"
h2 = "b"
a = MetricsDict(hues=[h1, h2])
a.add_metric("foo", 1.0, hue=h1)
a.add_metric("bar", 2.0, hue=h2)
a.delete_hue(h1)
assert a.get_hue_names(include_default=False) == [h2]
assert list(a.enumerate_single_values()) == [(h2, "bar", 2.0)]
def add_average_foreground_dice(metrics: MetricsDict) -> None:
"""
If the given metrics dictionary contains an entry for Dice score, and only one value for the Dice score per class,
then add an average Dice score for all foreground classes to the metrics dictionary (modified in place).
:param metrics: The object that holds metrics. The average Dice score will be written back into this object.
"""
all_dice = []
for structure_name in metrics.get_hue_names(include_default=False):
if structure_name != BACKGROUND_CLASS_NAME:
all_dice.append(metrics.get_single_metric(MetricType.DICE, hue=structure_name))
metrics.add_metric(MetricType.DICE, np.nanmean(all_dice).item())
def test_metrics_dict_with_default_hue() -> None:
hue_name = "foo"
metrics_dict = MetricsDict(hues=[hue_name, MetricsDict.DEFAULT_HUE_KEY])
assert metrics_dict.get_hue_names(include_default=True) == [hue_name, MetricsDict.DEFAULT_HUE_KEY]
assert metrics_dict.get_hue_names(include_default=False) == [hue_name]
class ModelTrainingStepsForSegmentation(
ModelTrainingStepsBase[SegmentationModelBase, DeviceAwareModule]):
"""
This class implements all steps necessary for training an image segmentation model during a single epoch.
"""
def __init__(self, model_config: SegmentationModelBase,
train_val_params: TrainValidateParameters[DeviceAwareModule]):
"""
Creates a new instance of the class.
:param model_config: The configuration of a segmentation model.
:param train_val_params: The parameters for training the model, including the optimizer and the data loaders.
"""
super().__init__(model_config, train_val_params)
self.example_to_save = np.random.randint(
0, len(train_val_params.data_loader))
self.pipeline = SegmentationForwardPass(
model=self.train_val_params.model,
model_config=self.model_config,
batch_size=self.model_config.train_batch_size,
optimizer=self.train_val_params.optimizer,
in_training_mode=self.train_val_params.in_training_mode,
criterion=self.compute_loss,
gradient_scaler=train_val_params.gradient_scaler)
self.metrics = MetricsDict(hues=[BACKGROUND_CLASS_NAME] +
model_config.ground_truth_ids)
def create_loss_function(self) -> torch.nn.Module:
"""
Returns a torch module that computes a loss function.
"""
return self.construct_loss_function(self.model_config)
@classmethod
def construct_loss_function(
cls, model_config: SegmentationModelBase
) -> SupervisedLearningCriterion:
"""
Returns a loss function from the model config; mixture losses are constructed as weighted combinations of
other loss functions.
"""
if model_config.loss_type == SegmentationLoss.Mixture:
components = model_config.mixture_loss_components
assert components is not None
sum_weights = sum(component.weight for component in components)
weights_and_losses = []
for component in components:
normalized_weight = component.weight / sum_weights
loss_function = cls.construct_non_mixture_loss_function(
model_config, component.loss_type,
component.class_weight_power)
weights_and_losses.append((normalized_weight, loss_function))
return MixtureLoss(weights_and_losses)
return cls.construct_non_mixture_loss_function(
model_config, model_config.loss_type,
model_config.loss_class_weight_power)
@classmethod
def construct_non_mixture_loss_function(
cls, model_config: SegmentationModelBase,
loss_type: SegmentationLoss,
power: Optional[float]) -> SupervisedLearningCriterion:
"""
:param model_config: model configuration to get some parameters from
:param loss_type: type of loss function
:param power: value for class_weight_power for the loss function
:return: instance of loss function
"""
if loss_type == SegmentationLoss.SoftDice:
return SoftDiceLoss(class_weight_power=power)
elif loss_type == SegmentationLoss.CrossEntropy:
return CrossEntropyLoss(
class_weight_power=power,
smoothing_eps=model_config.label_smoothing_eps,
focal_loss_gamma=None)
elif loss_type == SegmentationLoss.Focal:
return CrossEntropyLoss(
class_weight_power=power,
smoothing_eps=model_config.label_smoothing_eps,
focal_loss_gamma=model_config.focal_loss_gamma)
else:
raise NotImplementedError(
"Loss type {} is not implemented".format(loss_type))
def forward_and_backward_minibatch(
self, sample: Dict[str, Any], batch_index: int,
epoch: int) -> ModelForwardAndBackwardsOutputs:
"""
Runs training for a single minibatch of training data, and computes all metrics.
:param sample: The batched sample on which the model should be trained.
:param batch_index: The index of the present batch (supplied only for diagnostics).
:param epoch: The number of the present epoch.
"""
cropped_sample: CroppedSample = CroppedSample.from_dict(sample=sample)
labels = self.model_config.get_gpu_tensor_if_possible(
cropped_sample.labels_center_crop)
mask = None if self.train_val_params.in_training_mode else cropped_sample.mask_center_crop
forward_pass_result = self.pipeline.forward_pass_patches(
patches=cropped_sample.image, labels=labels, mask=mask)
# Clear the GPU cache between forward and backward passes to avoid possible out-of-memory
torch.cuda.empty_cache()
dice_for_all_classes = metrics.compute_dice_across_patches(
segmentation=torch.tensor(
forward_pass_result.segmentations).long(),
ground_truth=labels,
use_cuda=self.model_config.use_gpu,
allow_multiple_classes_for_each_pixel=True).cpu().numpy()
foreground_voxels = metrics_util.get_number_of_voxels_per_class(
cropped_sample.labels)
# loss is a scalar, also when running the forward pass over multiple crops.
# dice_for_all_structures has one row per crop.
if forward_pass_result.loss is None:
raise ValueError(
"During training, the loss should always be computed, but the value is None."
)
loss = forward_pass_result.loss
# store metrics per batch
self.metrics.add_metric(MetricType.LOSS, loss)
for i, ground_truth_id in enumerate(
self.metrics.get_hue_names(include_default=False)):
for b in range(dice_for_all_classes.shape[0]):
self.metrics.add_metric(MetricType.DICE,
dice_for_all_classes[b, i].item(),
hue=ground_truth_id,
skip_nan_when_averaging=True)
self.metrics.add_metric(MetricType.VOXEL_COUNT,
foreground_voxels[i],
hue=ground_truth_id)
# store diagnostics per batch
center_indices = cropped_sample.center_indices
if isinstance(center_indices, torch.Tensor):
center_indices = center_indices.cpu().numpy()
self.metrics.add_diagnostics(MetricType.PATCH_CENTER.value,
np.copy(center_indices))
if self.train_val_params.in_training_mode:
# store the sample train patch from this epoch for visualization
if batch_index == self.example_to_save and self.model_config.store_dataset_sample:
_store_dataset_sample(self.model_config,
self.train_val_params.epoch,
forward_pass_result, cropped_sample)
return ModelForwardAndBackwardsOutputs(
loss=loss,
logits=forward_pass_result.posteriors,
labels=forward_pass_result.segmentations)
def get_epoch_results_and_store(self,
epoch_time_seconds: float) -> MetricsDict:
"""
Assembles all training results that were achieved over all minibatches, writes them to Tensorboard and
AzureML, and returns them as a MetricsDict object.
:param epoch_time_seconds: For diagnostics, this is the total time in seconds for training the present epoch.
:return: A dictionary that holds all metrics averaged over the epoch.
"""
self.metrics.add_metric(MetricType.SECONDS_PER_EPOCH,
epoch_time_seconds)
assert len(self.train_val_params.epoch_learning_rate
) == 1, "Expected a single entry for learning rate."
self.metrics.add_metric(MetricType.LEARNING_RATE,
self.train_val_params.epoch_learning_rate[0])
result = metrics.aggregate_segmentation_metrics(self.metrics)
metrics.store_epoch_metrics(self.azure_and_tensorboard_logger,
self.df_logger,
self.train_val_params.epoch, result,
self.train_val_params.epoch_learning_rate,
self.model_config)
return result
