cellmap_flow.finetune.lora_trainer

LoRA finetuning trainer for CellMap-Flow models.

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

Attributes

logger

Classes

`DiceLoss`	Dice Loss for segmentation tasks.
`CombinedLoss`	Combined Dice + BCE loss for better convergence.
`MarginLoss`	Margin-based loss for sparse/scribble annotations.
`LoRAFinetuner`	Trainer for finetuning models with LoRA adapters.

Module Contents

cellmap_flow.finetune.lora_trainer.logger

class cellmap_flow.finetune.lora_trainer.DiceLoss(smooth: float = 1.0)

Dice Loss for segmentation tasks.

Dice loss is effective for imbalanced datasets where the target class may be sparse (e.g., mitochondria in EM images).

Formula: 1 - (2 * |X ∩ Y| + smooth) / (|X| + |Y| + smooth)

smooth = 1.0

apply_sigmoid = True

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor

Compute Dice loss.

Parameters:

pred – Predictions (B, C, Z, Y, X) - raw logits or probabilities
target – Targets (B, C, Z, Y, X) - binary masks [0, 1]
mask – Optional mask (B, 1, Z, Y, X) - if provided, only compute loss on masked regions

Returns:

Dice loss value (scalar)

class cellmap_flow.finetune.lora_trainer.CombinedLoss(dice_weight: float = 0.5, bce_weight: float = 0.5)

Combined Dice + BCE loss for better convergence.

Uses both Dice loss (for overlap) and BCE loss (for pixel-wise accuracy).

dice_loss

bce_loss

dice_weight = 0.5

bce_weight = 0.5

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor

Compute combined loss.

Parameters:

pred – Predictions (B, C, Z, Y, X) - raw logits
target – Targets (B, C, Z, Y, X) - binary masks [0, 1]
mask – Optional mask (B, 1, Z, Y, X) - if provided, only compute loss on masked regions

Returns:

Combined loss value (scalar)

class cellmap_flow.finetune.lora_trainer.MarginLoss(margin: float = 0.3, balance_classes: bool = False)

Margin-based loss for sparse/scribble annotations.

Only penalizes predictions on the wrong side of a margin threshold. For post-sigmoid outputs in [0, 1]: - Foreground (target=1): loss = relu(threshold - pred)^2, threshold = 1 - margin - Background (target=0): loss = relu(pred - margin)^2 - No loss when prediction is already correct with sufficient confidence.

margin = 0.3

balance_classes = False

apply_sigmoid = True

forward(pred: torch.Tensor, target: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor

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 = 0.0, distillation_all_voxels: bool = False, margin: float = 0.3, balance_classes: bool = False, target_transform=None)

Trainer for finetuning models with LoRA adapters.

Features: - Mixed precision (FP16) training for memory efficiency - Gradient accumulation to simulate larger batch sizes - Checkpointing with best model tracking - Progress logging - Partial annotation support (mask unannotated regions)

Parameters:

model – PEFT model with LoRA adapters
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 FP16 training (default: True)
loss_type – Loss function (“dice”, “bce”, or “combined”)
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.
target_transform – Optional TargetTransform instance that converts raw annotations to (target, mask) pairs. Overrides mask_unannotated when provided. See cellmap_flow.finetune.target_transforms.

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 = 0.0

distillation_all_voxels = False

balance_classes = False

target_transform = None

optimizer

scaler

current_epoch = 0

global_step = 0

best_loss

training_stats = []

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_model.pth”

save_adapter(adapter_path: str | None = None)

Save only the LoRA adapter (not the full model).

Automatically loads the best checkpoint weights before saving so the exported adapter reflects the best training epoch.

Parameters:: adapter_path – Path to save adapter. If None, uses output_dir/lora_adapter

load_checkpoint(checkpoint_path: str)

Load training checkpoint to resume training.

Parameters:: checkpoint_path – Path to checkpoint file