cellmap_flow.finetune.export_merged

Fold a finetune into the base model’s weights and re-export it at a larger tile.

Why this exists

Every cellmap UNet on HuggingFace ships two inference artefacts, and both are locked to the 178^3 -> 56^3 tile they were exported at: the .pt2 (what type: finetune serves) refuses any other input, and the TorchScript (what type: huggingface serves) silently crops a larger input back to 56^3. A valid-padding UNet spends 32 input voxels per output voxel at that tile; at 306^3 -> 184^3 it spends 4.6. The eager model.pt is not shape locked, so: load it, fold the LoRA adapter (or load full-finetune weights), and export the plain module again at the tile you want. The result is an ordinary CellmapModel folder served with type: cellmap – no PEFT, no fly_organelles, no adapter at inference time.

Tile rule: these UNets pool 2x2x2 three times, so the tile must be 178 + 16k (290, 306, 322, …). Any other size shifts the pooling phase and the output no longer matches what the 178 tile computes at the same place (measured: 298 differs by up to 2.3 in logit space; 290/306 differ by 0.0). The 378 in the repos’ metadata.json is NOT such a size.

Usage:

python -m cellmap_flow.finetune.export_merged \
    --repo cellmap/mito-aff-unet-setup-16 \
    --lora-adapter-path <run>/lora_adapter \
    --tile 306 --output <session>/exports/<dataset>/<class>/lora_aug_tile306

then in a cellmap-flow yaml:

models:
  - type: cellmap
    folder_path: <that output dir>
    name: <anything>

Requires fly_organelles importable (the eager pickle references its classes) – present in the cellmap-flow-finetune env.

Attributes

`logger`
`POOL_FACTOR`
`TRAIN_TILE`
`CONTEXT`
`REL_TOL`

Functions

`valid_tile`(tile)	True when `tile` keeps the pooling phase of the 178 training tile.
`load_eager_base`(folder_path)	The eager `model.pt` of a CellmapModel folder, eval mode, on CPU.
`apply_finetune`(trainable[, lora_adapter_path, ...])	Return a plain nn.Module with the finetune folded in.
`check_equivalence`(module, reference, tile[, seed])	Difference between `module` and `reference` at 178^3, and between the
`export_folder`(module, output, input_shape, ...)	Write `module` into `output` as cellmap-models exports a model, the
`main`([argv])

Module Contents

cellmap_flow.finetune.export_merged.logger

cellmap_flow.finetune.export_merged.POOL_FACTOR = 8

cellmap_flow.finetune.export_merged.TRAIN_TILE = 178

cellmap_flow.finetune.export_merged.CONTEXT = 122

cellmap_flow.finetune.export_merged.REL_TOL = 1e-05

cellmap_flow.finetune.export_merged.valid_tile(tile): True when tile keeps the pooling phase of the 178 training tile.

cellmap_flow.finetune.export_merged.load_eager_base(folder_path): The eager model.pt of a CellmapModel folder, eval mode, on CPU.

cellmap_flow.finetune.export_merged.apply_finetune(trainable, lora_adapter_path=None, weights_path=None)

Return a plain nn.Module with the finetune folded in.

trainable is the module tree the finetune was trained on, as model_loading.load_trainable_model builds it: the adapter’s target names and a full finetune’s state dict keys are that tree’s, and load onto no other. A TorchScript base’s tree is BatchLoopWrapper(export), whose child is named model, so its names carry a model. prefix; the wrapper only loops over the batch, and the folded module comes back without it. A base’s eager model.pt wrapped the same way reproduces those names exactly (checked: 38/38 adapter tensors identical, 0.0 output difference at 178^3), which is how main uses it.

The adapter is folded in by LoraStrategy.merge, the merge training uses to continue from a finetuned model: every adapted layer, Conv3d, Conv2d or Linear, to within float rounding of the adapter’s output.

cellmap_flow.finetune.export_merged.check_equivalence(module, reference, tile, seed=0)

Difference between module and reference at 178^3, and between the centre 56^3 of module’s output at tile and its own 178 output there.

Both are reported relative to the output’s own magnitude. These models emit unbounded logits whose scale varies by orders of magnitude between finetunes (observed: ~5 for one mito run, ~1400 for another), so an absolute tolerance is meaningless – folding the adapter into the base weights reassociates a sum over the LoRA rank channels, which perturbs the result by a float32 ULP at whatever scale the output happens to live.

cellmap_flow.finetune.export_merged.export_folder(module, output, input_shape, output_shape, metadata, provenance)

Write module into output as cellmap-models exports a model, the files type: cellmap serves from a folder and type: huggingface from the root of a repository (what cellmap’s own model repositories hold, but the eager model.pt, which only the training package can unpickle). Returns the names of the files written, of these:

model.ts: the TorchScript that is served, traced at an input of input_shape voxels, which it is locked to;
model.pt2: the torch.export program a LoRA finetune of the export unflattens and trains (CellmapModel.train), at the same input. A module torch.export cannot export (TorchScript itself, or one whose forward branches on its data) has none, and the export then trains in full only, as TorchScript does;
metadata.json: metadata with the shapes made the export’s, which it must be for the TorchScript to be fed the input it was traced at (models.configs.cellmap’s runs_at);
export_record.json: provenance, how the export was made.

cellmap_flow.finetune.export_merged.main(argv=None)