cellmap_flow.finetune.trainer

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

• cellmap_flow.finetune.trainer.diagnostics
• cellmap_flow.finetune.trainer.distillation
• cellmap_flow.finetune.trainer.loop
• cellmap_flow.finetune.trainer.recovery
• cellmap_flow.finetune.trainer.supervised
• cellmap_flow.finetune.trainer.tensorboard

Attributes

LoRAFinetuner

Classes

Finetuner

Trains a model on user corrections: a LoRA adapter, or every weight.

Package Contents

class cellmap_flow.finetune.trainer.Finetuner(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, output_dir: str, options: cellmap_flow.finetune.options.TrainingOptions | None = None, *, device: str | None = None, target_transform=None, teacher_model: torch.nn.Module | None = None, initial_state: Dict[str, torch.Tensor] | None = None, tb_start_step: int = 0, tb_start_epoch: int = 0, stop_requested: Callable[[], bool] | None = 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

Parameters:

• model – The model to train: a PEFT model (LoRA) or a plain module (full finetune)

• dataloader – DataLoader for training data

• output_dir – Directory to save checkpoints and logs

• 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)

• device – Training device (“cuda” or “cpu”, auto-detected if None)

• target_transform – Optional TargetTransform instance that converts raw annotations to (target, mask) pairs. Overrides mask_unannotated when provided. See cellmap_flow.finetune.target_transforms.

• teacher_model – A full finetune’s frozen teacher from a previous iteration, to reuse

• initial_state – A full finetune’s starting weights, to reset to; copied from the model if None

• tb_start_step – Where the previous iteration’s TensorBoard curves ended

• tb_start_epoch – Where the previous iteration’s TensorBoard curves ended

• 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

• **settings – Fields of TrainingOptions to set over options (num_epochs=3, say)

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()

options

model

dataloader

output_dir

num_epochs

gradient_accumulation_steps

select_channel

stop_requested = None

device

strategy

initial_state

precision

supervised

distillation

current_epoch = 0

global_step = 0

best_loss

last_supervised_loss

last_others_loss

training_stats = []

tb

property distillation_lambda: float

The distillation term’s weight as used (see Distillation).

property teacher_model: torch.nn.Module | None

A full finetune’s distillation teacher, the frozen copy of its starting weights; else None.

static log(msg)

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.

close()

Close the TensorBoard writer; the CLI calls this when an iteration is done.

property schedule_steps: 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.

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

Return type:

Training statistics dictionary with

save_checkpoint(is_best: bool = False)

Save training checkpoint.

Parameters:

is_best – If True, saves as “best_checkpoint.pth”, which save_adapter exports; otherwise as “checkpoint_epoch_<N>.pth”

save_adapter(adapter_path: str | None = None, export_dir: str | None = 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.

Parameters:

• adapter_path – Path to save adapter. If None, uses output_dir/lora_adapter. A full finetune ignores it and writes output_dir/full_finetune.

• 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.

load_checkpoint(checkpoint_path: str)

Load training checkpoint to resume training.

Parameters:

checkpoint_path – Path to checkpoint file

cellmap_flow.finetune.trainer.LoRAFinetuner