cellmap_flow.finetune

Human-in-the-loop finetuning for CellMap-Flow models.

This package provides lightweight LoRA-based finetuning for pre-trained models using user corrections as training data.

Submodules

Classes

`CorrectionDataset`	PyTorch Dataset for user corrections stored in Zarr format.
`LoRAFinetuner`	Trainer for finetuning models with LoRA adapters.
`DiceLoss`	Dice Loss for segmentation tasks.
`CombinedLoss`	Combined Dice + BCE loss for better convergence.

Functions

`detect_adaptable_layers`(→ List[str])	Automatically detect layers suitable for LoRA adaptation.
`wrap_model_with_lora`(→ torch.nn.Module)	Wrap a PyTorch model with LoRA adapters using HuggingFace PEFT.
`print_lora_parameters`(model)	Print statistics about trainable and total parameters in a LoRA model.
`load_lora_adapter`(→ torch.nn.Module)	Load a pretrained LoRA adapter into a base model.
`save_lora_adapter`(model, output_path)	Save only the LoRA adapter parameters (not the full model).
`create_dataloader`(→ torch.utils.data.DataLoader)

Package Contents

cellmap_flow.finetune.detect_adaptable_layers(model: torch.nn.Module, include_patterns: List[str] | None = None, exclude_patterns: List[str] | None = None) → List[str]

Automatically detect layers suitable for LoRA adaptation.

Searches for Conv2d, Conv3d, and Linear layers, filtering by name patterns. By default, only excludes batch/layer-norm style modules. Output/head layers are deliberately INCLUDED so the model can fully adapt its feature→output mapping for cross-domain finetuning. (Previously ‘final’, ‘head’, ‘output’ were excluded; that left the output projection frozen, which prevented learning when the base model’s predictions on the target dataset were poor.)

Parameters:

model – PyTorch model to inspect
include_patterns – List of regex patterns for layer names to include If None, includes all Conv/Linear layers
exclude_patterns – List of substrings for layer names to exclude Default: [‘bn’, ‘norm’]

Returns:

List of layer names suitable for LoRA adaptation

cellmap_flow.finetune.wrap_model_with_lora(model: torch.nn.Module, target_modules: List[str] | None = None, lora_r: int = 8, lora_alpha: int = 16, lora_dropout: float = 0.1, modules_to_save: List[str] | None = None, task_type: str | None = None) → torch.nn.Module

Wrap a PyTorch model with LoRA adapters using HuggingFace PEFT.

This creates a PEFT model with LoRA adapters on specified layers. The base model is frozen, and only LoRA parameters are trainable.

Parameters:

model – PyTorch model to wrap (e.g., UNet, CNN)
target_modules – List of layer names to adapt. If None, auto-detects.
lora_r – LoRA rank (number of low-rank dimensions) Higher = more capacity, more parameters Typical values: 4-32, default 8
lora_alpha – LoRA alpha (scaling factor) Controls strength of LoRA updates Typical: 2*r, default 16
lora_dropout – Dropout probability for LoRA layers (0.0-0.5, default 0.1)
modules_to_save – Additional modules to make trainable (e.g., final layer)
task_type – PEFT task type. Options: - “FEATURE_EXTRACTION” (default, for general models) - “SEQ_CLS” (sequence classification) - “TOKEN_CLS” (token classification) - “CAUSAL_LM” (causal language modeling)

Returns:

PEFT model with LoRA adapters

Raises:

ImportError – If peft library is not installed
ValueError – If no adaptable layers found

Examples

>>> # Auto-detect and wrap all Conv/Linear layers
>>> lora_model = wrap_model_with_lora(model, lora_r=8)

>>> # Wrap specific layers with custom config
>>> lora_model = wrap_model_with_lora(
...     model,
...     target_modules=["encoder.conv1", "encoder.conv2"],
...     lora_r=16,
...     lora_alpha=32,
...     modules_to_save=["final_conv"]
... )

>>> # Check trainable parameters
>>> print_lora_parameters(lora_model)

cellmap_flow.finetune.print_lora_parameters(model: torch.nn.Module)

Print statistics about trainable and total parameters in a LoRA model.

Parameters:: model – PEFT model with LoRA adapters

Examples

>>> lora_model = wrap_model_with_lora(model)
>>> print_lora_parameters(lora_model)
Trainable params: 294,912 (1.2% of total)
Total params: 24,567,890

cellmap_flow.finetune.load_lora_adapter(model: torch.nn.Module, adapter_path: str, is_trainable: bool = False) → torch.nn.Module

Load a pretrained LoRA adapter into a base model.

Parameters:

model – Base PyTorch model (without LoRA)
adapter_path – Path to saved LoRA adapter directory
is_trainable – If True, adapter parameters are trainable (for continued training) If False, adapter parameters are frozen (for inference)

Returns:

PEFT model with loaded adapter

Examples

>>> # Load adapter for inference
>>> model = load_lora_adapter(
...     base_model,
...     "models/fly_organelles/v1.1.0/lora_adapter"
... )

>>> # Load adapter for continued training
>>> model = load_lora_adapter(
...     base_model,
...     "models/fly_organelles/v1.1.0/lora_adapter",
...     is_trainable=True
... )

cellmap_flow.finetune.save_lora_adapter(model: torch.nn.Module, output_path: str)

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

This saves only the trained LoRA weights (~5-20 MB) rather than the entire model (~200-500 MB).

Parameters:

model – PEFT model with LoRA adapters
output_path – Directory to save adapter

Examples

>>> save_lora_adapter(
...     lora_model,
...     "models/fly_organelles/v1.1.0/lora_adapter"
... )

class cellmap_flow.finetune.CorrectionDataset(corrections_zarr_path: str, patch_shape: Tuple[int, int, int] | None = None, augment: bool = True, model_name: str | None = None)

PyTorch Dataset for user corrections stored in Zarr format.

Loads raw EM data and corrected masks from corrections.zarr/, with optional 3D augmentation.

Parameters:

corrections_zarr_path – Path to corrections.zarr directory
patch_shape – Shape of patches to extract (Z, Y, X) If None, uses full correction size
augment – Whether to apply 3D augmentation
model_name – If specified, only load corrections for this model

Examples

>>> dataset = CorrectionDataset(
...     "test_corrections.zarr",
...     patch_shape=(64, 64, 64),
...     augment=True
... )
>>> print(f"Dataset size: {len(dataset)}")
>>> raw, target = dataset[0]
>>> print(f"Raw shape: {raw.shape}, Target shape: {target.shape}")

corrections_path

patch_shape = None

augment = True

model_name = None

corrections = []

cellmap_flow.finetune.create_dataloader(corrections_zarr_path: str, batch_size: int = 2, patch_shape: Tuple[int, int, int] | None = None, augment: bool = True, num_workers: int = 4, shuffle: bool = True, model_name: str | None = None) → torch.utils.data.DataLoader

class cellmap_flow.finetune.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

class cellmap_flow.finetune.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.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)