cellmap_flow.models.configs.finetune

FinetuneModelConfig: a base model with a finetune’s weights on top.

The weights are a LoRA adapter or, from a full finetune, a whole state dict. Either is loaded onto the module the trainer trained (finetune.model_loading.load_trainable_model), and the served model keeps the base model’s geometry, at the voxel sizes the finetune was trained at when those differ (trained_at_geometry).

Classes

FinetuneModelConfig

Configuration class for a finetuned model.

Functions

trained_at_geometry(→ dict)

The base model's geometry at the voxel sizes a finetune was trained at.

Module Contents

cellmap_flow.models.configs.finetune.trained_at_geometry(base_cfg, input_voxel_size=None, output_voxel_size=None) → dict

The base model’s geometry at the voxel sizes a finetune was trained at.

A volume made without resampling lies on the raw level nearest the model’s voxel size, and the trainer read that level as it is: the model learnt that scale, so the finetuned model really is a model at the level’s voxel size (16 nm data read as 8 nm by an 8 nm model gives a 16 nm model). Declaring it, the served model reads that level exactly, instead of the nearest one relabelled with a warning. The shapes in nm scale with the voxel size; the voxel counts are the base’s.

Returns input_voxel_size, output_voxel_size, read_shape and write_shape; each size None keeps the base’s.

class cellmap_flow.models.configs.finetune.FinetuneModelConfig(lora_adapter_path: str = None, base_model: dict = None, name: str = None, scale=None, weights_path: str = None, input_voxel_size: tuple = None, output_voxel_size: tuple = None)

Configuration class for a finetuned model.

Wraps any base ModelConfig with the finetune’s weights, which are exactly one of: a LoRA adapter (lora_adapter_path), applied with PEFT, or a full finetune’s state dict (weights_path, from a run with --lora-r 0), loaded strictly. The base model is built by its own ModelConfig, as the trainable module the trainer trained, and the weights go on top of it.

cli_name = 'finetune'

finetunable = True

lora_adapter_path = None

weights_path = None

base_model_dict = None

name = None

scale = None

input_voxel_size = None

output_voxel_size = None

property env

the weights go on the base model, so they need the base’s packages.

Type:

Its own env if it was given one, else its base model’s

property default_env

The base model type’s default_env, as the weights need the base model’s packages. A type that decides per model decides for the base.

property base_model_config

Lazily build the base ModelConfig from the stored dict.

property display_channel

a finetune serves the same channels.

Type:

Its base’s

finetune_modes()

Its base’s: continuing a finetune trains the same network.

trainable_model()

The module it serves (its base’s trainable module with the adapter or weights on), to train on from where the finetune left off; not config.model, which a type serving through its own process_chunk (cellpose) keeps as its own model object.

to_dict()

This config as a model entry, which registry.build_model rebuilds.

Surfaces key base model fields at the top level so the pipeline builder UI can display them alongside the finetune-specific fields.