Training

Training fits the downstream model (an aggregator or a decoder, and a task head) to the features of a source. A config sets it in three sections, each holding the keywords of one callable: loaders those of soma.learn.loaders(), learner the name and params of the soma.learn.Learner, and trainer those of the soma.learn.Trainer. From Python, soma.fit() takes the same values.

Budget, optimizer and schedulers

loaders:
  batch_size: 1
learner:
  name: mil
  params:
    aggregator: {name: abmil}
    head: {name: binary_classification}
    optimizer:
      name: adamw
      params: {lr: 1.0e-4, weight_decay: 1.0e-5}
    scheduler:
      - {name: constant, params: {factor: 0.333, total_iters: 5}}
      - cosine
trainer:
  max_epochs: 50
  callbacks:
    - {name: checkpoint, params: {monitor: auroc, mode: max}}
    - {name: early_stopping, params: {monitor: auroc, mode: max, patience: 10}}

loaders holds batch_size, num_workers, pin_memory, persistent_workers and the optional sampler and augment leaves. learner.params holds the model’s aggregator or decoder leaf and its head leaf, the optimizer and scheduler leaves, the feature normalization and projection, and the keywords of a custom learner. learner.name is mil for the set shape and dense for the grid shape; a registered or dotted soma.learn.Learner subclass owns custom step logic (see A custom learner).

Choose one budget: trainer.max_epochs counts passes through the training loader; trainer.max_steps counts optimizer updates after accumulate_grad_batches, and replaces the epoch budget when set. A cosine scheduler anneals over the budget, per epoch or per update.

Every component is a {name, params} leaf whose name is a registry key or a dotted import path; a bare name keeps the component’s default params. The optimizer defaults to Adam (lr 1e-4, weight_decay 1e-5) when the learner names none. scheduler may be one leaf or a list, each stepped at its own interval (epoch unless set); a warm-up is a constant leaf. With no scheduler the learning rate stays constant.

A model with no trainable parameter, such as one with the attention_map detection decoder, is not fitted: the trainer returns it as built, with an empty training history. It still predicts on the trainer’s device.

Checkpoint selection and early stopping

The checkpoint callback selects the weights: selection: best (the default) keeps the epoch that is best on monitor in direction mode (default: minimize the tune loss; a metric name selects that metric), and selection: last keeps the final weights; it refuses an early_stopping callback and needs trainer.val_every_n_epochs: 1. early_stopping stops after patience tune evaluations without improvement; leave it out to train the full budget.

trainer.val_every_n_epochs: N evaluates the tune split after every N-th epoch and after the final epoch, for runs where tune evaluation dominates the epoch cost; selection and patience then count tune evaluations, not epochs. Tune metrics are always recorded for diagnostics.

While a fold trains, a summary panel shows the epoch, loss, learning rate, tune metrics, patience, trainable parameter count, timing and, for cross-validation, the active fold. Each fold writes the same history to training_history.json, one row per epoch: epoch, train_loss, lr, elapsed_seconds, then loss and one column per tune metric, null on an epoch without a tune evaluation.

The tune split

By default, training requires a tune split. The cohort section offers two explicit alternatives:

• cohort.tune_is_test: true uses one held-out split for both checkpoint selection and test reporting. Provide either tune or test, not both. Use this only when reproducing a benchmark protocol that selects on its reported cohort.

• cohort.allow_missing_tune: true reuses train as tune when tune is absent and emits a warning. This controls the diagnostic split independently of checkpoint selection.

Feature normalization

learner.params.normalization applies a feature adaptor before the trainable model, so encoders with different activation scales can share training settings.

learner:
  params:
    normalization:
      method: zscore   # none | zscore | l2 | layernorm
      eps: 1.0e-6

zscore estimates per-feature center and scale from the feature rows of the train split alone; held-out samples use the fitted transform without updating it. eps floors the scale of constant or near-constant channels and must be finite and positive. l2 and layernorm are stateless. none adds no adaptor and requires eps to keep its default.

The adaptor runs on cached tile-encoder MIL, slide-encoder embeddings and single-encoder dense grids; other paths refuse an active transform. A composite encoder normalizes each member’s block before concatenation instead (member_norm, see Composite encoders).

Feature projection

learner.params.projection maps frozen features to a common width without labels. With the same downstream input width, encoders can be compared without the aggregator’s capacity changing with their embedding widths.

learner:
  params:
    projection:
      method: pca      # none | pca | random
      target_dim: 512  # required when method != none
      seed: 0          # configurable only for random

Projection follows normalization; either can be enabled alone.

pca fits centered principal components on the train split alone, without whitening, and fixes each component’s sign by making its largest-magnitude entry positive. target_dim must not exceed the encoder dimension or the number of training feature rows; the adaptor refuses either before training starts.

random uses a fixed Gaussian matrix scaled by 1/sqrt(target_dim). It can reduce or expand the input while approximately preserving inner products and distances. Its seed is derived from the configured seed, the encoder identity and the dimensions, without touching the global random state. pca requires the default seed. The aggregator and head take target_dim as their input width when a projection is active.

Both stages are buffers of the model, so the fold’s bundle stores the fitted transform and evaluation reuses it. Neither changes the feature cache. The plan spells both out (method: none when off), so they are part of the experiment identity. A fold with an active transform also writes feature_adapter.json; see Run outputs.

Live re-encoding and augmentation

Segmentation and detection read cached feature grids by default (features.kind: cached): the frozen encoder runs once and every epoch reuses its output. With features.kind: live, soma re-encodes each pre-cropped tile through the frozen encoder at every step instead, so pixel augmentation reaches the features. This costs one encoder forward pass per tile and step.

features:
  kind: live
  encoder: {name: uni2}
loaders:
  augment:
    name: segmentation
    params:
      horizontal_flip: 0.5
      vertical_flip: 0.5
      rotation_degrees: 90
      crop_scale: [0.08, 1.0]

loaders.augment is null by default. It is a {name, params} leaf: the name is a registry key (segmentation is soma.learn.SegmentationAugment) or a dotted import path, and params go to its constructor. It applies to the train split only. Geometric operations transform the image and its label mask together (the mask with nearest-neighbour resampling); photometric operations change the image only. Pixels that a rotation or translation brings into the tile get ignore_index in the mask, so they carry no supervision. crop_scale takes a random crop of that area fraction and resizes it to the tile size before the other operations. For detection, the points move with the image: a point that leaves the tile is dropped, and the ignore mask moves like a label mask, so the pixels brought into the tile are ignored.

Live mode has these limits:

• loaders.augment requires features.kind: live; cached grids cannot be augmented.

• Live mode reads pre-cropped tiles only. It does not support annotation-sampled whole slides, composite encoders, an active normalization or projection, or a loaders.sampler.

Training from Python

soma.fit() trains a learner, or a bare model wrapped in its shape’s learner. It installs a Checkpoint and an EarlyStopping built from monitor, mode and patience (patience=None installs none), passes its other keywords to the soma.learn.Trainer, and returns a soma.FitResult: the learner in its selected state, the history frame and the selected epoch. soma.learn.loaders() builds one fold’s loaders from a source and a cohort: train shuffles, tune and each of tests keep split order.

fold = soma.learn.loaders(source, cohort, head=model.task_head, fold=0, batch_size=1)
result = soma.fit(model, fold.train, fold.tune, max_epochs=50, monitor="auroc", mode="max")
result.selected_epoch, result.history

Section 3 of the getting-started notebook, the same run from the verbs, runs a whole fold this way; Evaluation covers what follows fit.

The grid shape

The same verbs train a decoder over feature grids, with three differences:

• The source is a GridSource of cached grids, or a TileImageSource of pixels that the model re-encodes.

• soma.learn.loaders() needs targets=: a LabelMapSource for a segmentation head or a PointSource for a detection head. The head binds the grid geometry when it is built, and the loaders check that every sample matches it.

• The model is a soma.models.DenseModel, trained by soma.learn.DenseLearner. With an encoder stage it reads pixels; without one it reads cached grids.

soma.extract.live_encoder() wraps a slide2vec encoder as a parameter-free soma.extract.FrozenEncoder stage, so the foundation model never enters a checkpoint or an optimizer. preprocess=stage.preprocessor() runs the encoder’s preprocessing in the loader workers, after augmentation; augment is refused for cached grids:

from soma.models.decoders.registry import build_decoder_for_grid

stage = soma.extract.live_encoder("uni2", target_size=448, spacing_um=0.5, window=224)
images = soma.data.TileImageSource.from_manifest(soma.ImageManifest.from_csv("dataset.csv"), size=448, spacing_um=0.5)
labels = soma.data.LabelMapSource.from_manifests(soma.AnnotationManifest.from_csv("dataset.csv"), size=(448, 448))
head = soma.models.SegmentationHead(num_classes=3, geometry=stage.geometry)
decoder = build_decoder_for_grid(
    "lightweight_conv", {}, geometry=stage.geometry, input_dim=stage.feature_dim, num_classes=3
)
model = soma.models.DenseModel(decoder, head, encoder=stage)
fold = soma.learn.loaders(
    images, cohort, head=head, targets=labels, batch_size=8,
    preprocess=stage.preprocessor(),
    augment=soma.learn.SegmentationAugment(horizontal_flip=0.5, vertical_flip=0.5),
)
result = soma.fit(model, fold.train, fold.tune, max_epochs=50)

soma.testing.FakeEncoder is a weight-free stage with the same contract. The end-to-end scenarios grid_live_segmentation and grid_live_detection (tests/e2e/test_grid.py) run the whole grid flow on synthetic tiles with it, from loaders to evaluation.

Reference

class soma.config.LoadersConfig(batch_size=1, num_workers=0, pin_memory=False, persistent_workers=False, sampler=None, augment=None)

Bases: Editable

loaders: soma.learn.loaders()’ keywords. sampler is a training-batch sampler leaf (cached segmentation ROI strategies); augment augments a live grid run’s train split.

class soma.config.LearnerConfig(name='mil', params=<factory>)

Bases: Component

learner: the Learner and the model it trains.

params holds the model’s leaves (aggregator for the set shape, decoder for the grid shape, head with its targets: target key -> cohort column), the feature adaptor (normalization, projection), the optimizer leaf, the scheduler leaf or list of leaves, and any other keyword of the learner class.

class soma.config.TrainerConfig(max_epochs=50, max_steps=None, accumulate_grad_batches=1, val_every_n_epochs=1, device='auto', precision='fp32', callbacks=<factory>)

Bases: Editable

trainer: Trainer’s keywords. A step budget (max_steps) replaces the epoch budget; callbacks are the selection policy (checkpoint, early_stopping; leave the latter out to train the whole budget).

class soma.learn.SegmentationAugment(*, crop_scale=None, horizontal_flip=0.0, vertical_flip=0.0, rotation_degrees=0.0, translate=0.0, scale=0.0, brightness=0.0, contrast=0.0, saturation=0.0, hue=0.0, ignore_index=255, target_size=None)

Bases: object

Joint torchvision-v2 image + label-map augmentation.

Geometric ops transform the pixels and the label map together (the map resampled nearest-neighbour); photometric ops touch the pixels only. Out-of-canvas affine pixels fill the image with 0 and the label map with ignore_index, so synthesised borders carry no supervision. An optional RandomResizedCrop to target_size runs first, so images of any size reach the encoder at the supervision size.

Parameters:

• crop_scale (tuple[float, float] | None) – RandomResizedCrop area range (None = no crop); needs target_size.

• vertical_flip (horizontal_flip /) – Flip probabilities.

• rotation_degrees (float) – RandomAffine symmetric rotation range (±value).

• translate (float) – RandomAffine max translation fraction (both axes).

• scale (float) – RandomAffine scale jitter, scale ∈ (1 - value, 1 + value).

• hue (brightness / contrast / saturation /) – ColorJitter strengths.

• ignore_index (int) – Label-map fill for synthesised pixels.

• target_size (int | tuple[int, int] | None) – (H, W) output of crop_scale.

A Points target (in the pixel frame) rides as tv_tensors.KeyPoints: its points move with the pixels, points that leave the canvas are dropped (never clamped onto the edge), and its ignore raster moves like a label map, filled with 255 outside the canvas. Points without an ignore raster get one when the transform synthesises border pixels, so those carry no supervision.

class soma.learn.GridBatch(features, pixels, targets, sample_ids)

Bases: object

One collated grid batch.

features (B, D, h, w) are cached token grids; pixels (B, 3, H, W) are image pixels. Exactly one of them is the model’s input: features on the cached path, pixels on the live path, where the model’s encoder stage turns them into grids. An image-prior decoder reads pixels as a named side input, so a cached batch for such a decoder carries both. targets holds supervision only: {"mask": (B, H, W) long} or {"points": (B, K, 3) level-0 (x, y, class) NaN-padded, "ignore": (B, H, W) bool} (ignore absent when no sample has an ignore mask).

class soma.models.DenseModel(decoder, head, *, encoder=None, adaptor=None)

Bases: Module

The grid-shape model: pixels -> encoder -> adaptor -> decoder -> head.

Every stage but the decoder and the head is optional. Without an encoder the model reads cached feature grids (B, d, h, w); with one (a parameter-free FrozenEncoder, or any stage with the same contract) it re-encodes pixels every forward, so augmentation reaches the features. The adaptor (a fitted FeatureAdaptor: normalisation, projection) transforms the grid channel-axis before the decoder. The head resizes and crops the decoder’s logits to the supervision frame.

An image-prior decoder (consumes_image) also receives the batch’s raw uint8 pixels at the mask size, as [0, 1] floats. On the live path those pixels must reach the model unpreprocessed (the encoder stage then preprocesses them itself).

Parameters:

• decoder (Decoder) – Maps (B, d, h, w) to (B, C, h', w') logits, built against the adaptor’s output width.

• head (TaskHead) – A grid head (segmentation or detection) bound to the run’s geometry.

• encoder (Module | None) – Optional frozen encoder stage over pixels.

• adaptor (Module | None) – Optional channel-axis feature adaptor.

soma.extract.live_encoder(encoder, *, target_size, spacing_um, window=None, overlap=0.0, feature_kind='patch_features', precision=None, output_variant=None, allow_non_recommended_settings=False, tolerance=0.05, backend='auto', pad_mode='reflect', image_pad_value=None, attention_blocks=(-1,), attention_include_registers=False, num_gpus=None, device=None)

Prepare a slide2vec dense encoder for live re-encoding and wrap it as a stage.

Parameters:

• encoder (str) – slide2vec preset name (e.g. "uni2").

• target_size (int) – Supervision tile side in pixels; the stage pads it to the encoder’s patch multiple.

• spacing_um (float) – µm/px the pixels are read at.

• overlap (window /) – Sliding-window encoding of a tile larger than the encoder’s input (None = one forward over the whole tile).

• feature_kind (str) – slide2vec dense feature kind (patch_features, patch_features_prenorm, cls_attention).

• precision (str | None) – Encoder precision; None resolves the preset’s recommendation.

• image_pad_value (pad_mode /) – How the tile is padded to the encoder’s patch multiple, as DenseOptions states it for the cached path.

• attention_include_registers (attention_blocks /) – The cls_attention recipe.

Return type:

FrozenEncoder

The stage probe-encodes one blank tile to learn its feature width and to check that the kit’s grid matches its declared geometry.

class soma.data.TileImageSource(entries, *, size, spacing_um=None, spacing_policy='strict', backend='auto', tolerance=0.05)

Bases: object

Pixels of pre-cropped tiles (or whole small images), one (3, H, W) uint8 per id.

Parameters:

• entries (Mapping[str, ImageEntry] | Sequence[ImageEntry]) – The images, keyed by sample id (an ImageManifest’s entries).

• size (int | tuple[int, int]) – The declared (H, W) (or side) the encoder stage expects; images larger than it reach that size only through an augmentation crop.

• spacing_um (float | None) – µm/px to read pyramidal images at, and flat ones whose declared spacing_at_level_0 lies outside tolerance of it; None reads every image as stored.

• spacing_policy (str) – "strict" reads every image at spacing_um; "native_if_coarser" reads one whose declared spacing_at_level_0 is coarser than it at that native spacing (no upsampling).

• backend (str) – hs2p backend for spacing-aware reads.

• tolerance (float) – hs2p spacing tolerance: a level within it is read without resampling.