cellmap_flow.finetune.trainable

Building blocks for handing the finetune trainer a model’s network.

The trainer trains a torch module that takes the normalized patch (B, C, Z, Y, X) and gives the model’s output (B, C’, Z’, Y’, X’). Many networks are not quite that: a 2D network wants (N, C, Y, X), a model normalizes its own input, crops a halo off its output or ends in a sigmoid its serving applies. A model type’s trainable_model() wraps its network in these, composed with torch.nn.Sequential, rather than writing its own wrapper; they hold no parameters of their own, so what is trained is the network’s.

finetune_modes says what a module can be finetuned with: LoRA attaches adapters beside a network’s Conv and Linear layers, which needs a Python module tree, while a full finetune trains every parameter, which a compiled TorchScript module has too.

Classes

SliceWise	Run a 2D network on each z slice of (B, C, Z, Y, X).
Crop	Cut crop voxels off each side of the spatial axes (the last
ZeroMeanUnitVariance	(x - mean) / (std + eps), mean and std over dims of each sample
FixedZeroMeanUnitVariance	(x - mean) / (std + eps) with a given mean and std (one number, or one
ScaleRange	(x - lo) / (hi - lo + eps), lo and hi the min_percentile and
ScaleLinear	x * gain + offset (one number, or one per channel): bioimage.io's scale_linear.
Clip	x clamped to [min, max] (either may be None): bioimage.io's clip.

Functions

finetune_modes(→ Tuple[str, ...])

What module can be finetuned with: ("lora", "full"), ("full",) or ().

Module Contents

cellmap_flow.finetune.trainable.finetune_modes(module: torch.nn.Module) → Tuple[str, ...]

What module can be finetuned with: (“lora”, “full”), (“full”,) or ().

A TorchScript module trains in full but takes no adapters (its layers are compiled, so PEFT cannot wrap them); a module without parameters cannot be trained at all; any other with Conv or Linear layers takes both.

class cellmap_flow.finetune.trainable.SliceWise(net: torch.nn.Module)

Run a 2D network on each z slice of (B, C, Z, Y, X).

The slices go through net as one batch of B * Z images and come back as (B, C’, Z, Y’, X’): a 2D model served slice by slice is trained the way it is served.

net

forward(x)

class cellmap_flow.finetune.trainable.Crop(crop: Sequence[int])

Cut crop voxels off each side of the spatial axes (the last len(crop)): a model’s halo, which its serving cuts off too.

crop

forward(x)

class cellmap_flow.finetune.trainable.ZeroMeanUnitVariance(dims: Sequence[int] | None = None, eps: float = 1e-06)

(x - mean) / (std + eps), mean and std over dims of each sample (all but the batch by default): bioimage.io’s zero_mean_unit_variance.

forward(x)

class cellmap_flow.finetune.trainable.FixedZeroMeanUnitVariance(mean, std, eps: float = 1e-06, channel_dim: int = 1)

(x - mean) / (std + eps) with a given mean and std (one number, or one per channel): bioimage.io’s fixed_zero_mean_unit_variance.

forward(x)

class cellmap_flow.finetune.trainable.ScaleRange(min_percentile: float = 0.0, max_percentile: float = 100.0, dims: Sequence[int] | None = None, eps: float = 1e-06)

(x - lo) / (hi - lo + eps), lo and hi the min_percentile and max_percentile (0-100) over dims of each sample: bioimage.io’s scale_range, and Cellpose’s 1-99% normalization.

forward(x)

class cellmap_flow.finetune.trainable.ScaleLinear(gain=1.0, offset=0.0, channel_dim: int = 1)

x * gain + offset (one number, or one per channel): bioimage.io’s scale_linear.

channel_dim = 1

forward(x)

class cellmap_flow.finetune.trainable.Clip(min: float | None = None, max: float | None = None)

x clamped to [min, max] (either may be None): bioimage.io’s clip.

forward(x)