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 (models.geometry.trained_at_geometry).

Attributes

logger

Classes

FinetuneModelConfig

Configuration class for a finetuned model.

Module Contents

cellmap_flow.models.configs.finetune.logger

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

property normalizes_input

a finetune serves through the same preprocessing.

Type:

Its base’s

finetune_modes()

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

declared_geometry()

Its base’s, at the voxel sizes the finetune was trained at (trained_at_geometry); None when the base declares none.

output_graph()

Its base’s: a finetune’s weights (an adapter beside each layer, or the whole network’s) change what its layers compute, not the op that ends it, and the module it serves (PEFT’s, around the base’s trainable module) is not one a torch.fx trace reads.

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.