diff --git a/distillation/fast_nnunet_primus_distillation_export_onnx.py b/distillation/fast_nnunet_primus_distillation_export_onnx.py new file mode 100644 index 0000000..2ef1f8b --- /dev/null +++ b/distillation/fast_nnunet_primus_distillation_export_onnx.py @@ -0,0 +1,351 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +## Title: Fast-nnUNet — Primus distillation ONNX export entry point +## Authors: Justin Lee +## Description: Export a distilled Primus student checkpoint to ONNX. + +import argparse +import json +import os +import sys + +import numpy as np +import torch +from batchgenerators.utilities.file_and_folder_operations import isfile, join + +nnunet_dir = os.path.join(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, nnunet_dir) + +from nnunetv2.paths import nnUNet_preprocessed, nnUNet_raw, nnUNet_results +from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels +from nnunetv2.utilities.plans_handling.plans_handler import PlansManager +from primus_distillation_trainer import ( + PRIMUS_TEACHER_SPECS, + LitePrimusStudent, + reduce_primus_dims, +) + + +def get_dataset_name_from_id(dataset_id): + try: + dataset_id = int(dataset_id) + except ValueError: + return dataset_id + for dataset_dir in os.listdir(nnUNet_raw): + if dataset_dir.startswith(f"Dataset{dataset_id:03d}_"): + return dataset_dir + default_name = f"Dataset{dataset_id:03d}" + print( + f"Warning: No dataset directory with ID {dataset_id} found in {nnUNet_raw}, " + f"using default name {default_name}" + ) + return default_name + + +def export_primus_distillation_to_onnx( + dataset_id, + configuration="3d_fullres", + fold=0, + teacher_size="M", + feature_reduction_factor=2, + checkpoint_name="checkpoint_final.pth", + output_path=None, + device=None, + dynamic_axes=True, + input_shape=None, + simplify_onnx=False, + verbose=False, + use_da5=False, +): + """Export a distilled Primus student model to ONNX. + + Args: + dataset_id: Dataset ID, e.g. 793. + configuration: nnUNet configuration. Primus is 3D-only. + fold: Student model fold to export. + teacher_size: Teacher Primus size used during training (S/B/M/L). Drives the student + architecture reconstruction so the checkpoint weights load. + feature_reduction_factor: Same value used during training. + checkpoint_name: Filename inside the fold folder. + output_path: Custom .onnx path. Auto-named if None. + device: Torch device. Defaults to CUDA if available, else CPU. + dynamic_axes: If True, mark batch + spatial axes as dynamic. + input_shape: Custom dummy input shape; defaults to (1, C, *patch_size). + simplify_onnx: Run onnx-simplifier afterwards if installed. + verbose: Extra logging. + use_da5: Set if the checkpoint was trained with DA5 (only affects the trainer name + used for the folder path). + """ + if teacher_size not in PRIMUS_TEACHER_SPECS: + raise ValueError( + f"teacher_size must be one of {sorted(PRIMUS_TEACHER_SPECS)}; got {teacher_size!r}" + ) + + dataset_name = get_dataset_name_from_id(dataset_id) + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(device) + + trainer_name = ( + "nnUNetDistillationPrimusTrainerDA5" if use_da5 else "nnUNetDistillationPrimusTrainer" + ) + model_folder = join( + nnUNet_results, dataset_name, f"{trainer_name}__nnUNetPlans__{configuration}" + ) + model_folder_fold = join(model_folder, f"fold_{fold}") + if not os.path.exists(model_folder_fold): + raise FileNotFoundError(f"Model folder does not exist: {model_folder_fold}") + + checkpoint_path = join(model_folder_fold, checkpoint_name) + if not isfile(checkpoint_path): + raise FileNotFoundError(f"Checkpoint file does not exist: {checkpoint_path}") + + plans_file = join(nnUNet_preprocessed, dataset_name, "nnUNetPlans.json") + if not isfile(plans_file): + raise FileNotFoundError(f"Plans file does not exist: {plans_file}") + dataset_json_file = join(nnUNet_raw, dataset_name, "dataset.json") + if not isfile(dataset_json_file): + raise FileNotFoundError(f"Dataset json file does not exist: {dataset_json_file}") + + with open(plans_file, "r") as f: + plans = json.load(f) + with open(dataset_json_file, "r") as f: + dataset_json = json.load(f) + + plans_manager = PlansManager(plans) + cfg_manager = plans_manager.get_configuration(configuration) + patch_size = tuple(cfg_manager.patch_size) + if len(patch_size) != 3: + raise ValueError(f"Primus is 3D-only; got patch_size {patch_size}") + if any(p % 8 != 0 for p in patch_size): + raise ValueError( + f"Primus requires patch_size divisible by 8; got {patch_size}" + ) + + num_input_channels = determine_num_input_channels(plans_manager, cfg_manager, dataset_json) + num_output_channels = plans_manager.get_label_manager(dataset_json).num_segmentation_heads + print( + f"Configuration: {configuration}. patch_size={patch_size}. " + f"Input channels: {num_input_channels}, output channels: {num_output_channels}" + ) + + # Reconstruct student architecture from teacher_size + reduction_factor + spec = PRIMUS_TEACHER_SPECS[teacher_size] + s_embed, s_depth, s_heads = reduce_primus_dims( + spec["embed_dim"], spec["depth"], spec["num_heads"], feature_reduction_factor + ) + print( + f"Primus student (teacher {teacher_size}, r={feature_reduction_factor}): " + f"embed_dim={s_embed}, depth={s_depth}, num_heads={s_heads}" + ) + + model = LitePrimusStudent( + input_channels=num_input_channels, + num_classes=num_output_channels, + embed_dim=s_embed, + depth=s_depth, + num_heads=s_heads, + patch_size=patch_size, + ) + + # Load checkpoint + try: + import numpy.core.multiarray as _np_core + if hasattr(torch.serialization, "add_safe_globals"): + torch.serialization.add_safe_globals([_np_core.scalar]) + checkpoint = torch.load(checkpoint_path, map_location=device) + else: + checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) + except Exception: + checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) + + if "network_weights" in checkpoint: + state_dict = checkpoint["network_weights"] + elif "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + else: + raise ValueError( + f"Checkpoint {checkpoint_path} contains neither 'network_weights' nor 'state_dict'" + ) + + cleaned = {} + for k, v in state_dict.items(): + cleaned[k[len("module."):]] = v if k.startswith("module.") else cleaned.setdefault(k, v) + # Simpler: strip `module.` and load with strict=False + cleaned = {(k[len("module."):] if k.startswith("module.") else k): v for k, v in state_dict.items()} + load_result = model.load_state_dict(cleaned, strict=False) + if load_result.missing_keys or load_result.unexpected_keys: + if verbose: + print( + f"Load issues — missing={len(load_result.missing_keys)}, " + f"unexpected={len(load_result.unexpected_keys)}" + ) + + model.eval() + model.to(device) + total_params = sum(p.numel() for p in model.parameters()) + print(f"Model: {total_params:,} params (~{total_params * 4 / 1024 / 1024:.1f} MB)") + + if output_path is None: + output_dir = join(model_folder_fold, "exported_models") + os.makedirs(output_dir, exist_ok=True) + da5_suffix = "_da5" if use_da5 else "" + output_path = join( + output_dir, + f"primus_{teacher_size}_{configuration}_fold{fold}_r{feature_reduction_factor}{da5_suffix}.onnx", + ) + + if input_shape is None: + input_shape = (1, num_input_channels, *patch_size) + torch.manual_seed(42) + dummy_input = torch.randn(input_shape, dtype=torch.float32, device=device) + + if dynamic_axes: + dynamic_axes_dict = { + "input": {0: "batch_size", 2: "depth", 3: "height", 4: "width"}, + "output": {0: "batch_size", 2: "depth", 3: "height", 4: "width"}, + } + else: + dynamic_axes_dict = None + + print(f"Exporting Primus student to ONNX (input: {tuple(dummy_input.shape)})...") + with torch.no_grad(): + torch_output = model(dummy_input) + + torch.onnx.export( + model, + dummy_input, + output_path, + export_params=True, + opset_version=17, + do_constant_folding=True, + input_names=["input"], + output_names=["output"], + dynamic_axes=dynamic_axes_dict, + training=torch.onnx.TrainingMode.EVAL, + verbose=verbose, + ) + print(f"Exported to: {output_path}") + + # Validate + try: + import onnx + from onnx import checker + from onnxruntime import InferenceSession + + onnx_model = onnx.load(output_path) + checker.check_model(onnx_model) + + ort_session = InferenceSession(output_path, providers=["CPUExecutionProvider"]) + ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.cpu().numpy()} + ort_outputs = ort_session.run(None, ort_inputs) + torch_output_np = torch_output.detach().cpu().numpy() + abs_diff = np.abs(torch_output_np - ort_outputs[0]) + max_diff = float(np.max(abs_diff)) + mean_diff = float(np.mean(abs_diff)) + if max_diff < 0.01: + print(f" Excellent match (max={max_diff:.6f}, mean={mean_diff:.6f})") + elif max_diff < 0.5: + print(f" Good match (max={max_diff:.6f}, mean={mean_diff:.6f})") + else: + print(f" Difference detected (max={max_diff:.6f}, mean={mean_diff:.6f})") + except ImportError as exc: + print(f"Skipping ONNX validation ({exc}); install with `pip install -e '.[onnx]'`.") + + if simplify_onnx: + try: + from onnxsim import simplify # type: ignore + print("Simplifying ONNX model...") + onnx_model = onnx.load(output_path) + simplified, ok = simplify(onnx_model) + if ok: + onnx.save(simplified, output_path) + print(f"Simplified ONNX written to: {output_path}") + else: + print("onnx-simplifier reported failure; keeping original.") + except ImportError: + print("onnx-simplifier not installed; skipping --simplify path.") + + return output_path + + +def main(): + parser = argparse.ArgumentParser(description="Export distilled Primus student to ONNX") + parser.add_argument("-d", "--dataset_id", type=str, required=True, help="Dataset ID") + parser.add_argument( + "-c", + "--configuration", + type=str, + default="3d_fullres", + help="nnUNet configuration (Primus is 3D-only) (default: 3d_fullres)", + ) + parser.add_argument("-f", "--fold", type=int, default=0, help="Fold number (default: 0)") + parser.add_argument( + "-ts", + "--teacher_size", + type=str, + default="M", + choices=list(PRIMUS_TEACHER_SPECS), + help="Primus teacher size used during training (default: M)", + ) + parser.add_argument( + "-r", "--reduction_factor", type=int, default=2, help="Reduction factor used during training (default: 2)" + ) + parser.add_argument( + "-cp", + "--checkpoint", + type=str, + default="checkpoint_final.pth", + help="Checkpoint filename (default: checkpoint_final.pth)", + ) + parser.add_argument("-o", "--output_path", type=str, help="Custom output .onnx path") + parser.add_argument("-d_device", "--device", type=str, help="Torch device, e.g. cuda:0") + parser.add_argument( + "-no_da", + "--no_dynamic_axes", + action="store_true", + help="Disable dynamic batch+spatial axes (export fixed-shape ONNX)", + ) + parser.add_argument( + "-is", + "--input_shape", + type=int, + nargs="+", + help="Custom input shape, e.g. 1 1 128 128 128", + ) + parser.add_argument("-sim", "--simplify", action="store_true", help="Run onnx-simplifier") + parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging") + parser.add_argument( + "-da5", "--use_da5", action="store_true", help="Set if checkpoint was trained with DA5" + ) + + args = parser.parse_args() + + export_primus_distillation_to_onnx( + dataset_id=args.dataset_id, + configuration=args.configuration, + fold=args.fold, + teacher_size=args.teacher_size, + feature_reduction_factor=args.reduction_factor, + checkpoint_name=args.checkpoint, + output_path=args.output_path, + device=args.device, + dynamic_axes=not args.no_dynamic_axes, + input_shape=tuple(args.input_shape) if args.input_shape else None, + simplify_onnx=args.simplify, + verbose=args.verbose, + use_da5=args.use_da5, + ) + + +if __name__ == "__main__": + main() diff --git a/distillation/fast_nnunet_primus_distillation_train.py b/distillation/fast_nnunet_primus_distillation_train.py new file mode 100644 index 0000000..8d44e52 --- /dev/null +++ b/distillation/fast_nnunet_primus_distillation_train.py @@ -0,0 +1,341 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +## Title: Fast-nnUNet — Primus distillation training entry point +## Authors: Justin Lee +## Description: CLI driver for distilling an upstream nnunetv2 Primus teacher into a smaller +## Primus student. Mirrors the standard / ResEnc distillation entry points and +## adds a `-ts/--teacher_size` flag to pick S / B / M / L. + +import argparse +import json +import os +import sys + +import torch +from batchgenerators.utilities.file_and_folder_operations import join + +# Make sure the vendored nnunetv2 (on main) and our trainer module are importable +nnunet_dir = os.path.join(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, nnunet_dir) + +from nnunetv2.paths import nnUNet_preprocessed, nnUNet_raw, nnUNet_results +from primus_distillation_trainer import ( + PRIMUS_TEACHER_SPECS, + PRIMUS_TEACHER_TRAINER_CLASS, + nnUNetDistillationPrimusTrainer, + nnUNetDistillationPrimusTrainerDA5, +) + + +def get_dataset_name_from_id(dataset_id): + try: + dataset_id = int(dataset_id) + except ValueError: + return dataset_id + for dataset_dir in os.listdir(nnUNet_raw): + if dataset_dir.startswith(f"Dataset{dataset_id:03d}_"): + return dataset_dir + default_name = f"Dataset{dataset_id:03d}" + print( + f"Warning: No dataset directory with ID {dataset_id} found in {nnUNet_raw}, " + f"using default name {default_name}" + ) + return default_name + + +def run_primus_distillation_training( + dataset_id, + configuration="3d_fullres", + fold=0, + teacher_size="M", + teacher_model_folder=None, + teacher_folds=None, + teacher_checkpoint_name="checkpoint_final.pth", + teacher_plans_identifier="nnUNetPlans", + alpha=0.3, + temperature=3.0, + feature_reduction_factor=2, + warmup_duration_whole_net=50, + continue_training=False, + val_with_mirroring=True, + rotate_training_folds=False, + rotate_folds_frequency=5, + device=None, + epochs=1000, + use_da5=False, +): + """Run Primus knowledge-distillation training. + + Args mostly mirror ``fast_nnunet_resenc_distillation_train.run_resenc_distillation_training``. + Primus-specific extras: + teacher_size: One of "S", "B", "M", "L". Picks the upstream Primus teacher trainer. + warmup_duration_whole_net: Number of epochs of linear LR warmup (matches upstream + AbstractPrimus default of 50). + """ + dataset_name = get_dataset_name_from_id(dataset_id) + + # Default teacher folder from teacher_size + if teacher_model_folder is None: + teacher_trainer_cls = PRIMUS_TEACHER_TRAINER_CLASS[teacher_size] + teacher_model_folder = join( + nnUNet_results, + dataset_name, + f"{teacher_trainer_cls}__{teacher_plans_identifier}__{configuration}", + ) + print( + f"Teacher model folder not specified, using upstream Primus default for size " + f"{teacher_size}: {teacher_model_folder}" + ) + + if not os.path.exists(teacher_model_folder): + raise FileNotFoundError( + f"Primus teacher model folder does not exist: {teacher_model_folder}" + ) + + # Auto-detect teacher folds if not specified + if teacher_folds is None: + teacher_folds = [] + for item in os.listdir(teacher_model_folder): + if os.path.isdir(join(teacher_model_folder, item)) and item.startswith("fold_"): + try: + fold_num = int(item.split("_")[1]) + if os.path.exists(join(teacher_model_folder, item, teacher_checkpoint_name)): + teacher_folds.append(fold_num) + except ValueError: + continue + teacher_folds.sort() + if not teacher_folds: + print("Warning: No available Primus teacher folds found, falling back to fold_0") + teacher_folds = [0] + else: + print(f"Detected {len(teacher_folds)} available Primus teacher folds: {teacher_folds}") + else: + for tf in teacher_folds: + cp = join(teacher_model_folder, f"fold_{tf}", teacher_checkpoint_name) + if not os.path.exists(cp): + raise FileNotFoundError(f"Primus teacher checkpoint missing: {cp}") + + dataset_json_file = join(nnUNet_raw, dataset_name, "dataset.json") + if not os.path.exists(dataset_json_file): + raise FileNotFoundError(f"Dataset json file does not exist: {dataset_json_file}") + + student_plans_file = join(nnUNet_preprocessed, dataset_name, "nnUNetPlans.json") + if not os.path.exists(student_plans_file): + raise FileNotFoundError( + f"Student plans file does not exist: {student_plans_file}, " + "make sure preprocessing for the configuration has been run" + ) + + with open(student_plans_file, "r") as f: + plans = json.load(f) + with open(dataset_json_file, "r") as f: + dataset_json = json.load(f) + + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(device) + + trainer_cls = ( + nnUNetDistillationPrimusTrainerDA5 if use_da5 else nnUNetDistillationPrimusTrainer + ) + trainer = trainer_cls( + plans=plans, + configuration=configuration, + fold=fold, + dataset_json=dataset_json, + teacher_model_folder=teacher_model_folder, + teacher_fold=teacher_folds, + teacher_checkpoint_name=teacher_checkpoint_name, + alpha=alpha, + temperature=temperature, + feature_reduction_factor=feature_reduction_factor, + rotate_training_folds=rotate_training_folds, + rotate_folds_frequency=rotate_folds_frequency, + teacher_size=teacher_size, + warmup_duration_whole_net=warmup_duration_whole_net, + device=device, + ) + + teacher_spec = PRIMUS_TEACHER_SPECS[teacher_size] + print(f"\n============ Primus Knowledge Distillation Training Configuration ============") + print(f"Dataset name: {dataset_name}") + print(f"Configuration: {configuration}") + print(f"Training fold: {fold}") + print( + f"Teacher Primus size: {teacher_size} " + f"(embed_dim={teacher_spec['embed_dim']}, depth={teacher_spec['depth']}, " + f"num_heads={teacher_spec['num_heads']})" + ) + print(f"Teacher model folder: {teacher_model_folder}") + print(f"Teacher model folds: {teacher_folds}") + print(f"Teacher model checkpoint: {teacher_checkpoint_name}") + print(f"Distillation loss weight alpha: {alpha}") + print(f"Distillation temperature: {temperature}") + print(f"Feature reduction factor (embed_dim/depth/heads divisor): {feature_reduction_factor}") + print(f"Warmup duration (epochs): {warmup_duration_whole_net}") + print(f"Continue training: {continue_training}") + print(f"Validation with mirroring: {val_with_mirroring}") + if rotate_training_folds: + print(f"Rotating folds every {rotate_folds_frequency} epochs") + print(f"Device: {device}") + print(f"Maximum training epochs: {epochs}") + print(f"================================================================================\n") + + trainer.num_epochs = epochs + if not val_with_mirroring: + trainer.inference_allowed_mirroring_axes = [] + + expected_checkpoint_file = join(trainer.output_folder, "checkpoint_latest.pth") + checkpoint_exists = os.path.exists(expected_checkpoint_file) + + if continue_training and checkpoint_exists: + print(f"Continuing previous training, loading checkpoint: {expected_checkpoint_file}") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if not trainer.was_initialized: + trainer.initialize() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + trainer.load_student_checkpoint(expected_checkpoint_file) + print(f"Loaded checkpoint, resuming from epoch {trainer.current_epoch}") + else: + if continue_training and not checkpoint_exists: + print( + f"Warning: could not find checkpoint {expected_checkpoint_file}, " + "starting from scratch" + ) + if not trainer.was_initialized: + trainer.initialize() + + trainer.run_training() + trainer.perform_actual_validation(False) + + +def main(): + parser = argparse.ArgumentParser( + description="nnUNetv2 Primus knowledge-distillation training" + ) + parser.add_argument("-d", "--dataset_id", type=str, required=True, help="Dataset ID (e.g., 793)") + parser.add_argument( + "-c", + "--configuration", + type=str, + default="3d_fullres", + help="nnUNet configuration. Primus is 3D-only; patch_size must be divisible by 8 " + "(default: 3d_fullres)", + ) + parser.add_argument("-f", "--fold", type=int, default=0, help="Fold number for training (default: 0)") + parser.add_argument( + "-ts", + "--teacher_size", + type=str, + default="M", + choices=list(PRIMUS_TEACHER_SPECS), + help="Upstream Primus teacher size: S / B / M / L (default: M)", + ) + parser.add_argument( + "-t", + "--teacher_model_folder", + type=str, + help="Custom Primus teacher folder (auto-derived from --teacher_size if omitted)", + ) + parser.add_argument( + "-tf", + "--teacher_folds", + type=int, + nargs="+", + help="List of teacher fold numbers, e.g. 0 or 0 1 2 3 4. Default: auto-detect all available folds", + ) + parser.add_argument( + "-tcp", + "--teacher_checkpoint", + type=str, + default="checkpoint_final.pth", + help="Teacher model checkpoint filename (default: checkpoint_final.pth)", + ) + parser.add_argument( + "-tpl", + "--teacher_plans", + type=str, + default="nnUNetPlans", + help="Teacher plans identifier (default: nnUNetPlans)", + ) + parser.add_argument("-a", "--alpha", type=float, default=0.3, help="Distillation loss weight (default: 0.3)") + parser.add_argument("-temp", "--temperature", type=float, default=3.0, help="Distillation temperature (default: 3.0)") + parser.add_argument( + "-r", + "--reduction_factor", + type=int, + default=2, + help="Reduction factor applied to embed_dim/depth/num_heads (default: 2)", + ) + parser.add_argument( + "-w", + "--warmup_epochs", + type=int, + default=50, + help="Linear LR warmup duration in epochs (default: 50, matches upstream Primus)", + ) + parser.add_argument("-d_device", "--device", type=str, help="Device to use, e.g., \"cuda:0\"") + parser.add_argument("-c_continue", "--continue_training", action="store_true", help="Continue previous training") + parser.add_argument( + "-disable_mirroring", + "--disable_val_mirroring", + action="store_true", + help="Disable mirroring during validation", + ) + parser.add_argument( + "-rotate_folds", + "--rotate_training_folds", + action="store_true", + help="Rotate training folds periodically", + ) + parser.add_argument( + "-rotate_freq", + "--rotate_folds_frequency", + type=int, + default=5, + help="How often to rotate folds (in epochs, default: 5)", + ) + parser.add_argument("-e", "--epochs", type=int, default=1000, help="Maximum training epochs (default: 1000)") + parser.add_argument( + "--use_da5", + action="store_true", + help="Use DA5 strong data augmentation (recommended for small datasets)", + ) + + args = parser.parse_args() + + run_primus_distillation_training( + dataset_id=args.dataset_id, + configuration=args.configuration, + fold=args.fold, + teacher_size=args.teacher_size, + teacher_model_folder=args.teacher_model_folder, + teacher_folds=args.teacher_folds, + teacher_checkpoint_name=args.teacher_checkpoint, + teacher_plans_identifier=args.teacher_plans, + alpha=args.alpha, + temperature=args.temperature, + feature_reduction_factor=args.reduction_factor, + warmup_duration_whole_net=args.warmup_epochs, + continue_training=args.continue_training, + val_with_mirroring=not args.disable_val_mirroring, + rotate_training_folds=args.rotate_training_folds, + rotate_folds_frequency=args.rotate_folds_frequency, + device=args.device, + epochs=args.epochs, + use_da5=args.use_da5, + ) + + +if __name__ == "__main__": + main() diff --git a/distillation/primus_distillation_trainer.py b/distillation/primus_distillation_trainer.py new file mode 100644 index 0000000..42625dc --- /dev/null +++ b/distillation/primus_distillation_trainer.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +## Title: Fast-nnUNet — Primus knowledge-distillation trainer +## Authors: Justin Lee +## Description: Distillation trainer + lightweight student for the upstream nnunetv2 Primus +## transformer architecture (S/B/M/L teachers). + +"""Primus-flavored knowledge-distillation trainer. + +Provides: +- ``reduce_primus_dims``: shrink (embed_dim, depth, num_heads) by an integer factor while + preserving the ``embed_dim % num_heads == 0`` constraint required by Primus' attention. +- ``LitePrimusStudent``: thin wrapper around + ``dynamic_network_architectures.architectures.primus.Primus`` that accepts the reduced + hyperparameters. +- ``nnUNetDistillationPrimusTrainer`` / ``nnUNetDistillationPrimusTrainerDA5``: distillation + trainers that reuse the existing multi-teacher KL pipeline from + ``nnUNetDistillationTrainer`` but build a Primus student and use upstream Primus' AdamW + + warmup optimizer schedule. + +The teacher must be a trained upstream Primus model +(``nnUNet_Primus_{S,B,M,L}_Trainer__nnUNetPlans__``). Multi-teacher ensembling, +fold rotation, and DA5 augmentation are inherited unchanged from the existing distillation +trainer family. +""" + +import os +import sys +from typing import Tuple + +import torch +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP + +# Ensure the vendored nnunetv2 tree (on main) wins over a pip-installed copy. +# On the refactor branch this is a no-op because the module lives next to the entry-point +# scripts and the vendored tree is gone, but keeping the insertion is harmless. +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from dynamic_network_architectures.architectures.primus import Primus +from nnunetv2.training.lr_scheduler.warmup import ( + Lin_incr_LRScheduler, + PolyLRScheduler_offset, +) +from nnunetv2.training.nnUNetTrainer.variants.data_augmentation.nnUNetTrainerDA5 import ( + nnUNetTrainerDA5, +) +from nnunetv2.training.nnUNetTrainer.variants.nnUNetDistillationTrainer import ( + nnUNetDistillationTrainer, +) +from nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels + + +# Reference architecture specs for each upstream Primus size. +PRIMUS_TEACHER_SPECS = { + "S": dict(embed_dim=396, depth=12, num_heads=6), + "B": dict(embed_dim=792, depth=12, num_heads=12), + "M": dict(embed_dim=864, depth=16, num_heads=12), + "L": dict(embed_dim=1056, depth=24, num_heads=16), +} + +# Upstream Primus trainer class names, used to construct the default teacher folder. +PRIMUS_TEACHER_TRAINER_CLASS = { + "S": "nnUNet_Primus_S_Trainer", + "B": "nnUNet_Primus_B_Trainer", + "M": "nnUNet_Primus_M_Trainer", + "L": "nnUNet_Primus_L_Trainer", +} + + +def reduce_primus_dims( + embed_dim: int, depth: int, num_heads: int, factor: int +) -> Tuple[int, int, int]: + """Shrink Primus hyperparameters by ``factor`` while keeping the model valid. + + Primus' 3D rotary positional embedding requires ``head_dim`` (= embed_dim // num_heads) + to be divisible by 6. The reliable way to keep that invariant under reduction is to + **hold ``head_dim`` constant at the teacher's value** and shrink only ``num_heads`` + (and ``depth``): + + - ``num_heads`` is floor-divided by ``factor`` but never goes below 1. + - ``embed_dim`` is set to ``head_dim * num_heads`` so head_dim is preserved exactly. + - ``depth`` is floor-divided by ``factor`` but never goes below 2 — we still want at + least one transformer block with skip + norm. + """ + if factor < 1: + raise ValueError(f"reduction factor must be >= 1, got {factor}") + if num_heads < 1 or embed_dim % num_heads != 0: + raise ValueError( + f"teacher embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})" + ) + head_dim = embed_dim // num_heads + new_heads = max(num_heads // factor, 1) + new_embed = head_dim * new_heads + new_depth = max(depth // factor, 2) + return new_embed, new_depth, new_heads + + +class LitePrimusStudent(nn.Module): + """Primus with reduced embed_dim/depth/heads for knowledge distillation.""" + + def __init__( + self, + input_channels: int, + num_classes: int, + embed_dim: int, + depth: int, + num_heads: int, + patch_size: Tuple[int, ...], + patch_embed_size: Tuple[int, ...] = (8, 8, 8), + drop_path_rate: float = 0.2, + init_values: float = 0.1, + scale_attn_inner: bool = True, + ): + super().__init__() + if any(p % 8 != 0 for p in patch_size): + raise ValueError( + f"Primus requires patch_size divisible by 8; got {tuple(patch_size)}" + ) + if embed_dim % num_heads != 0: + raise ValueError( + f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})" + ) + self.input_channels = input_channels + self.num_classes = num_classes + self.embed_dim = embed_dim + self.depth = depth + self.num_heads = num_heads + self.patch_size = tuple(patch_size) + self.patch_embed_size = tuple(patch_embed_size) + self.network = Primus( + input_channels, + embed_dim, + tuple(patch_embed_size), + num_classes, + eva_depth=depth, + eva_numheads=num_heads, + input_shape=tuple(patch_size), + drop_path_rate=drop_path_rate, + scale_attn_inner=scale_attn_inner, + init_values=init_values, + ) + + def forward(self, x): + return self.network(x) + + +class nnUNetDistillationPrimusTrainer(nnUNetDistillationTrainer): + """Knowledge-distillation trainer with a Primus student. + + Inherits the multi-teacher / KL-loss / fold-rotation machinery from + ``nnUNetDistillationTrainer`` and swaps in: + - A Primus student instead of a CNN U-Net student. + - AdamW + warmup optimizer schedule (matches upstream ``AbstractPrimus``). + - Deep supervision disabled — Primus emits a single full-resolution map. + """ + + DEFAULT_WARMUP_DURATION_WHOLE_NET = 50 + + def __init__( + self, + plans, + configuration, + fold, + dataset_json, + teacher_model_folder=None, + teacher_fold=0, + teacher_checkpoint_name="checkpoint_final.pth", + alpha=0.3, + temperature=3.0, + feature_reduction_factor=2, + rotate_training_folds=False, + rotate_folds_frequency=5, + teacher_size="M", + warmup_duration_whole_net=DEFAULT_WARMUP_DURATION_WHOLE_NET, + device=torch.device("cuda"), + ): + if teacher_size not in PRIMUS_TEACHER_SPECS: + raise ValueError( + f"teacher_size must be one of {sorted(PRIMUS_TEACHER_SPECS)}; got {teacher_size!r}" + ) + # block_reduction_strategy is ignored for Primus (no residual blocks); we pass a + # placeholder so the parent constructor doesn't complain. + super().__init__( + plans=plans, + configuration=configuration, + fold=fold, + dataset_json=dataset_json, + teacher_model_folder=teacher_model_folder, + teacher_fold=teacher_fold, + teacher_checkpoint_name=teacher_checkpoint_name, + alpha=alpha, + temperature=temperature, + feature_reduction_factor=feature_reduction_factor, + block_reduction_strategy="keep", + rotate_training_folds=rotate_training_folds, + rotate_folds_frequency=rotate_folds_frequency, + device=device, + ) + + # Primus-specific overrides (parent's patched_init flipped some of these on). + self.enable_deep_supervision = False + self.teacher_size = teacher_size + self.warmup_duration_whole_net = warmup_duration_whole_net + self.training_stage = None + # Mirror upstream AbstractPrimus defaults. + self.initial_lr = 3e-4 + self.weight_decay = 5e-2 + self.my_init_kwargs.update( + { + "teacher_size": teacher_size, + "warmup_duration_whole_net": warmup_duration_whole_net, + } + ) + + # ------------------------------------------------------------------ + # Network architecture + # ------------------------------------------------------------------ + def build_network_architecture( + self, + architecture_class_name=None, + arch_init_kwargs=None, + arch_init_kwargs_req_import=None, + num_input_channels: int = None, + num_output_channels: int = None, + enable_deep_supervision: bool = False, + ): + if num_input_channels is None: + num_input_channels = determine_num_input_channels( + self.plans_manager, self.configuration_manager, self.dataset_json + ) + if num_output_channels is None: + num_output_channels = self.label_manager.num_segmentation_heads + + patch_size = tuple(self.configuration_manager.patch_size) + if len(patch_size) != 3: + raise ValueError( + f"Primus is 3D-only; got patch_size {patch_size} (configuration={self.configuration_name})" + ) + if any(p % 8 != 0 for p in patch_size): + raise ValueError( + f"Primus requires patch_size divisible by 8; got {patch_size}. " + "Pick a configuration whose patch_size matches, or adjust the plans." + ) + + teacher_spec = PRIMUS_TEACHER_SPECS[self.teacher_size] + s_embed, s_depth, s_heads = reduce_primus_dims( + teacher_spec["embed_dim"], + teacher_spec["depth"], + teacher_spec["num_heads"], + self.feature_reduction_factor, + ) + self.print_to_log_file( + f"Primus student: embed_dim={s_embed}, depth={s_depth}, num_heads={s_heads} " + f"(reduced from teacher {self.teacher_size} " + f"embed_dim={teacher_spec['embed_dim']}, depth={teacher_spec['depth']}, " + f"num_heads={teacher_spec['num_heads']}, factor={self.feature_reduction_factor})" + ) + return LitePrimusStudent( + input_channels=num_input_channels, + num_classes=num_output_channels, + embed_dim=s_embed, + depth=s_depth, + num_heads=s_heads, + patch_size=patch_size, + ) + + # ------------------------------------------------------------------ + # Optimizer / scheduler — mirrors upstream AbstractPrimus + # ------------------------------------------------------------------ + def configure_optimizers(self, stage: str = "warmup_all"): + if stage not in {"warmup_all", "train"}: + raise ValueError(f"unknown stage {stage!r}") + if self.training_stage == stage and self.optimizer is not None: + return self.optimizer, self.lr_scheduler + + params = ( + self.network.module.parameters() + if isinstance(self.network, DDP) + else self.network.parameters() + ) + if stage == "warmup_all": + optimizer = torch.optim.AdamW( + params, + self.initial_lr, + weight_decay=self.weight_decay, + amsgrad=False, + betas=(0.9, 0.98), + ) + lr_scheduler = Lin_incr_LRScheduler( + optimizer, self.initial_lr, self.warmup_duration_whole_net + ) + self.print_to_log_file( + f"Primus distill: warmup_all optimizer at epoch {self.current_epoch}" + ) + else: + optimizer = torch.optim.AdamW( + params, + self.initial_lr, + weight_decay=self.weight_decay, + amsgrad=False, + betas=(0.9, 0.98), + ) + lr_scheduler = PolyLRScheduler_offset( + optimizer, + self.initial_lr, + self.num_epochs, + self.warmup_duration_whole_net, + ) + self.print_to_log_file( + f"Primus distill: train optimizer at epoch {self.current_epoch}" + ) + self.training_stage = stage + return optimizer, lr_scheduler + + def set_deep_supervision_enabled(self, enabled: bool): + # Primus is single-resolution; this is a no-op. + return + + +class nnUNetDistillationPrimusTrainerDA5( + nnUNetDistillationPrimusTrainer, nnUNetTrainerDA5 +): + """Primus distillation trainer with DA5 strong data augmentation.""" + + def __init__(self, *args, **kwargs): + nnUNetDistillationPrimusTrainer.__init__(self, *args, **kwargs) + self.print_to_log_file( + "Using DA5 strong data augmentation for Primus knowledge distillation" + ) diff --git a/distillation/setup.py b/distillation/setup.py index 4df881d..5a22176 100644 --- a/distillation/setup.py +++ b/distillation/setup.py @@ -2,7 +2,7 @@ setup( name="nnunetv2_distillation", - version="1.2.3", + version="1.2.4", packages=find_packages(), install_requires=[ "torch>=1.6.0", @@ -12,15 +12,20 @@ 'console_scripts': [ 'nnUNetv2_distillation_train=fast_nnunet_distillation_train:main', 'nnUNetv2_resenc_distillation_train=fast_nnunet_resenc_distillation_train:main', + 'nnUNetv2_primus_distillation_train=fast_nnunet_primus_distillation_train:main', 'nnUNetv2_distillation_export_onnx=fast_nnunet_distillation_export_onnx:main', 'nnUNetv2_resenc_distillation_export_onnx=fast_nnunet_resenc_distillation_export_onnx:main', + 'nnUNetv2_primus_distillation_export_onnx=fast_nnunet_primus_distillation_export_onnx:main', ], }, py_modules=[ 'fast_nnunet_distillation_train', 'fast_nnunet_resenc_distillation_train', + 'fast_nnunet_primus_distillation_train', 'fast_nnunet_distillation_export_onnx', - 'fast_nnunet_resenc_distillation_export_onnx' + 'fast_nnunet_resenc_distillation_export_onnx', + 'fast_nnunet_primus_distillation_export_onnx', + 'primus_distillation_trainer', ], python_requires='>=3.7', author="Justin", diff --git a/docs/Distillation.md b/docs/Distillation.md index 0f84b50..57944c3 100644 --- a/docs/Distillation.md +++ b/docs/Distillation.md @@ -201,6 +201,59 @@ nnUNetv2_resenc_distillation_train -d DATASET_ID -f 0 -a 0.3 -temp 3.0 -r 2 --us nnUNetv2_resenc_distillation_train -d DATASET_ID -f 0 -a 0.3 -temp 3.0 -r 2 -tpl nnUNetResEncUNetLPlans -spl nnUNetResEncUNetMPlans --use_da5 ``` +#### Primus Knowledge Distillation + +Distill an upstream nnUNetv2 Primus transformer teacher into a smaller Primus student. Primus +is 3D-only and requires `patch_size` divisible by 8. + +Four teacher sizes are supported (select with `-ts / --teacher_size`): + +| Size | embed_dim | depth | num_heads | head_dim | +| --- | --- | --- | --- | --- | +| `S` | 396 | 12 | 6 | 66 | +| `B` | 792 | 12 | 12 | 66 | +| `M` (default) | 864 | 16 | 12 | 72 | +| `L` | 1056 | 24 | 16 | 66 | + +The student is produced by holding `head_dim` constant and dividing `num_heads` / `depth` by +`-r / --reduction_factor`. For example, `-ts M -r 2` yields a student with embed_dim=432, +depth=8, num_heads=6 (head_dim still 72, so the 3D rotary positional embedding stays valid). + +Teacher folder is auto-derived from `-ts`: +``` +{nnUNet_results}/{Dataset}/nnUNet_Primus_{S|B|M|L}_Trainer__nnUNetPlans__{configuration}/ +``` + +```bash +# 1) Train the upstream Primus teacher with stock nnUNetv2 first: +nnUNetv2_train DATASET_ID 3d_fullres 0 -tr nnUNet_Primus_M_Trainer +# (repeat for folds 1-4 if you want a multi-teacher ensemble) + +# 2) Distill into a Primus student (default teacher size M, reduction 2x): +nnUNetv2_primus_distillation_train -d DATASET_ID -f 0 -a 0.3 -temp 3.0 -r 2 + +# Pick a smaller / larger teacher: +nnUNetv2_primus_distillation_train -d DATASET_ID -ts S -f 0 -a 0.3 -temp 3.0 -r 2 +nnUNetv2_primus_distillation_train -d DATASET_ID -ts L -f 0 -a 0.3 -temp 3.0 -r 2 + +# Multi-teacher ensemble across folds: +nnUNetv2_primus_distillation_train -d DATASET_ID -ts M -f 0 -tf 0 1 2 3 4 -a 0.3 -temp 3.0 -r 2 + +# Continue previous training, custom warmup duration, DA5 augmentation: +nnUNetv2_primus_distillation_train -d DATASET_ID -ts M -f 0 -r 2 -w 50 -c_continue +nnUNetv2_primus_distillation_train -d DATASET_ID -ts M -f 0 -r 2 --use_da5 + +# Custom teacher folder (e.g. one of the BS8 / 96_BS1 variants): +nnUNetv2_primus_distillation_train -d DATASET_ID -ts M -f 0 -r 2 -t /path/to/custom/primus/teacher +``` + +Primus-specific notes: +- The student uses AdamW + linear warmup → polynomial schedule (matches upstream + `AbstractPrimus`); `-w / --warmup_epochs` controls the warmup duration (default 50). +- Deep supervision is disabled — Primus emits a single full-resolution map. +- Cascade configurations (`3d_cascade_fullres`) are accepted as long as the underlying + Primus teacher was trained on the same configuration with the lowres predictions in place. + Parameter description: - `-d, --dataset_id`: Dataset ID - `-f, --fold`: Fold number used to train the student model @@ -295,6 +348,32 @@ nnUNetv2_resenc_distillation_export_onnx -d DATASET_ID -f 0 -r 2 -da5 -fix nnUNetv2_resenc_distillation_export_onnx -d DATASET_ID -f 0 -r 2 -sim ``` +#### Primus Distillation Model Export + +```bash +# Basic export (default teacher size M, reduction 2x) +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -f 0 -r 2 + +# Pick a different teacher size used during training +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts S -f 0 -r 2 +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts L -f 0 -r 2 + +# Custom output path +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts M -f 0 -r 2 -o /path/to/primus.onnx + +# Custom input shape (must satisfy patch_size divisibility constraints of the teacher) +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts M -f 0 -r 2 -is 1 1 128 128 128 + +# Fixed-shape export (disable dynamic axes) +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts M -f 0 -r 2 -no_da + +# Export DA5-trained checkpoint +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts M -f 0 -r 2 -da5 + +# Simplify with onnx-simplifier +nnUNetv2_primus_distillation_export_onnx -d DATASET_ID -ts M -f 0 -r 2 -sim +``` + Parameter description: - `-d, --dataset_id`: Dataset ID - `-f, --fold`: Model fold number