cellmap_flow.finetune.adaptation

How a finetune adapts the model: a LoRA adapter, or every weight.

Everything that differs between the two lives here, behind one interface:

prepare: make the model trainable – wrap it in a fresh adapter, or unfreeze every weight – folding in any adapter it already carries;
train_mode: train mode, less what stays frozen;
teacher: the model as it was before this run, for distillation;
initial_state / reset: put the model back where training started, in place (the trainer’s retry after a NaN);
restart: the model the next iteration of a job trains (the CLI’s restart, which for LoRA re-makes the adapter so its rank can change);
trainable_state / checkpoint: what a checkpoint keeps;
export: what is written for serving, under export_name;
merge: the finetune folded into plain weights (export_merged).

The request decides the strategy once, when the model is prepared (--lora-r, 0 for full). After that the model does (strategy_for): a restart can change args.lora_r but not the model.

The adapter’s names depend on the module tree it was trained on – BatchLoopWrapper’s model. prefix included – so the wrappers stay in lora_wrapper, and this module only calls them. peft is imported lazily, as everywhere: a full finetune never needs it.

Classes

`AdaptationStrategy`	Base class for protocol classes.
`LoraStrategy`	Train a LoRA adapter on the frozen model; export only the adapter.
`FullStrategy`	Train every weight; export the whole state dict.

Functions

`is_peft_model`(→ bool)	Whether `model` carries a LoRA adapter (is a PeftModel); False without peft.
`cpu_state_copy`(→ Dict[str, torch.Tensor])	A CPU copy of `model`'s state dict, to reset a full finetune to.
`frozen_teacher_copy`(→ torch.nn.Module)	A frozen, eval-mode copy of `model`: the distillation teacher of a full finetune.
`strategy_for`(→ AdaptationStrategy)	The strategy `model` is trained with: LoRA if it carries an adapter, else full.

Module Contents

cellmap_flow.finetune.adaptation.is_peft_model(model) → bool: Whether model carries a LoRA adapter (is a PeftModel); False without peft.

cellmap_flow.finetune.adaptation.cpu_state_copy(model: torch.nn.Module) → Dict[str, torch.Tensor]: A CPU copy of model’s state dict, to reset a full finetune to.

cellmap_flow.finetune.adaptation.frozen_teacher_copy(model: torch.nn.Module) → torch.nn.Module: A frozen, eval-mode copy of model: the distillation teacher of a full finetune.

class cellmap_flow.finetune.adaptation.AdaptationStrategy

Base class for protocol classes.

Protocol classes are defined as:

class Proto(Protocol):
    def meth(self) -> int:
        ...

Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing).

For example:

class C:
    def meth(self) -> int:
        return 0

def func(x: Proto) -> int:
    return x.meth()

func(C())  # Passes static type check

See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as:

class GenProto(Protocol[T]):
    def meth(self) -> T:
        ...

kind: Literal['lora', 'full']

export_name: str

r: int

teacher_model: torch.nn.Module | None

prepare(model: torch.nn.Module) → torch.nn.Module

train_mode(model: torch.nn.Module) → None

teacher(model: torch.nn.Module) → ContextManager[torch.nn.Module]

initial_state(model: torch.nn.Module) → dict | None

reset(model: torch.nn.Module, initial_state) → torch.nn.Module

restart(model: torch.nn.Module, initial_state) → torch.nn.Module

trainable_state(model: torch.nn.Module) → dict

checkpoint(model: torch.nn.Module, optimizer, scaler, is_best: bool) → dict | None

export(model: torch.nn.Module, export_dir: pathlib.Path, path=None) → pathlib.Path

merge(model: torch.nn.Module) → torch.nn.Module

class cellmap_flow.finetune.adaptation.LoraStrategy(r: int, alpha: float, dropout: float, min_channels=None, target_modules=None)

Train a LoRA adapter on the frozen model; export only the adapter.

kind = 'lora'

export_name = 'lora_adapter'

teacher_model = None

r

alpha

dropout

min_channels = None

target_modules = None

prepare(model): model in a fresh adapter; one it already carries is folded into its weights first.

train_mode(model)

Train mode, except for the frozen base’s norm layers.

LoRA freezes the base, but model.train() would also put its BatchNorm layers in train mode: they would normalize by each (tiny) batch and keep updating running statistics that are not part of the adapter, so the model served in this process would drift from adapter + fresh base, and the base would not be frozen after all. They stay in eval mode, as they are served.

teacher(model): The model as it was before this run: itself, adapter off. Costs nothing.

initial_state(model): None: LoRA resets by re-initialising its adapter, not by reloading weights.

reset(model, initial_state=None): Re-initialise the adapter in place: B to zero, so it adds nothing, and A as peft does.

restart(model, initial_state=None)

A fresh adapter, with this strategy’s rank, on the base the old one was trained on.

The old adapter is unloaded, not merged: the next iteration starts from the model the job started from.

trainable_state(model): The adapter’s weights: the parameters that train.

checkpoint(model, optimizer, scaler, is_best)

The adapter’s weights, with the optimizer and scaler states, so a run can resume.

Only the trainable parameters: writing the full 800M+ param base to disk at every checkpoint is what this avoids.

export(model, export_dir, path=None): save_pretrained into export_dir/lora_adapter (or path, the adapter’s own directory).

static merge(model)

The adapter folded into the base weights: a plain module, no peft needed to serve it.

peft merges most layers itself. For a 1x1x1 Conv3d (an affinity head) lora_wrapper computes the delta instead, since peft’s shortcut for 1x1 kernels is conv2d-only and fails on it.

class cellmap_flow.finetune.adaptation.FullStrategy(teacher_model: torch.nn.Module | None = None)

Train every weight; export the whole state dict.

Measured against LoRA r=64 on mito-aff-unet-setup-16 (2026-09-23): faster per step (0.50 vs 0.90 s), lower memory (31 vs 50 GB at batch 8), and lower training loss at every checkpoint – the adapter’s savings are in parameters, which is not where this model’s cost is. The export is a full state dict under full_finetune/, served via FinetuneModelConfig(weights_path=…).

kind = 'full'

export_name = 'full_finetune'

r = 0

teacher_model = None

prepare(model): Every weight trainable; an adapter the model carries (a finetuned base) is folded into them first.

train_mode(model): Train mode throughout: a full finetune trains its norm layers too.

teacher(model)

A frozen copy of the weights training starts from.

Made at the first call – the trainer’s, before any step – and kept: the CLI hands it to the next iteration’s trainer, so a restart neither copies the model again nor distils toward weights an earlier iteration already changed. It costs one extra set of parameters on the device, and no activations, since the teacher runs under no_grad.

initial_state(model): The weights training starts from, on the CPU: what resetting it means.

reset(model, initial_state): Load the starting weights back, in place.

restart(model, initial_state)

The same model, back at the weights the job started from (initial_state).

Every iteration starts from them, as a LoRA restart starts from the base: nothing an earlier iteration learned, or diverged into, carries over.

trainable_state(model): Every weight.

checkpoint(model, optimizer, scaler, is_best)

Only the best epoch’s weights, without optimizer state (so no resume); None otherwise.

Every parameter is trainable, so a LoRA-style checkpoint would be the whole model plus two Adam moments – ~9.5 GB for an 800M-param UNet, at every periodic checkpoint. The best weights are what the export writes: the trainer’s save_adapter loads them first.

export(model, export_dir, path=None)

The whole state dict, to export_dir/full_finetune/model_state_dict.pt.

There FinetuneModelConfig(weights_path=…) expects it. path (an adapter’s own directory) does not apply.

static merge(model): model itself: its weights are the finetune.

cellmap_flow.finetune.adaptation.strategy_for(model: torch.nn.Module, lora_r: int | None = None, *, alpha=None, dropout=None, min_channels=None, target_modules=None, teacher_model: torch.nn.Module | None = None) → AdaptationStrategy

The strategy model is trained with: LoRA if it carries an adapter, else full.

The model decides, not the request. A restart can ask for another rank, or for none, but the model – which the inference server shares – is built once, so a job never switches between the two.

For a LoRA model, lora_r (when positive), alpha, dropout, min_channels and target_modules set the adapter restart makes; what is not given is the active adapter’s own (rank, alpha, dropout) or the default. teacher_model is a full finetune’s existing teacher, to reuse.