cellmap_flow.finetune.trainer ============================= .. py:module:: cellmap_flow.finetune.trainer .. autoapi-nested-parse:: The finetune trainer, in its parts. - ``loop.Finetuner``: the epoch loop, the best epoch, the stop signal; - ``supervised.SupervisedLoss``: the supervised term, and how the model's output is read for it (``output_class_of``: off the network's graph, else ``probe_output_class``); - ``distillation.Distillation``: the term that holds the model to the one it started from, and its teacher; - ``recovery``: mixed precision and the way back to fp32, out-of-memory mitigation, the finite-gradient check; - ``diagnostics``: the gradient-flow watch and the single-class warning; - ``tensorboard.TensorBoard``: the run's event files. ``LoRAFinetuner`` is ``Finetuner``'s old name, kept for the CLI and tests. Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/cellmap_flow/finetune/trainer/diagnostics/index /autoapi/cellmap_flow/finetune/trainer/distillation/index /autoapi/cellmap_flow/finetune/trainer/loop/index /autoapi/cellmap_flow/finetune/trainer/recovery/index /autoapi/cellmap_flow/finetune/trainer/supervised/index /autoapi/cellmap_flow/finetune/trainer/tensorboard/index Attributes ---------- .. autoapisummary:: cellmap_flow.finetune.trainer.LoRAFinetuner Classes ------- .. autoapisummary:: cellmap_flow.finetune.trainer.Finetuner Package Contents ---------------- .. py:class:: Finetuner(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, output_dir: str, options: Optional[cellmap_flow.finetune.options.TrainingOptions] = None, *, device: Optional[str] = None, target_transform=None, teacher_model: Optional[torch.nn.Module] = None, initial_state: Optional[Dict[str, torch.Tensor]] = None, tb_start_step: int = 0, tb_start_epoch: int = 0, stop_requested: Optional[Callable[[], bool]] = None, **settings) Trains a model on user corrections: a LoRA adapter, or every weight. Which of the two the model decides: a PEFT model trains its adapter, a plain module trains in full (see adaptation.strategy_for). Features: - Mixed precision on CUDA: bf16 where the GPU has it, else fp16 with a GradScaler; it falls back to fp32 when the model NaNs under either - Gradient accumulation to simulate larger batch sizes - The learning rate eased from the form's down to LEARNING_RATE_FLOOR of it over the run, or kept at the form's (learning_rate_schedule, learning_rate_factor) - Partial annotation support (mask unannotated regions) - Distillation toward the starting model, on unlabeled voxels, on all voxels, or on the good regions' rehearsal patches - One channel of a multi-channel model trained (select_channel), every other one held to the starting model's prediction - A best checkpoint, chosen by the supervised loss alone; the export loads it first - Recovery from an OOM (halve the batch, then drop distillation and the hold) and from a NaN under mixed precision (restart in fp32) - TensorBoard logging, continuing across a job's iterations :param model: The model to train: a PEFT model (LoRA) or a plain module (full finetune) :param dataloader: DataLoader for training data :param output_dir: Directory to save checkpoints and logs :param options: The training settings (finetune/options.py); the defaults if None. Every one of them is read here: the epochs, batch, learning rate and its schedule, the loss and its margin, smoothing and class balance, the distillation weight and scope, mask_unannotated, select_channel, mixed_precision, tensorboard, and output_class, what the model's output is (None reads it as logits; the training job probes the model when nothing declared it) :param device: Training device ("cuda" or "cpu", auto-detected if None) :param target_transform: Optional TargetTransform instance that converts raw annotations to (target, mask) pairs. Overrides mask_unannotated when provided. See cellmap_flow.finetune.target_transforms. :param teacher_model: A full finetune's frozen teacher from a previous iteration, to reuse :param initial_state: A full finetune's starting weights, to reset to; copied from the model if None :param tb_start_step: Where the previous iteration's TensorBoard curves ended :param tb_start_epoch: Where the previous iteration's TensorBoard curves ended :param stop_requested: Asked before each epoch, the first included, whether to stop there (Stop Early); the training job's says whether the dashboard asked (session_loop.TrainingSession). None: never :param \*\*settings: Fields of TrainingOptions to set over ``options`` (``num_epochs=3``, say) .. rubric:: Examples >>> lora_model = wrap_model_with_lora(model) >>> dataloader = create_dataloader("corrections.zarr") >>> trainer = Finetuner(lora_model, dataloader, output_dir="output/fly_organelles_v1.1") >>> trainer.train() >>> trainer.save_adapter() .. py:attribute:: options .. py:attribute:: model .. py:attribute:: dataloader .. py:attribute:: output_dir .. py:attribute:: num_epochs .. py:attribute:: gradient_accumulation_steps .. py:attribute:: select_channel .. py:attribute:: stop_requested :value: None .. py:attribute:: device .. py:attribute:: strategy .. py:attribute:: initial_state .. py:attribute:: precision .. py:attribute:: supervised .. py:attribute:: distillation .. py:attribute:: current_epoch :value: 0 .. py:attribute:: global_step :value: 0 .. py:attribute:: best_loss .. py:attribute:: last_supervised_loss .. py:attribute:: last_others_loss .. py:attribute:: training_stats :value: [] .. py:attribute:: tb .. py:property:: distillation_lambda :type: float The distillation term's weight as used (see Distillation). .. py:property:: teacher_model :type: Optional[torch.nn.Module] A full finetune's distillation teacher, the frozen copy of its starting weights; else None. .. py:method:: log(msg) :staticmethod: Log to console (the job's tee writes the log file). Timestamped like the logger's lines so epoch duration can be read off the log: the 2026-09-23 A/B runs had none on the per-epoch summaries, and per-arm speed had to come from LSF's start/end times instead. The progress parsers in finetune/job_manager use unanchored re.findall, so the prefix does not affect them. .. py:method:: close() Close the TensorBoard writer; the CLI calls this when an iteration is done. .. py:property:: schedule_steps :type: int The optimizer steps the run takes, which the learning rate eases over: an epoch's batches over the gradient accumulation, its last partial accumulation a step of its own (see _train_epoch), times the epochs. .. py:method:: train() -> Dict[str, Any] Run the training loop. :returns: - final_loss: Final epoch loss - best_loss: Best loss achieved - total_epochs: Number of epochs trained - total_steps: Total training steps :rtype: Training statistics dictionary with .. py:method:: save_checkpoint(is_best: bool = False) Save training checkpoint. :param is_best: If True, saves as "best_checkpoint.pth", which save_adapter exports; otherwise as "checkpoint_epoch_.pth" .. py:method:: save_adapter(adapter_path: Optional[str] = None, export_dir: Optional[str] = None) Export the finetune: only the LoRA adapter, or a full finetune's weights. Automatically loads the best checkpoint weights before saving so the exported adapter reflects the best training epoch. Without one from this run -- no epoch supervised anything, or none ran -- it exports the weights as they are, those of the last epoch. :param adapter_path: Path to save adapter. If None, uses output_dir/lora_adapter. A full finetune ignores it and writes output_dir/full_finetune. :param export_dir: Instead, the directory to export into: the adapter goes to export_dir/lora_adapter, full weights to export_dir/full_finetune/model_state_dict.pt. The CLI gives every iteration its own. :returns: The adapter directory, or the full-finetune weights file. .. py:method:: load_checkpoint(checkpoint_path: str) Load training checkpoint to resume training. :param checkpoint_path: Path to checkpoint file .. py:data:: LoRAFinetuner