cellmap_flow.finetune.lora_trainer

Finetuning trainer for CellMap-Flow models: a LoRA adapter, or every weight.

This module provides a trainer class for finetuning models using user corrections with mixed-precision training and gradient accumulation.

Attributes

logger

Classes

LoRAFinetuner

Trainer for finetuning a model on user corrections: a LoRA adapter, or

Module Contents

cellmap_flow.finetune.lora_trainer.logger
class cellmap_flow.finetune.lora_trainer.LoRAFinetuner(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, output_dir: str, learning_rate: float = 0.0001, num_epochs: int = 10, gradient_accumulation_steps: int = 1, use_mixed_precision: bool = True, loss_type: str = 'combined', device: str | None = None, select_channel: int | None = None, mask_unannotated: bool = True, label_smoothing: float = 0.0, distillation_lambda: float | None = None, distillation_all_voxels: bool = False, margin: float = 0.3, balance_classes: bool = False, slope_weight: float = 1.0, target_transform=None, tensorboard: bool = True, 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)

Trainer for finetuning a model on user corrections: a LoRA adapter, or every weight (a full finetune).

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

  • Partial annotation support (mask unannotated regions)

  • Distillation toward the starting model, on unlabeled voxels, on all voxels, or on the good regions’ rehearsal patches

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

  • learning_rate – Learning rate (default: 1e-4)

  • num_epochs – Number of training epochs (default: 10)

  • gradient_accumulation_steps – Steps to accumulate gradients (default: 1)

  • use_mixed_precision – Enable mixed precision on CUDA (default: True; off on the CPU)

  • loss_type – Loss function: “dice”, “bce”, “combined” (Dice + BCE), “mse”, “margin”, “interval” (a distance model on scribbles; needs an IntervalTargetTransform) or “flow” (a flow model such as Cellpose; needs a FlowTargetTransform)

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

  • select_channel – Optional channel index to select from multi-channel output (default: None)

  • mask_unannotated – If True (default), only compute loss on annotated regions (target > 0). Targets are shifted down by 1 (e.g., 1->0, 2->1) after masking. This allows partial annotations where 0=unannotated, 1=background, 2=foreground, etc. Ignored if target_transform is provided.

  • label_smoothing – s moves the targets to s/2 and 1 - s/2 (default: 0; forced to 0 for margin, interval and flow)

  • distillation_lambda – Weight of the distillation term. None means 1.0 when the dataset has good regions and 0 otherwise; an explicit 0 is honoured.

  • distillation_all_voxels – Distil on every voxel, not only the unlabeled ones (no effect when the dataset has good regions: they decide where)

  • margin – The margin of the “margin” loss (default: 0.3)

  • balance_classes – Weight foreground and background voxels equally

  • slope_weight – The weight of the interval loss’s slope limit (default: 1)

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

  • tensorboard – Write TensorBoard logs to output_dir/tensorboard (default: True)

  • 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

Examples

>>> lora_model = wrap_model_with_lora(model)
>>> dataloader = create_dataloader("corrections.zarr")
>>> trainer = LoRAFinetuner(
...     lora_model,
...     dataloader,
...     output_dir="output/fly_organelles_v1.1"
... )
>>> trainer.train()
>>> trainer.save_adapter()
model
dataloader
output_dir
num_epochs = 10
gradient_accumulation_steps = 1
use_mixed_precision = True
select_channel = None
mask_unannotated = True
label_smoothing = 0.0
distillation_lambda = None
distillation_all_voxels = False
balance_classes = False
target_transform = None
loss_type = 'combined'
margin = 0.3
learning_rate = 0.0001
strategy
initial_state
optimizer
amp_dtype = Ellipsis
scaler
current_epoch = 0
global_step = 0
best_loss
last_supervised_loss
training_stats = []
tb = None
tb_dir
tb_image_every = 5
property teacher_model: torch.nn.Module | None

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

The CLI hands it to the next iteration’s trainer. Setting it to None frees it (the OOM handler does, with distillation).

close()

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

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