From d4c89dbf0ea4c6c6882df497b68d6fb591a46b12 Mon Sep 17 00:00:00 2001 From: Nathan Molinier Date: Wed, 5 Aug 2026 14:38:35 -0400 Subject: [PATCH 1/2] save weights with different amount of augmentation applied during validation --- auglab/configs/transform_params_gpu.json | 4 + auglab/trainers/nnUNetTrainerDAExt.py | 165 ++++++++++++++++++++++- 2 files changed, 167 insertions(+), 2 deletions(-) diff --git a/auglab/configs/transform_params_gpu.json b/auglab/configs/transform_params_gpu.json index 397eb91..803eea5 100644 --- a/auglab/configs/transform_params_gpu.json +++ b/auglab/configs/transform_params_gpu.json @@ -205,5 +205,9 @@ }, "ZscoreNormalizationTransform": { "probability": 0.00 + }, + "ValidationAugmentationCheckpoints": { + "comment": "Percentage of validation samples that will have training augmentations applied. A separate checkpoint will be saved during training for each percentage.", + "checkpoints": [0, 0.5, 1.0] } } diff --git a/auglab/trainers/nnUNetTrainerDAExt.py b/auglab/trainers/nnUNetTrainerDAExt.py index b3caa55..c6ce086 100644 --- a/auglab/trainers/nnUNetTrainerDAExt.py +++ b/auglab/trainers/nnUNetTrainerDAExt.py @@ -6,6 +6,7 @@ import numpy as np import torch +from batchgenerators.utilities.file_and_folder_operations import join from batchgeneratorsv2.helpers.scalar_type import RandomScalar from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform @@ -19,9 +20,12 @@ from batchgeneratorsv2.transforms.utils.random import RandomTransform from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform +from nnunetv2.training.loss.dice import get_tp_fp_fn_tn from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer +from nnunetv2.utilities.collate_outputs import collate_outputs from nnunetv2.utilities.helpers import dummy_context from torch import autocast +from torch import distributed as dist from auglab import configs from auglab.trainers.utils import DownsampleSegForDSTransformCustom @@ -159,6 +163,13 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic self.transforms = AugTransformsGPU(json_path=json_path).to(self.device) print(f"Using AugLab GPU transforms with parameters from: {json_path}") + # Load JSON parameters for validation augmentation checkpoints + with open(json_path) as f: + config = json.load(f) + self.validation_augmentation_checkpoints = config.get("ValidationAugmentationCheckpoints", {}).get("checkpoints", [0]) + self.ema_dice_validation = [None] * len(self.validation_augmentation_checkpoints) + self.best_ema_dice_validation = [None] * len(self.validation_augmentation_checkpoints) + # Copy json transfrom parameters to output folder shutil.copy(json_path, os.path.join(self.output_folder, "transform_params_gpu_used_for_training.json")) @@ -299,8 +310,9 @@ def get_validation_transforms( # transforms.append(ZscoreNormalization()) - if deep_supervision_scales is not None: - transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) + # NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations + # if deep_supervision_scales is not None: + # transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales)) return ComposeTransforms(transforms) def train_step(self, batch: dict) -> dict: @@ -346,6 +358,155 @@ def train_step(self, batch: dict) -> dict: self.optimizer.step() return {"loss": l.detach().cpu().numpy()} + def validation_step(self, batch: dict, augmentation_prob: float) -> dict: + data = batch["data"] + target = batch["target"] + + data = data.to(self.device, non_blocking=True) + # Now target should be a single tensor, not a list + target = target.to(self.device, non_blocking=True) + # if isinstance(target, list): + # target = [i.to(self.device, non_blocking=True) for i in target] + # else: + # target = target.to(self.device, non_blocking=True) + + # Autocast can be annoying + # If the device_type is 'cpu' then it's slow as heck and needs to be disabled. + # If the device_type is 'mps' then it will complain that mps is not implemented, even if enabled=False is set. Whyyyyyyy. (this is why we don't make use of enabled=False) + # So autocast will only be active if we have a cuda device. + with autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context(): + # Apply GPU augmentations to full-resolution data/target + if torch.rand(1).item() < augmentation_prob: + data, target = self.transforms(data, target) + + # Create multi-scale targets for deep supervision after augmentation + deep_supervision_scales = self._get_deep_supervision_scales() + if deep_supervision_scales is not None: + ds_transform = DownsampleSegForDSTransformCustom(ds_scales=deep_supervision_scales) + target = ds_transform(target) + + output = self.network(data) + del data + l = self.loss(output, target) + + # we only need the output with the highest output resolution (if DS enabled) + if self.enable_deep_supervision: + output = output[0] + target = target[0] + + # the following is needed for online evaluation. Fake dice (green line) + axes = [0, *list(range(2, output.ndim))] + + if self.label_manager.has_regions: + predicted_segmentation_onehot = (torch.sigmoid(output) > 0.5).long() + else: + # no need for softmax + output_seg = output.argmax(1)[:, None] + predicted_segmentation_onehot = torch.zeros(output.shape, device=output.device, dtype=torch.float32) + predicted_segmentation_onehot.scatter_(1, output_seg, 1) + del output_seg + + if self.label_manager.has_ignore_label: + if not self.label_manager.has_regions: + mask = (target != self.label_manager.ignore_label).float() + # CAREFUL that you don't rely on target after this line! + target[target == self.label_manager.ignore_label] = 0 + else: + mask = ~target[:, -1:] if target.dtype == torch.bool else 1 - target[:, -1:] + # CAREFUL that you don't rely on target after this line! + target = target[:, :-1] + else: + mask = None + + tp, fp, fn, _ = get_tp_fp_fn_tn(predicted_segmentation_onehot, target, axes=axes, mask=mask) + + tp_hard = tp.detach().cpu().numpy() + fp_hard = fp.detach().cpu().numpy() + fn_hard = fn.detach().cpu().numpy() + if not self.label_manager.has_regions: + # if we train with regions all segmentation heads predict some kind of foreground. In conventional + # (softmax training) there needs tobe one output for the background. We are not interested in the + # background Dice + # [1:] in order to remove background + tp_hard = tp_hard[1:] + fp_hard = fp_hard[1:] + fn_hard = fn_hard[1:] + + return {"loss": l.detach().cpu().numpy(), "tp_hard": tp_hard, "fp_hard": fp_hard, "fn_hard": fn_hard} + + def compute_validation_metrics(self, val_outputs: list[dict]): + """ + Based on on_validation_epoch_end nnUNetTrainer + """ + outputs_collated = collate_outputs(val_outputs) + tp = np.sum(outputs_collated["tp_hard"], 0) + fp = np.sum(outputs_collated["fp_hard"], 0) + fn = np.sum(outputs_collated["fn_hard"], 0) + + if self.is_ddp: + world_size = dist.get_world_size() + + tps = [None for _ in range(world_size)] + dist.all_gather_object(tps, tp) + tp = np.vstack([i[None] for i in tps]).sum(0) + + fps = [None for _ in range(world_size)] + dist.all_gather_object(fps, fp) + fp = np.vstack([i[None] for i in fps]).sum(0) + + fns = [None for _ in range(world_size)] + dist.all_gather_object(fns, fn) + fn = np.vstack([i[None] for i in fns]).sum(0) + + losses_val = [None for _ in range(world_size)] + dist.all_gather_object(losses_val, outputs_collated["loss"]) + loss_here = np.vstack(losses_val).mean() + else: + loss_here = np.mean(outputs_collated["loss"]) + + global_dc_per_class = [2 * i / (2 * i + j + k) for i, j, k in zip(tp, fp, fn)] + mean_fg_dice = np.nanmean(global_dc_per_class) + val_losses = loss_here + return mean_fg_dice, global_dc_per_class, val_losses + + def run_training(self): + self.on_train_start() + + for _epoch in range(self.current_epoch, self.num_epochs): + self.on_epoch_start() + + self.on_train_epoch_start() + train_outputs = [self.train_step(next(self.dataloader_train)) for _batch_id in range(self.num_iterations_per_epoch)] + self.on_train_epoch_end(train_outputs) + + with torch.no_grad(): + self.on_validation_epoch_start() + val_outputs = [[]] * len(self.validation_augmentation_checkpoints) + for _batch_id in range(self.num_val_iterations_per_epoch): + batch = next(self.dataloader_val) + for prob_id, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): + val_outputs[prob_id].append(self.validation_step(batch, augmentation_prob)) + for val_idx, val_output in enumerate(val_outputs): + mean_fg_dice, global_dc_per_class, val_losses = self.compute_validation_metrics(val_output) + self.ema_dice_validation[val_idx] = ( + self.ema_dice_validation[val_idx] * 0.9 + 0.1 * mean_fg_dice + if self.ema_dice_validation[val_idx] is not None + else mean_fg_dice + ) + if val_idx == 0: + self.logger.log("mean_fg_dice", mean_fg_dice, self.current_epoch) + self.logger.log("dice_per_class_or_region", global_dc_per_class, self.current_epoch) + self.logger.log("val_losses", val_losses, self.current_epoch) + + self.on_epoch_end() + for val_idx, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): + ema_dice = self.ema_dice_validation[val_idx] + if ema_dice is None or ema_dice > self.best_ema_dice_validation[val_idx]: + self.best_ema_dice_validation[val_idx] = ema_dice + self.save_checkpoint(join(self.output_folder, f"checkpoint_best_validation_aug_{augmentation_prob}.pth")) + + self.on_train_end() + class nnUNetTrainerDAExtHybrid(nnUNetTrainer): def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dict, device: torch.device = torch.device("cuda")): From f09897bc5d9a9c0667ff56c340357b67070c5b11 Mon Sep 17 00:00:00 2001 From: Nathan Molinier Date: Wed, 5 Aug 2026 15:33:56 -0400 Subject: [PATCH 2/2] fix bug --- auglab/trainers/nnUNetTrainerDAExt.py | 34 +++++++++++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/auglab/trainers/nnUNetTrainerDAExt.py b/auglab/trainers/nnUNetTrainerDAExt.py index c6ce086..f24b102 100644 --- a/auglab/trainers/nnUNetTrainerDAExt.py +++ b/auglab/trainers/nnUNetTrainerDAExt.py @@ -434,6 +434,35 @@ def validation_step(self, batch: dict, augmentation_prob: float) -> dict: return {"loss": l.detach().cpu().numpy(), "tp_hard": tp_hard, "fp_hard": fp_hard, "fn_hard": fn_hard} + @staticmethod + def _valaug_sidecar_path(checkpoint_path: str) -> str: + root, ext = os.path.splitext(checkpoint_path) + return f"{root}_valaug{ext}" + + def save_checkpoint(self, filename: str) -> None: + super().save_checkpoint(filename) + if self.local_rank == 0 and not self.disable_checkpointing: + torch.save( + { + "ema_dice_validation": self.ema_dice_validation, + "best_ema_dice_validation": self.best_ema_dice_validation, + "validation_augmentation_checkpoints": self.validation_augmentation_checkpoints, + }, + self._valaug_sidecar_path(filename), + ) + + def load_checkpoint(self, filename_or_checkpoint: Union[dict, str]) -> None: + super().load_checkpoint(filename_or_checkpoint) + if isinstance(filename_or_checkpoint, str): + sidecar = self._valaug_sidecar_path(filename_or_checkpoint) + if os.path.isfile(sidecar): + state = torch.load(sidecar, map_location="cpu") + # Only adopt persisted EMA state if the configured checkpoints list is unchanged; + # otherwise the per-index slots no longer correspond and we restart tracking. + if state.get("validation_augmentation_checkpoints") == self.validation_augmentation_checkpoints: + self.ema_dice_validation = state["ema_dice_validation"] + self.best_ema_dice_validation = state["best_ema_dice_validation"] + def compute_validation_metrics(self, val_outputs: list[dict]): """ Based on on_validation_epoch_end nnUNetTrainer @@ -481,7 +510,7 @@ def run_training(self): with torch.no_grad(): self.on_validation_epoch_start() - val_outputs = [[]] * len(self.validation_augmentation_checkpoints) + val_outputs = [[] for _ in range(len(self.validation_augmentation_checkpoints))] for _batch_id in range(self.num_val_iterations_per_epoch): batch = next(self.dataloader_val) for prob_id, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): @@ -501,7 +530,8 @@ def run_training(self): self.on_epoch_end() for val_idx, augmentation_prob in enumerate(self.validation_augmentation_checkpoints): ema_dice = self.ema_dice_validation[val_idx] - if ema_dice is None or ema_dice > self.best_ema_dice_validation[val_idx]: + best = self.best_ema_dice_validation[val_idx] + if ema_dice is not None and (best is None or ema_dice > best): self.best_ema_dice_validation[val_idx] = ema_dice self.save_checkpoint(join(self.output_folder, f"checkpoint_best_validation_aug_{augmentation_prob}.pth"))