cellmap_flow.finetune.lora_wrapper

Generic LoRA wrapper for PyTorch models.

This module provides automatic detection of adaptable layers and wraps PyTorch models with LoRA (Low-Rank Adaptation) adapters using the HuggingFace PEFT library.

LoRA enables efficient finetuning by training only a small number of additional parameters (typically 1-2% of the original model) while keeping the base model frozen.

Attributes

logger

Classes

`BatchLoopWrapper`	Wraps a model with fixed batch_size=1 (e.g. UnflattenedModule from
`SequentialWrapper`	Wrapper for Sequential models to make them compatible with PEFT.

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

Module Contents

cellmap_flow.finetune.lora_wrapper.logger

cellmap_flow.finetune.lora_wrapper.detect_adaptable_layers(model: torch.nn.Module, include_patterns: List[str] | None = None, exclude_patterns: List[str] | None = None, min_channels: int = 0) → 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’]
min_channels – Skip layers whose narrower side (in or out) is below this. 0 adapts everything. Narrow layers are where LoRA is expensive for nothing: PEFT builds lora_A as a full-kernel conv Cin -> r, so on a 16-channel layer at full resolution the adapter is 4x the FLOPs of the layer it adapts and runs bandwidth-bound. On mito-aff-unet-setup-16 the seven layers under 96 channels hold ~1% of the adapter’s parameters and cost 41% of every training step (0.095 -> 0.056 s at 178^3).

Returns:

List of layer names suitable for LoRA adaptation

class cellmap_flow.finetune.lora_wrapper.BatchLoopWrapper(model: torch.nn.Module)

Wraps a model with fixed batch_size=1 (e.g. UnflattenedModule from torch.export without dynamic shapes) so it accepts arbitrary batch sizes by looping over the batch dim.

model

forward(x, *args, **kwargs)

class cellmap_flow.finetune.lora_wrapper.SequentialWrapper(model: torch.nn.Module)

Wrapper for Sequential models to make them compatible with PEFT.

PEFT expects models to accept **kwargs, but Sequential only accepts positional args. This wrapper provides that interface.

model

forward(x=None, input_ids=None, **kwargs)

cellmap_flow.finetune.lora_wrapper.wrap_model_with_lora(model: torch.nn.Module, target_modules: List[str] | None = None, lora_r: int = 64, lora_alpha: int = 128, lora_dropout: float = 0.1, modules_to_save: List[str] | None = None, task_type: str | None = None, lora_min_channels: int = 0) → 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)
lora_min_channels – When auto-detecting, skip layers narrower than this on either side; see detect_adaptable_layers. Ignored when target_modules is given explicitly. Default 0 (adapt all). A model with a lora_exclude_patterns attribute (name substrings, as detect_adaptable_layers’ exclude_patterns) has those layers left out of auto-detection too, besides the norm layers: a trainable_model() sets it for layers its network uses in a way an adapter would not reach.
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.lora_wrapper.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.lora_wrapper.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.lora_wrapper.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"
... )