cellmap_flow.finetune.model_loading

The trainable nn.Module behind a model config.

A served model is often not trainable as it stands: the cellmap and Hugging Face models are TorchScript, which LoRA cannot wrap and autograd cannot train. cellmap_model.train() rebuilds them as an unflattened torch.export module, which is what the trainer needs. The adapter and full-finetune exports are keyed by the names in that exact module tree (BatchLoopWrapper’s model. prefix included), so training and serving must build it the same way: both call load_trainable_model (serving through FinetuneModelConfig).

The finetune CLI used to do this itself for fly/dacapo/huggingface/script models only, so a type: cellmap model – which is exactly what export_merged produces – or a finetuned model registered back into the dashboard could not be finetuned at all.

Attributes

logger

Functions

decode_model_entry(→ dict)	A model entry given on the command line: JSON, or encode_to_str() of it.
model_config_from_entry(entry[, name])	Build a ModelConfig from a model entry, as ModelConfig.to_dict() writes it.
load_trainable_model(→ torch.nn.Module)	The module to train for model_config.
root_base_model_dict(→ dict)	The innermost non-finetune model entry under model_config.

Module Contents

cellmap_flow.finetune.model_loading.logger

cellmap_flow.finetune.model_loading.decode_model_entry(text: str) → dict

A model entry given on the command line: JSON, or encode_to_str() of it.

cellmap_flow.finetune.model_loading.model_config_from_entry(entry: dict, name=None)

Build a ModelConfig from a model entry, as ModelConfig.to_dict() writes it.

Keys the class does not take are dropped first: to_dict() surfaces some for display (FinetuneModelConfig copies its base model’s channels and voxel sizes up, say), and the constructor would reject them.

cellmap_flow.finetune.model_loading.load_trainable_model(model_config) → torch.nn.Module

The module to train for model_config.

• TorchScript (cellmap, Hugging Face): the unflattened module from cellmap_model.train(), wrapped in BatchLoopWrapper when it is fixed at batch 1.

• A finetuned model (type: finetune): the model it serves – the base tree with its LoRA adapter attached (a PeftModel), or with its full finetuned weights loaded. Wrapping it in a new adapter folds the old one in first (lora_wrapper._merge_existing_adapters), so training continues from the finetuned model rather than from its base.

• A type whose trainable_model() gives one (bioimage, cellpose): that module.

• Anything else: model_config.config.model as it is.

cellmap_flow.finetune.model_loading.root_base_model_dict(model_config) → dict

The innermost non-finetune model entry under model_config.

A full finetune’s weights replace every parameter, so what it needs from its base is only the module tree – and a finetune-of-a-finetune base would hand the weights a PeftModel, whose names no longer match.