@attrs.define
class ModelTrainer:
"""Train sleap-nn model using PyTorch Lightning.
This class is used to create dataloaders, train a sleap-nn model and save the model checkpoints/ logs with options to logging
with wandb and csvlogger.
Args:
config: OmegaConf dictionary which has the following:
(i) data_config: data loading pre-processing configs.
(ii) model_config: backbone and head configs to be passed to `Model` class.
(iii) trainer_config: trainer configs like accelerator, optimiser params, etc.
train_labels: List of `sio.Labels` objects for training dataset.
val_labels: List of `sio.Labels` objects for validation dataset.
skeletons: List of `sio.Skeleton` objects in a single slp file.
lightning_model: One of the child classes of `sleap_nn.training.lightning_modules.LightningModel`.
model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
trainer: Instance of the `lightning.Trainer` initialized with loggers and callbacks.
"""
config: DictConfig
_initial_config: Optional[DictConfig] = None
train_labels: List[sio.Labels] = attrs.field(factory=list)
val_labels: List[sio.Labels] = attrs.field(factory=list)
skeletons: Optional[List[sio.Skeleton]] = None
lightning_model: Optional[LightningModel] = None
model_type: Optional[str] = None
backbone_type: Optional[str] = None
_profilers: dict = {
"advanced": AdvancedProfiler(),
"passthrough": PassThroughProfiler(),
"pytorch": PyTorchProfiler(),
"simple": SimpleProfiler(),
}
trainer: Optional[L.Trainer] = None
# Dry-run state (see `train(dry_run=True)`): the scratch folders to remove,
# as `(path, base, outermost folder created)`, the full data's labeled-frame
# counts (the labels are cut to one batch), and whether the real in-memory
# image cache would fit (`None` unless the pipeline caches in memory).
_dry_run_dirs: list = attrs.field(factory=list, init=False)
_dry_run_frame_counts: tuple = attrs.field(default=(0, 0), init=False)
_dry_run_cache_fits: Optional[bool] = attrs.field(default=None, init=False)
@classmethod
def get_model_trainer_from_config(
cls,
config: DictConfig,
train_labels: Optional[List[sio.Labels]] = None,
val_labels: Optional[List[sio.Labels]] = None,
):
"""Create a model trainer instance from config."""
model_trainer = cls(config=config)
model_trainer._initialize_from_config(
train_labels=train_labels, val_labels=val_labels
)
return model_trainer
def _initialize_from_config(
self,
train_labels: Optional[List[sio.Labels]] = None,
val_labels: Optional[List[sio.Labels]] = None,
):
"""Bring a bare trainer up to a trainable state.
The single initializer behind BOTH entry points --
:meth:`get_model_trainer_from_config` and the bare-constructor fallback in
:meth:`train`. It used to be duplicated between them, and the copy in
``train()`` had already drifted: it omitted ``_set_seed()`` (so
``ModelTrainer(config).train()`` ran with unseeded weight init and
augmentation even with ``trainer_config.seed`` set) and the
video-existence check (so a missing video surfaced later, as a confusing
failure deep in data loading).
Args:
train_labels: Labels to train on. ``None`` (with ``val_labels`` also
``None``) loads them from ``data_config.train_labels_path``.
val_labels: Labels to validate on, or ``None`` to split from
``train_labels`` / load from ``data_config.val_labels_path``.
Raises:
FileNotFoundError: If any video referenced by the labels is missing.
"""
# Normalize the config first: it fills in optional sections (e.g.
# `preprocessing.tiling`), and the bare constructor does no verification
# of its own, so without this the first access to a missing key raises.
self.config = verify_training_cfg(self.config)
# Derived from the config, but only when not explicitly provided to the
# constructor -- an explicit override must survive.
if self.model_type is None:
self.model_type = get_model_type_from_cfg(self.config)
if self.backbone_type is None:
self.backbone_type = get_backbone_type_from_cfg(self.config)
if self.config.trainer_config.seed is not None:
self._set_seed()
if train_labels is None and val_labels is None:
# read labels from paths provided in the config
train_labels = [
sio.load_slp(path) for path in self.config.data_config.train_labels_path
]
val_labels = (
[sio.load_slp(path) for path in self.config.data_config.val_labels_path]
if self.config.data_config.val_labels_path is not None
else None
)
self._setup_train_val_labels(labels=train_labels, val_labels=val_labels)
# Snapshot the pre-`setup_config` config: it is written out as
# `initial_config.yaml` at the end of training.
if self._initial_config is None:
self._initial_config = self.config.copy()
# update config parameters
self.setup_config()
# Check if all videos exist across all labels
all_videos_exist = all(
video.exists(check_all=True)
for labels in [*self.train_labels, *self.val_labels]
for video in labels.videos
)
if not all_videos_exist:
raise FileNotFoundError(
"One or more video files do not exist or are not accessible."
)
def _set_seed(self):
"""Set seed for the current experiment."""
seed = self.config.trainer_config.seed
random.seed(seed)
# torch
torch.manual_seed(seed)
# if cuda is available
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
# lightning
L.seed_everything(seed)
# numpy
np.random.seed(seed)
def _get_trainer_devices(self):
"""Get trainer devices."""
trainer_devices = (
self.config.trainer_config.trainer_devices
if self.config.trainer_config.trainer_devices is not None
else "auto"
)
if (
trainer_devices == "auto"
and OmegaConf.select(
self.config, "trainer_config.trainer_device_indices", default=None
)
is not None
):
trainer_devices = len(
OmegaConf.select(
self.config,
"trainer_config.trainer_device_indices",
default=None,
)
)
elif trainer_devices == "auto":
if torch.cuda.is_available():
trainer_devices = torch.cuda.device_count()
elif torch.backends.mps.is_available():
trainer_devices = 1
elif torch.xpu.is_available():
trainer_devices = torch.xpu.device_count()
else:
trainer_devices = 1
return trainer_devices
# Model types whose training targets are segmentation masks (``lf.masks``)
# rather than keypoint instances. Mask-only labels — no pose skeleton at all —
# are a normal input for these, so a frame with user masks and zero instances
# is a perfectly good training frame.
_MASK_TARGET_MODEL_TYPES = (
"bottomup_segmentation",
"centered_instance_segmentation",
"semantic_segmentation",
"embedding",
)
def _is_training_frame(self, lf: "sio.LabeledFrame") -> bool:
"""Whether a frame carries a usable training target.
Normally that means user instances. Two model families read something else:
- The centroid model can train on frames carrying only user centroid
annotations (``frame.centroids``) with no pose instance — the
pure-centroid seeding case (label a body-center per animal before any
keypoints exist).
- Mask-target models (see ``_MASK_TARGET_MODEL_TYPES``) train on
``lf.masks``. Mask-only datasets carry **zero** instances, so counting
only user instances rejected them outright: every split came back empty
and ``_validate_nonempty_labels`` raised "No labeled frames available
for train" before setup finished. That made the `embedding` model
untrainable on its primary dataset, and would do the same to
`semantic_segmentation` on mask-only labels.
"""
if self.model_type == "embedding":
# The embedding model trains on any detection its dataset keeps: user
# labels, and predictions no user label on the frame supersedes (proofread
# tracks on predicted poses stay `PredictedInstance`s).
instances, masks = embedding_detections(
lf,
OmegaConf.select(
self.config, "data_config.user_instances_only", default=True
),
)
return bool(instances or masks)
if self.model_type in BBOX_MODEL_TYPES:
# Boxes were resolved/derived by `_apply_bbox_source` before counting.
return bool(frame_user_bboxes(lf))
if lf.has_user_instances:
return True
if self.model_type == "centroid":
return any(not c.is_predicted for c in lf.centroids)
if self.model_type in self._MASK_TARGET_MODEL_TYPES:
return any(not m.is_predicted for m in (getattr(lf, "masks", None) or []))
return False
def _count_labeled_frames(
self, labels_list: List[sio.Labels], user_only: bool = True
) -> int:
"""Count labeled frames, optionally filtering to trainable frames only.
Args:
labels_list: List of Labels objects to count frames from.
user_only: If True, count only frames with a training target
(user instances, or — for the centroid model — user centroids).
Returns:
Total count of labeled frames.
"""
total = 0
for label in labels_list:
if user_only:
total += sum(1 for lf in label if self._is_training_frame(lf))
else:
total += len(label)
return total
def _filter_to_user_labeled(self, labels: sio.Labels) -> sio.Labels:
"""Filter a Labels object to only include trainable frames.
Args:
labels: Labels object to filter.
Returns:
New Labels object containing only frames with a training target
(user instances, or user centroids for the centroid model).
"""
# Filter labeled frames to only trainable ones
user_lfs = [lf for lf in labels if self._is_training_frame(lf)]
# Set instances to user instances only (empty for centroid-only frames).
# Not for `embedding`: it trains on unsuperseded predictions too, its dataset
# and post-training eval apply `user_instances_only` themselves, and these are
# the TRAINING labels' own frames (this runs on rank 0 before the datasets are
# built), so stripping them here would drop training samples.
if self.model_type != "embedding":
for lf in user_lfs:
lf.instances = lf.user_instances
# Create new Labels with filtered frames
return sio.Labels(
labeled_frames=user_lfs,
videos=labels.videos,
skeletons=labels.skeletons,
tracks=labels.tracks,
suggestions=labels.suggestions,
provenance=labels.provenance,
)
def _split_embedding_labels(
self, label: sio.Labels, val_fraction: float, seed: Optional[int]
):
"""Random frame split for the `embedding` model that keeps its samples.
`sio.Labels.make_training_splits` step for step, except that where it removes
every prediction this removes only the predictions a user label on the same
frame supersedes (with ``user_instances_only``; nothing without it), the rule
the embedding dataset applies. On labels without such predictions the split
is the one `make_training_splits` makes.
"""
user_instances_only = OmegaConf.select(
self.config, "data_config.user_instances_only", default=True
)
labels = deepcopy(label)
for lf in labels.labeled_frames:
lf.instances, lf.masks = embedding_detections(lf, user_instances_only)
labels.suggestions = []
labels.clean()
train_split, rest = labels.split(1 - val_fraction, seed=seed)
n_val = val_fraction
if n_val < 1:
n_val = (n_val * len(labels)) / len(rest)
if isinstance(n_val, float) and n_val == 1.0:
val_split = rest
else:
val_split, _ = rest.split(n=n_val, seed=seed)
source_labels = label.provenance.get("filename", None)
train_split.provenance["source_labels"] = source_labels
val_split.provenance["source_labels"] = source_labels
return train_split, val_split
def _split_centroid_labels(
self, label: sio.Labels, val_fraction: float, seed: Optional[int]
):
"""Train/val split for the centroid model: the other models' draw, filtered.
The frames are drawn by `sio.Labels.make_training_splits`, the split every
other pose model type uses, so a centroid model and a centered-instance
model trained on the same labels, `validation_fraction` and seed hold out
the same frames and a top-down pipeline can be scored end to end on frames
neither model trained on (#833). A separate draw, even with the same seed,
overlaps the other models' validation set only by chance.
- It keeps centroid-only frames (user centroids, no pose instance): it
removes predictions (predicted centroids included), then only frames
left with no annotation at all.
- Training frames are the other models' exactly. Validation is the whole
remainder (no ``n_val``): the other models' validation frames plus the at
most one frame their rounding leaves in neither split, which none of
them trains on either. So no trainable frame is lost and every
validation frame is held out from both models.
- Each side is then filtered to the frames the centroid model can train on
(`_is_training_frame`) plus negative frames, which it reads with
``use_negative_frames`` and skips otherwise; this drops e.g. box-only
frames and moves no frame to the other side. Only when one side is left
without a trainable frame (tiny projects) is one lent across from the
other, with a warning.
"""
if not any(self._is_training_frame(lf) for lf in label):
# `make_training_splits` divides by zero when nothing survives its
# cleanup; return empty splits so `_validate_nonempty_labels` raises
# its clear error instead (or another labels file supplies frames).
empty = sio.Labels(videos=label.videos, skeletons=label.skeletons)
return empty, deepcopy(empty)
splits = label.make_training_splits(n_train=1 - val_fraction, seed=seed)
frames = {
name: [
lf
for lf in splits[name]
if self._is_training_frame(lf) or lf.is_negative
]
for name in ("train", "val")
}
empty_side = [
name
for name, lfs in frames.items()
if not any(self._is_training_frame(lf) for lf in lfs)
]
if empty_side:
# The draw put no frame the centroid model can train on on this side.
# Lend it one from the other side (shared when that is the only one,
# as `Labels.split` does for a single frame) so training still runs.
(short,) = empty_side
donor = "val" if short == "train" else "train"
donor_lfs = [lf for lf in frames[donor] if self._is_training_frame(lf)]
lent = donor_lfs[-1]
if len(donor_lfs) > 1:
frames[donor] = [lf for lf in frames[donor] if lf is not lent]
# Each split holds its own copies; point the lent copy at this side's.
lent = deepcopy(lent)
lent.video = next(
(v for v in splits[short].videos if v.matches_path(lent.video)),
lent.video,
)
frames[short].append(lent)
side = "training" if short == "train" else "validation"
logger.warning(
f"Centroid train/val split: no {side} frame drawn has a centroid "
"target (e.g. negative or box-only frames), so one frame was moved "
f"to {side} from the other side. This split no longer matches the "
"other model types' split exactly."
)
def mk(name):
split = splits[name]
return sio.Labels(
labeled_frames=frames[name],
videos=split.videos,
skeletons=split.skeletons,
tracks=split.tracks,
suggestions=split.suggestions,
provenance=split.provenance,
)
return mk("train"), mk("val")
def _setup_train_val_labels(
self,
labels: Optional[List[sio.Labels]] = None,
val_labels: Optional[List[sio.Labels]] = None,
):
"""Create train and val labels objects. (Initialize `self.train_labels` and `self.val_labels`)."""
logger.info(f"Creating train-val split...")
total_train_lfs = 0
total_val_lfs = 0
self.skeletons = labels[0].skeletons
# Mask-only labels: derive the centroid annotations a centroid model needs
# from the segmentation masks. Must run BEFORE splitting and counting —
# such labels have no trainable frame at all until it has (#586 / #674) —
# and on the raw lists, so every split inherits the derived centroids.
self._apply_centroids_from_masks(labels, val_labels)
# Bounding-box models: give every frame its GT boxes (`data_config.
# bbox_source`), deriving them from poses/masks when the labels carry no
# boxes. Same ordering constraint: before splitting and counting.
self._apply_bbox_source(labels, val_labels)
# Check if we should count only user-labeled frames
user_instances_only = OmegaConf.select(
self.config, "data_config.user_instances_only", default=True
)
# check if all `.slp` files have the same skeleton structure (if multiple slp
# file paths are provided). Node ORDER must match too: points are read in
# each file's own node order, so the same names in a different order would
# silently train each channel on a different node.
# Mask-only labels (e.g. instance segmentation) may carry no skeleton.
if self.skeletons:
skeleton = self.skeletons[0]
for split, label_list in (("training", labels), ("validation", val_labels)):
for index, label in enumerate(label_list or []):
if not label.skeletons:
continue
skel_temp = label.skeletons[0]
if skeleton.matches(skel_temp, require_same_order=True):
continue
message = (
f"The skeleton in the {split} labels: {index + 1} does not "
"match the skeleton in the first training label file."
)
if skeleton.matches(skel_temp):
message += (
" The nodes are the same but in a different order "
f"({skel_temp.node_names} vs {skeleton.node_names}); "
"reorder them to match, e.g. with "
"`labels.reorder_nodes(...)` in sleap-io."
)
logger.error(message)
raise ValueError(message)
# Check for same-data mode (train = val, for intentional overfitting)
use_same = OmegaConf.select(
self.config, "data_config.use_same_data_for_val", default=False
)
# Group-aware split (SPEC §5.3): when `data_config.split` is configured and no
# explicit val labels are provided, partition the training labels by a group key
# (frame/video/identity) instead of the frame-level random validation split.
split_cfg = OmegaConf.select(self.config, "data_config.split", default=None)
use_group_split = (
split_cfg is not None
and not use_same
and (val_labels is None or not len(val_labels))
)
if use_same:
# Same mode: use identical data for train and val (for overfitting)
logger.info("Using same data for train and val (overfit mode)")
self.train_labels = labels
self.val_labels = labels
total_train_lfs = self._count_labeled_frames(labels, user_instances_only)
total_val_lfs = total_train_lfs
elif use_group_split:
from sleap_nn.data.splitting import split_labels_list_train_val
logger.info(
"Using group-aware train/val split "
f"(split_by={OmegaConf.select(self.config, 'data_config.split.split_by', default='frame')})"
)
self.train_labels, self.val_labels = split_labels_list_train_val(
labels, split_cfg
)
total_train_lfs = self._count_labeled_frames(
self.train_labels, user_instances_only
)
total_val_lfs = self._count_labeled_frames(
self.val_labels, user_instances_only
)
elif val_labels is None or not len(val_labels):
# if val labels are not provided, split from train
val_fraction = OmegaConf.select(
self.config, "data_config.validation_fraction", default=0.1
)
seed = self.config.trainer_config.seed
if self.model_type == "embedding":
logger.warning(
"embedding: `validation_fraction` splits FRAMES at random, so "
"near-identical neighboring frames of one animal land on both "
"sides and inflate the validation rank-1 that selects the "
"checkpoint. Set `data_config.split` (split_by='video' measures "
"cross-session re-ID) or pass `val_labels_path`."
)
# Warn if resuming from a checkpoint with a potentially different seed
resume_ckpt = OmegaConf.select(
self.config, "trainer_config.resume_ckpt_path", default=None
)
if resume_ckpt is not None:
orig_config_path = Path(resume_ckpt).parent / "training_config.yaml"
if orig_config_path.exists():
try:
orig_cfg = OmegaConf.load(orig_config_path.as_posix())
orig_seed = OmegaConf.select(
orig_cfg, "trainer_config.seed", default=None
)
if orig_seed != seed:
logger.warning(
f"Current seed ({seed}) differs from the original "
f"training seed ({orig_seed}) in {orig_config_path}. "
f"This will produce a different train/val split and "
f"may cause data leakage between train and val sets. "
f"Set `trainer_config.seed: {orig_seed}` to preserve "
f"the original split."
)
except Exception:
pass
else:
logger.warning(
f"Resuming from checkpoint but could not find "
f"{orig_config_path} to verify the train/val split seed. "
f"Ensure `trainer_config.seed` matches the original "
f"training run to avoid data leakage."
)
for label in labels:
if self.model_type == "centroid":
# The same frame draw as the other models, filtered to frames
# with a centroid target (pure-centroid ones too) + negatives.
train_split, val_split = self._split_centroid_labels(
label, val_fraction, seed
)
elif self.model_type == "embedding":
# make_training_splits removes EVERY prediction, including the
# unsuperseded ones the embedding model trains on.
train_split, val_split = self._split_embedding_labels(
label, val_fraction, seed
)
else:
train_split, val_split = label.make_training_splits(
n_train=1 - val_fraction, n_val=val_fraction, seed=seed
)
self.train_labels.append(train_split)
self.val_labels.append(val_split)
total_train_lfs += len(train_split)
total_val_lfs += len(val_split)
else:
self.train_labels = labels
self.val_labels = val_labels
total_train_lfs = self._count_labeled_frames(labels, user_instances_only)
total_val_lfs = self._count_labeled_frames(val_labels, user_instances_only)
logger.info(f"# Train Labeled frames: {total_train_lfs}")
logger.info(f"# Val Labeled frames: {total_val_lfs}")
# Fail fast on an empty split instead of letting training run for several
# minutes of setup only to crash with a cryptic IndexError the first time
# something indexes into an empty Labels (e.g. `_verify_model_input_channels`
# accessing `self.train_labels[0][0]`).
self._validate_nonempty_labels(total_train_lfs, "train")
self._validate_nonempty_labels(total_val_lfs, "validation")
# Single-instance models assume exactly one instance per frame; fail fast
# with a clear error if any frame has more than one.
if self.model_type == "single_instance":
self._validate_single_instance_labels(self.train_labels, "train")
self._validate_single_instance_labels(self.val_labels, "validation")
def _apply_centroids_from_masks(self, *label_lists):
"""Derive user centroids from segmentation masks when configured.
No-op unless ``data_config.centroids_from_masks`` names a method. Mutates
the labels in place, before any split or frame count: a mask-only dataset
(no poses, no centroid annotations) has zero trainable frames until this
has run, so running it later fails the empty-split guard.
Args:
*label_lists: The raw ``List[sio.Labels]`` inputs (train, val).
``None`` entries are skipped.
"""
method = OmegaConf.select(
self.config, "data_config.centroids_from_masks", default=None
)
if not method:
return
from sleap_nn.data.instance_centroids import add_centroids_from_masks
for label_list in label_lists:
for labels in label_list or []:
add_centroids_from_masks(labels, method=method)
def _apply_bbox_source(self, *label_lists):
"""Resolve each labels file's ground-truth boxes for a ``bbox_*`` model.
No-op for every other model type. Mutates the labels in place (in memory)
via :func:`sleap_nn.data.bboxes.add_bboxes_to_labels`; see
``data_config.bbox_source`` / ``data_config.bbox_padding``.
Args:
*label_lists: The raw ``List[sio.Labels]`` inputs (train, val).
``None`` entries are skipped.
"""
if self.model_type not in BBOX_MODEL_TYPES:
return
source = OmegaConf.select(
self.config, "data_config.bbox_source", default="auto"
)
padding = OmegaConf.select(self.config, "data_config.bbox_padding", default=0.0)
for label_list in label_lists:
for labels in label_list or []:
add_bboxes_to_labels(labels, source=source, padding=padding)
# Box metrics a `bbox_*` model can select checkpoints on (all higher-is-better).
_BBOX_SELECT_METRICS = ("bbox_mAP", "bbox_AP50", "bbox_AP75", "bbox_AR", "bbox_f1")
def _bbox_select_metric(self) -> Optional[str]:
"""The eval metric a ``bbox_*`` model selects checkpoints on, or ``None``.
``trainer_config.eval.select_metric`` names one of ``_BBOX_SELECT_METRICS``,
or ``val_loss`` for the loss. Left at its schema default (``"rank1"``, an
embedding metric), detectors select on ``bbox_mAP`` when
``trainer_config.eval.enabled`` (the metric only exists then) and on the
val loss otherwise.
Raises:
ValueError: For an unknown metric name, or a box metric requested with
evaluation disabled (the metric would never be logged).
"""
select = OmegaConf.select(
self.config, "trainer_config.eval.select_metric", default="rank1"
)
eval_enabled = bool(
OmegaConf.select(self.config, "trainer_config.eval.enabled", default=False)
)
if select in ("rank1", None):
return "bbox_mAP" if eval_enabled else None
if select == "val_loss":
return None
if select not in self._BBOX_SELECT_METRICS:
raise ValueError(
f"trainer_config.eval.select_metric='{select}' is not a valid "
"selection metric for a bounding-box model; choose one of "
f"{'|'.join(self._BBOX_SELECT_METRICS + ('val_loss',))}."
)
if not eval_enabled:
raise ValueError(
f"trainer_config.eval.select_metric='{select}' needs "
"trainer_config.eval.enabled=true (box metrics are only computed "
"by the epoch-end evaluation)."
)
return select
def _setup_bbox_head_config(self):
"""Fill in a ``bbox_*`` model's categories (and SSD's anchor sizes) from labels.
``categories`` defaults to the training boxes' categories in first-seen order
(``["object"]`` when none has one). SSD ``anchor_sizes`` default to sizes
spanning the training boxes in model-input pixels
(:func:`sleap_nn.data.bboxes.auto_anchor_sizes`). Both are written into the
config, so the saved model carries them to inference.
"""
head_configs = self.config.model_config.head_configs
leaf = bbox_head_leaf(self.model_type, head_configs)
if leaf.categories is None:
leaf.categories = labels_bbox_categories(
self.train_labels + self.val_labels
)
logger.info(f"Bounding-box categories: {list(leaf.categories)}")
if self.model_type == "bbox_ssd" and leaf.anchor_sizes is None:
preprocessing = self.config.data_config.preprocessing
sizes = bbox_size_stats(
self.train_labels,
max_hw=(preprocessing.max_height, preprocessing.max_width),
scale=preprocessing.scale,
)
leaf.anchor_sizes = auto_anchor_sizes(
sizes, int(leaf.num_levels), ssd_level_strides(leaf)
)
logger.info(
"SSD anchor sizes (model-input px, per level "
f"{ssd_level_strides(leaf)}): "
f"{[round(a, 1) for a in leaf.anchor_sizes]}"
)
def _validate_nonempty_labels(self, n_labeled_frames: int, split_name: str):
"""Ensure a split has at least one trainable labeled frame.
Args:
n_labeled_frames: Count of labeled frames usable as training targets
for this split (user instances, or user centroids for the
centroid model).
split_name: Name of the split (e.g. "train", "validation") for the
error message.
Raises:
ValueError: If `n_labeled_frames` is zero.
"""
if n_labeled_frames == 0:
message = (
f"No labeled frames available for {split_name}: none of the "
f"labeled frame(s) in the provided {split_name} labels contain "
"user-labeled data usable by this model. Predicted instances "
"and suggestion frames are not used as training targets (nor "
"are standalone centroid annotations, except by centroid "
"models). Verify that the .slp file(s) passed for training "
"contain user-labeled instances, and that `validation_fraction` "
"and `use_same_data_for_val` are set as intended."
)
logger.error(message)
raise ValueError(message)
def _validate_single_instance_labels(
self, labels: List[sio.Labels], split_name: str
):
"""Ensure no frame has more than one instance for single-instance models.
Single-instance confidence-map generation flattens the instance dimension
(see `sleap_nn.data.confidence_maps.generate_confmaps`), so a frame with
more than one instance would silently merge multiple animals into a single
instance with `n_instances * n_nodes` "nodes" and corrupt training. This
raises a clear error instead, naming the offending frame.
Args:
labels: List of `sio.Labels` objects to validate.
split_name: Name of the split (e.g. "train", "validation") for the
error message.
Raises:
ValueError: If any frame contains more than one (non-empty) instance.
"""
user_instances_only = OmegaConf.select(
self.config, "data_config.user_instances_only", default=True
)
for label in labels:
for lf in label:
if user_instances_only and lf.user_instances is not None:
instances = lf.user_instances
else:
instances = lf.instances
instances = [inst for inst in instances if not inst.is_empty]
if len(instances) > 1:
video_idx = label.videos.index(lf.video)
raise ValueError(
f"Single-instance training requires at most one instance "
f"per frame, but the {split_name} frame at (video index "
f"{video_idx}, frame_idx {lf.frame_idx}) has "
f"{len(instances)} instances. Remove the extra instance(s) "
f"from this frame, or train a multi-instance (top-down or "
f"bottom-up) model instead."
)
def _resolve_crop_centroid(self):
"""Resolve the centroid definition the crops will be centered on.
The crop center is what makes a crop size sufficient or not, so sizing
the crop requires knowing it. Resolution is deliberately lenient: this
runs before `_setup_head_config`, which is where a bad ``anchor_part``
is reported, so an unresolvable anchor degrades to its fallback here and
lets that validation raise the good error a moment later.
Returns:
``(anchor_ind, anchor_part, method, fallback)`` -- the first, third
and fourth for `generate_centroids`, and the node's name for
reporting. ``anchor_part`` is ``None`` when no anchor is in force,
including when a configured one could not be resolved and was
degraded to its fallback.
"""
leaf_paths = {
"centered_instance": "centered_instance.confmaps",
"multi_class_topdown": "multi_class_topdown.confmaps",
"centered_instance_segmentation": (
"centered_instance_segmentation.segmentation"
),
"embedding": "embedding.embedding",
}
leaf_path = leaf_paths.get(self.model_type)
if leaf_path is None:
return None, None, None, None
head_cfg = OmegaConf.select(
self.config, f"model_config.head_configs.{leaf_path}", default=None
)
anchor_part = OmegaConf.select(head_cfg, "anchor_part", default=None)
anchor_ind = None
if anchor_part is not None:
# Resolve against the skeleton rather than the head's `part_names`,
# which `_setup_head_config` has not populated yet.
for labels in self.train_labels:
if labels.skeletons:
names = labels.skeletons[0].node_names
if anchor_part in names:
anchor_ind = names.index(anchor_part)
break
method, fallback = degrade_anchor_if_unresolved(
*centroid_method_from_config(head_cfg), anchor_ind
)
# Report the name only if the anchor actually took effect; an
# unresolvable one has degraded to `fallback` and naming it would
# describe a center that is not being used.
return (
anchor_ind,
(anchor_part if anchor_ind is not None else None),
method,
fallback,
)
def _compute_crop_padding(self, train_label, max_hw, bbox_size=None):
"""Return the augmentation margin to add to a computed crop size.
Args:
train_label: The `sio.Labels` to measure instances from.
max_hw: The resolved ``(max_height, max_width)``, so the bounding
box is measured in the space the crop is taken in.
bbox_size: The largest bounding-box side, when already measured (the
mask-mode embedding crop measures masks, not instances).
Returns:
Padding in pixels, from the config when set, else derived from the
enabled geometric augmentations.
"""
padding = self.config.data_config.preprocessing.crop_padding
if padding is not None:
return padding
aug_config = self.config.data_config.augmentation_config
if not (
self.config.data_config.use_augmentations_train
and aug_config is not None
and aug_config.geometric is not None
):
return 0
geo = aug_config.geometric
# Check if rotation is enabled (via rotation_p or affine_p)
rotation_enabled = (geo.rotation_p is not None and geo.rotation_p > 0) or (
geo.rotation_p is None
and geo.scale_p is None
and geo.translate_p is None
and geo.affine_p > 0
)
# Check if scale is enabled (via scale_p or affine_p)
scale_enabled = (geo.scale_p is not None and geo.scale_p > 0) or (
geo.rotation_p is None
and geo.scale_p is None
and geo.translate_p is None
and geo.affine_p > 0
)
if not (rotation_enabled or scale_enabled):
return 0
# First find the actual max bbox size from labels
if bbox_size is None:
bbox_size = find_max_instance_bbox_size(
train_label,
max_hw=max_hw,
user_instances_only=self.config.data_config.user_instances_only,
superseded_only=self.model_type == "embedding",
)
bbox_size = max(
bbox_size,
self.config.data_config.preprocessing.min_crop_size or 100,
)
rotation_max = (
max(abs(geo.rotation_min), abs(geo.rotation_max))
if rotation_enabled
else 0.0
)
scale_max = geo.scale_max if scale_enabled else 1.0
return compute_augmentation_padding(
bbox_size=bbox_size,
rotation_max=rotation_max,
scale_max=scale_max,
)
def _setup_preprocessing_config(self):
"""Setup preprocessing config.
Runs in two passes: ``max_height``/``max_width`` must be resolved across
*every* labels file before any crop size can be computed, because the
size matcher rescales each frame to them and the crop is taken in that
rescaled space. Measuring a crop in native pixels under-sizes it for any
video that gets scaled up (see #2862).
"""
max_height = self.config.data_config.preprocessing.max_height
max_width = self.config.data_config.preprocessing.max_width
# Pass 1: resolve the size-matcher target.
if max_height is None or max_width is None:
max_h, max_w = 0, 0
for train_label in self.train_labels:
current_max_h, current_max_w = get_max_height_width(train_label)
max_h = max(max_h, current_max_h)
max_w = max(max_w, current_max_w)
self.config.data_config.preprocessing.max_height = max_h
self.config.data_config.preprocessing.max_width = max_w
if self.model_type not in CROPPING_MODEL_TYPES:
return
max_hw = (
self.config.data_config.preprocessing.max_height,
self.config.data_config.preprocessing.max_width,
)
if (
self.model_type == "embedding"
and self.config.model_config.head_configs.embedding.embedding.detection_mode
== "mask"
):
# Mask-mode crops are centered on the mask and must hold the MASK: the
# keypoint sizing below finds no keypoints on mask-only data and fell
# back to `min_crop_size` (a 400x300 px animal got a 128 px crop).
self._setup_mask_crop_size(max_hw)
return
bbox_crop = BBoxCropSpec.from_config(self.config)
if bbox_crop is not None:
self._setup_bbox_crop(bbox_crop, max_hw)
return
(
anchor_ind,
anchor_part,
centroid_method,
centroid_fallback,
) = self._resolve_crop_centroid()
user_instances_only = self.config.data_config.user_instances_only
crop_size = self.config.data_config.preprocessing.crop_size
# Pass 2: size the crop when it was not given.
was_auto = crop_size is None
if was_auto:
max_crop_size = 0
for train_label in self.train_labels:
padding = self._compute_crop_padding(train_label, max_hw)
crop_sz = find_instance_crop_size(
labels=train_label,
padding=padding,
maximum_stride=self.config.model_config.backbone_config[
f"{self.backbone_type}"
]["max_stride"],
min_crop_size=self.config.data_config.preprocessing.min_crop_size,
max_hw=max_hw,
anchor_ind=anchor_ind,
centroid_method=centroid_method,
centroid_fallback=centroid_fallback,
user_instances_only=user_instances_only,
superseded_only=self.model_type == "embedding",
)
max_crop_size = max(max_crop_size, crop_sz)
self.config.data_config.preprocessing.crop_size = max_crop_size
crop_size = max_crop_size
self._log_crop_size(
max_crop_size, anchor_ind, anchor_part, centroid_method, max_hw
)
# Check the resolved crop size against every labeled instance, however it
# was arrived at. A computed size is derived from the TRAIN split, so a
# larger val instance can still clip -- and the validation images are
# exactly where a user would notice it (#2862).
self._warn_if_crop_size_clips(
crop_size,
max_hw,
anchor_ind,
centroid_method,
centroid_fallback,
was_auto=was_auto,
)
def _setup_embedding_carrier(self):
"""Decide and record the carrier (pose / mask) an embedding model trains on.
Resolved from the TRAIN labels with the dataset's own membership rule
(:func:`~sleap_nn.data.custom_datasets.resolve_embedding_detection_mode`):
a configured ``detection_mode``, else masks under ``burn_in``, else the
carrier holding more samples. Writing it into
``head_configs.embedding.embedding.detection_mode`` makes the train and val
datasets crop the same carrier, sizes the crop for it, and tells inference
which carrier to embed when a file holds both.
"""
leaf = self.config.model_config.head_configs.embedding.embedding
membership = embedding_membership_from_config(
self.config, self.train_labels + self.val_labels
)
# (Warns when the preferred carrier holds no sample and the other is used.)
mode = resolve_embedding_detection_mode(
self.train_labels,
membership.is_member,
preferred=embedding_preferred_carrier(self.config),
)
leaf.detection_mode = mode
logger.info(f"Embedding training carrier (detection_mode): {mode}")
def _setup_bbox_crop(self, spec: BBoxCropSpec, max_hw):
"""Size box-scaled crops (``crop_mode="bbox"``) from the training boxes.
Every box is resized to the crop, so ``crop_size`` no longer has to hold the
largest animal: it sets the resolution an animal is seen at. Left unset, it
is the median training window side rounded up to the backbone's max stride
(a typical animal keeps its pixel size), and ``crop_bbox_min_size`` is a
quarter of that median, bounding how far a tiny or partial box is
magnified. Both are written into the config, so inference cuts the same
crops.
Raises:
ValueError: For a model type whose crops cannot be box-scaled, or when
no training instance has a visible keypoint.
"""
if self.model_type not in BBOX_CROP_MODEL_TYPES:
message = (
f"preprocessing.crop_mode='bbox' is supported for "
f"{', '.join(BBOX_CROP_MODEL_TYPES)}, not {self.model_type!r}."
)
logger.error(message)
raise ValueError(message)
preprocessing = self.config.data_config.preprocessing
# Freeze `bbox_source: auto` to what it means for the training labels, so a
# prediction file that happens to carry user boxes cannot re-resolve it.
from sleap_nn.data.bboxes import resolve_bbox_source
if spec.bbox_source == "auto":
resolved = {resolve_bbox_source(lbl, "auto") for lbl in self.train_labels}
if len(resolved) == 1:
self.config.data_config.bbox_source = resolved.pop()
spec = BBoxCropSpec.from_config(self.config)
sides = np.concatenate(
[
bbox_crop_window_sides(
labels,
spec,
max_hw,
user_instances_only=self.config.data_config.user_instances_only,
)
for labels in self.train_labels
]
)
if sides.size == 0:
message = "crop_mode='bbox': no training instance has a visible keypoint."
logger.error(message)
raise ValueError(message)
p5, median, p95 = np.percentile(sides, [5, 50, 95])
if preprocessing.crop_bbox_min_size is None:
preprocessing.crop_bbox_min_size = round(float(median) / 4.0, 2)
stride = int(
self.config.model_config.backbone_config[f"{self.backbone_type}"][
"max_stride"
]
)
if preprocessing.crop_size is None:
preprocessing.crop_size = max(
stride, int(math.ceil(float(median) / stride) * stride)
)
crop_size = int(preprocessing.crop_size)
min_size = float(preprocessing.crop_bbox_min_size)
logger.info(
f"Box-scaled crops ({spec.aspect}, padding {spec.padding:g} of the longest "
f"side): crop_size {crop_size}px; training windows p5/p50/p95 "
f"{p5:.0f}/{median:.0f}/{p95:.0f}px, i.e. scales "
f"{crop_size / max(p95, min_size):.2f}-{crop_size / max(p5, min_size):.2f}"
f"x; windows floored at {min_size:g}px (at most "
f"{crop_size / min_size:.1f}x)."
)
def _setup_mask_crop_size(self, max_hw):
"""Size (when unset) and check the crop of a mask-mode embedding model.
The mask analog of the keypoint path in `_setup_preprocessing_config`: the
crop is centered on each mask (``crop_centering``) and must hold the whole
mask. Computed from the TRAIN split; both splits are checked for clipping.
"""
preprocessing = self.config.data_config.preprocessing
crop_centering = (
OmegaConf.select(
self.config, "data_config.preprocessing.crop_centering", default=None
)
or "auto"
)
stride = int(
self.config.model_config.backbone_config[f"{self.backbone_type}"][
"max_stride"
]
)
def extents(labels_list):
return [
extent
for labels in labels_list
for extent in iter_mask_extents(
labels,
max_hw=max_hw,
crop_centering=crop_centering,
user_instances_only=self.config.data_config.user_instances_only,
min_mask_area=EMBEDDING_MIN_MASK_AREA,
)
]
train = extents(self.train_labels)
val = extents(self.val_labels)
crop_size = preprocessing.crop_size
was_auto = crop_size is None
if was_auto:
padding = self._compute_crop_padding(
None, max_hw, bbox_size=max((bbox for _, bbox in train), default=0.0)
)
min_crop_size = preprocessing.min_crop_size or 0
length = max([float(min_crop_size - padding)] + [req for req, _ in train])
length = max(length, 0.0) + float(padding)
crop_size = int(math.ceil(length / float(stride)) * stride)
preprocessing.crop_size = crop_size
message = (
f"Computed crop size: {crop_size}px, sized to hold every training mask "
f"around its {'bounding-box midpoint' if crop_centering == 'bbox' else 'center of mass'}"
)
if max_hw[0] is not None:
message += (
f", measured after size matching to {max_hw[0]}x{max_hw[1]} (HxW)"
)
logger.warning(message + ".")
train_clipped = sum(req > crop_size for req, _ in train)
val_clipped = sum(req > crop_size for req, _ in val)
if not (train_clipped or val_clipped):
return
message = (
f"{'Computed' if was_auto else 'Configured'} crop size {crop_size}px clips "
f"{train_clipped + val_clipped} of {len(train) + len(val)} masks "
f"({train_clipped} in train, {val_clipped} in validation): the crop is "
f"centered on each mask, and these extend further than {crop_size // 2}px "
"from that center, so the crop cuts part of the animal off."
)
if train_clipped:
required = max(req for req, _ in train)
suggested = int(math.ceil(required / float(stride)) * stride)
message += (
f" Set crop_size to {suggested} to contain every training mask (a "
"smaller crop that keeps the identifying part of the animal can be a "
"deliberate choice)."
)
logger.warning(message)
def _log_crop_size(
self, crop_size, anchor_ind, anchor_part, centroid_method, max_hw
):
"""Report the computed crop size and what it was derived from.
The value depends on the size-matcher scale, the centroid the crop is
centered on and the augmentation margin, none of which a user can infer
from the number alone -- an off-center anchor alone can double it, so a
crop several times the animal's width is expected rather than a bug.
Warned rather than logged at info so it does not scroll past the person
who is trying to choose a crop size (#2862).
Args:
crop_size: The computed crop size, in size-matched pixels.
anchor_ind: Index of the anchor node, or ``None``.
anchor_part: Name of the anchor node, for reporting, or ``None``.
centroid_method: The resolved centroid method.
max_hw: The resolved ``(max_height, max_width)``.
"""
scale = self.config.data_config.preprocessing.scale or 1.0
if anchor_ind is not None:
center = (
f"anchor node {anchor_part!r}"
if anchor_part is not None
else f"anchor node index {anchor_ind}"
)
else:
center = f"{centroid_method or 'center_of_mass'} centroid"
message = (
f"Computed crop size: {crop_size}px, sized to reach every labeled "
f"node from the {center} it is centered on"
)
if max_hw[0] is not None:
message += (
f", measured after size matching to {max_hw[0]}x{max_hw[1]} (HxW)"
)
if scale != 1.0:
message += (
f". Input scaling {scale} is applied to the crop, so the network "
f"input is {int(crop_size * scale)}px"
)
logger.warning(message + ".")
def _warn_if_crop_size_clips(
self,
crop_size,
max_hw,
anchor_ind,
centroid_method,
centroid_fallback,
was_auto=False,
):
"""Warn when the crop size in force clips labeled instances.
Checks both splits and reports them separately, because which split an
instance falls in determines what the user should do about it.
The crop size is computed from the TRAINING split only -- the behavior
SLEAP has had since 1.x, where `find_instance_crop_size` was likewise
passed `training_labels`. Keeping the crop a property of the training
data means a validation instance larger than anything in training is not
covered, so it is called out here rather than left to be discovered in a
validation visualization (#2862). An explicit crop size is never
overridden either -- clipping an extremity may well be a deliberate trade
against GPU memory -- but it should not be silent, since the clipped
nodes are dropped from the targets.
Args:
crop_size: The crop size in force, in size-matched pixels.
max_hw: The resolved ``(max_height, max_width)``.
anchor_ind: Index of the anchor node, or ``None``.
centroid_method: The resolved centroid method.
centroid_fallback: The reduce method for a non-visible anchor node.
was_auto: Whether ``crop_size`` was computed rather than configured,
which changes what the user can do about it.
"""
def tally(labels_list):
"""Return ``(n_clipped, n_total, max_required)`` over one split."""
n_clipped = n_total = 0
max_required = 0.0
for labels in labels_list:
clipped, total, required = count_clipped_instances(
labels,
crop_size=crop_size,
max_hw=max_hw,
anchor_ind=anchor_ind,
centroid_method=centroid_method,
centroid_fallback=centroid_fallback,
user_instances_only=self.config.data_config.user_instances_only,
superseded_only=self.model_type == "embedding",
)
n_clipped += clipped
n_total += total
max_required = max(max_required, required)
return n_clipped, n_total, max_required
# Tallied per split, because which split an instance is in determines
# what the user should do about it.
train_clipped, train_total, train_required = tally(self.train_labels)
val_clipped, val_total, val_required = tally(self.val_labels)
n_clipped = train_clipped + val_clipped
if n_clipped == 0:
return
n_total = train_total + val_total
stride = self.config.model_config.backbone_config[f"{self.backbone_type}"][
"max_stride"
]
origin = "Computed" if was_auto else "Configured"
message = (
f"{origin} crop size {crop_size}px clips {n_clipped} of {n_total} "
f"labeled instances ({train_clipped} in train, {val_clipped} in "
f"validation): crops are centered on the instance centroid, and these "
f"instances have nodes further from it than {crop_size // 2}px. Those "
f"nodes are dropped from the targets."
)
if train_clipped:
# Only ever suggest a size that covers TRAIN. Deriving it from val
# too would fit a hyperparameter to held-out data, which is the very
# thing the train-only sizing exists to avoid.
suggested = math.ceil(train_required / float(stride)) * int(stride)
message += (
f" Set crop_size to {suggested} to contain every training instance."
)
if val_clipped:
# Deliberately no size to set: the crop is sized from train alone so
# that it generalizes to unseen data, and a val instance bigger than
# anything in train is a coverage problem in the labels, not a knob.
plural = "" if val_clipped == 1 else "s"
message += (
f" The crop size is sized from the training split alone, so that it "
f"generalizes to new data rather than being fitted to this project's "
f"validation set -- so {val_clipped} validation instance{plural}, "
f"larger than anything labeled for training, "
f"{'is' if val_clipped == 1 else 'are'} not covered."
)
logger.warning(message)
def _get_confmap_sigma(self, output_stride: int) -> float:
"""Return the active head's confmap sigma (input px), else ``output_stride``.
Used to size the tiling overlap so a keypoint's Gaussian blob is not split
across a seam. Segmentation / non-confmap heads have no sigma; callers only
rely on this for confmap-based tiled models (Phase A).
Args:
output_stride: Head output stride, used as the fallback value.
Returns:
The confmap Gaussian sigma in input pixels.
"""
head_cfg = self.config.model_config.head_configs[self.model_type]
if head_cfg is not None:
for head_layer in head_cfg:
sub = head_cfg[head_layer]
if sub is not None and "sigma" in sub and sub["sigma"] is not None:
return float(sub["sigma"])
return float(output_stride)
def _setup_tiling_config(self):
"""Auto-size tiling geometry from labels + backbone, writing back to config.
No-op unless ``data_config.preprocessing.tiling.enabled``. Explicit
``tile_size`` / ``overlap`` values are preserved; only ``None`` values are
auto-sized. Must run after :func:`check_output_strides` (so ``max_stride`` /
``output_stride`` are finalized). :func:`check_tiling` then validates and
snaps the resulting geometry.
"""
tiling = self.config.data_config.preprocessing.tiling
if tiling is None or not tiling.enabled:
return
backbone_type = self.backbone_type
backbone_cfg = self.config.model_config.backbone_config[f"{backbone_type}"]
max_stride = int(backbone_cfg["max_stride"])
output_stride = int(backbone_cfg["output_stride"])
convs_per_block = int(backbone_cfg.get("convs_per_block", 2))
kernel_size = int(backbone_cfg.get("kernel_size", 3))
# Object extent + instance count across train labels.
max_bbox_dim = 0.0
n_instances = 0
for train_label in self.train_labels:
bbox = find_max_instance_bbox_size(train_label)
if bbox > max_bbox_dim:
max_bbox_dim = bbox
n_instances += sum(len(lf.instances) for lf in train_label)
try:
backbone_margin = compute_backbone_context_margin(
backbone_type, max_stride, convs_per_block, kernel_size
)
except ValueError:
# Unsupported backbone under tiling; check_tiling emits the hard error.
backbone_margin = 0
# tile_size (preserve explicit).
if tiling.tile_size is None:
tile_size = compute_suggested_tile_size(
max_bbox_dim, max_stride, output_stride, backbone_margin
)
self.config.data_config.preprocessing.tiling.tile_size = tile_size
logger.info(
f"Auto-sized tiling.tile_size={tile_size} "
f"(max_bbox_dim={max_bbox_dim:.1f}, backbone_margin={backbone_margin})."
)
tile_size = int(self.config.data_config.preprocessing.tiling.tile_size)
# overlap (preserve explicit; conservative + warn when labels sparse).
if tiling.overlap is None:
if n_instances < _SPARSE_LABEL_THRESHOLD:
overlap = (
math.ceil(tiling.min_overlap_fraction * tile_size / output_stride)
* output_stride
)
logger.warning(
f"Only {n_instances} labeled instances "
f"(< {_SPARSE_LABEL_THRESHOLD}); the object-size estimate is "
f"unreliable. Using a conservative tiling.overlap={overlap} "
f"({tiling.min_overlap_fraction:.0%} of tile_size). Set "
"data_config.preprocessing.tiling.overlap explicitly to override."
)
else:
sigma = self._get_confmap_sigma(output_stride)
overlap = compute_suggested_tile_overlap(
tile_size,
max_bbox_dim,
sigma,
output_stride,
backbone_margin,
tiling.min_overlap_fraction,
)
logger.info(
f"Auto-sized tiling.overlap={overlap} "
f"(confmap_sigma={sigma}, backbone_margin={backbone_margin})."
)
self.config.data_config.preprocessing.tiling.overlap = overlap
overlap = int(self.config.data_config.preprocessing.tiling.overlap)
# samples_per_frame default: a conservative grid-tile count for a
# representative (max-sized) frame, so train coverage ~= one grid pass
# per frame. Written back into the live config before dataset creation.
if tiling.samples_per_frame is None:
scale = float(self.config.data_config.preprocessing.scale)
rep_h = self.config.data_config.preprocessing.max_height
rep_w = self.config.data_config.preprocessing.max_width
if rep_h is None or rep_w is None:
# Fall back to the first labeled frame's native size.
shape = getattr(self.train_labels[0].videos[0], "shape", None)
if shape is not None and len(shape) >= 3:
rep_h, rep_w = int(shape[1]), int(shape[2])
else:
img = self.train_labels[0][0].image
rep_h, rep_w = int(img.shape[0]), int(img.shape[1])
sized_hw = (int(rep_h * scale), int(rep_w * scale))
n_tiles = len(
generate_tile_grid(
sized_hw,
tile_size=tile_size,
overlap=overlap,
output_stride=output_stride,
max_stride=max_stride,
min_overlap_fraction=float(tiling.min_overlap_fraction),
)
)
self.config.data_config.preprocessing.tiling.samples_per_frame = max(
1, n_tiles
)
logger.info(
"Auto-sized tiling.samples_per_frame="
f"{max(1, n_tiles)} (grid tiles for a "
f"{sized_hw[0]}x{sized_hw[1]} frame)."
)
def _setup_head_config(self):
"""Setup node, edge and class names in head config."""
# if edges and part names aren't set in head configs, get it from labels object.
head_config = self.config.model_config.head_configs[self.model_type]
# Skeleton-less model types carry no skeleton at all: mask-only labels
# for a centroid model trained off `data_config.centroids_from_masks`,
# the `embedding` model type, and the segmentation heads. None of their
# heads declares part_names or edges, so an absent skeleton is correct --
# but indexing `self.skeletons[0]` unconditionally raised
# `IndexError: list index out of range` here before any guard below could
# be reached, making those model types unconfigurable. Heads that DO need
# a skeleton now say so by name instead.
skeleton = self.skeletons[0] if self.skeletons else None
skeleton_node_names = list(skeleton.node_names) if skeleton is not None else []
for key in head_config:
if "part_names" in head_config[key].keys():
if head_config[key]["part_names"] is None:
if skeleton is None:
message = (
f"model_config.head_configs.{self.model_type}.{key}"
".part_names is null and the labels carry no skeleton, "
"so the node names cannot be inferred. Provide labels "
"with a skeleton, or set part_names explicitly."
)
logger.error(message)
raise ValueError(message)
self.config.model_config.head_configs[self.model_type][key][
"part_names"
] = skeleton.node_names
elif list(head_config[key]["part_names"]) != skeleton_node_names:
# GT confidence-map generation always produces one channel
# per node in the skeleton (custom_datasets.py's
# generate_confmaps has no part_names/subset parameter), while
# the head's own output channel count is len(part_names). An
# explicit part_names that's shorter, longer, or reordered
# relative to the skeleton silently mismatches those two
# channel counts (a confusing tensor-shape error deep in the
# loss) or silently mislabels channels (if the same length
# but reordered). Catch it here, fail-fast, before any data
# loading/model construction.
message = (
f"model_config.head_configs.{self.model_type}.{key}"
f".part_names must exactly match the skeleton's node "
f"names (in order) -- partial/reordered subsets are not "
f"supported. Got {list(head_config[key]['part_names'])!r}, "
f"skeleton has {skeleton_node_names!r}. Set part_names to "
f"null to use the full skeleton automatically."
)
logger.error(message)
raise ValueError(message)
if (
"anchor_part" in head_config[key].keys()
and head_config[key]["anchor_part"] is not None
and self.model_type != "centroid"
):
# `centroid`'s anchor_part deliberately falls back to None (mean
# of visible nodes) when absent from the skeleton -- see the
# comment at custom_datasets.py's centroid branch ("must NOT
# crash"). Every other head type that consumes anchor_part
# (centered_instance, multi_class_topdown,
# centered_instance_segmentation) does `nodes.index(anchor_part)`
# with no such guard, so a typo'd/nonexistent anchor_part
# currently passes setup cleanly and only fails deep inside
# dataset construction -- with an error message that
# misleadingly blames `part_names`, not the actual offending
# `anchor_part` field. Catch it here instead.
if head_config[key]["anchor_part"] not in skeleton_node_names:
message = (
f"model_config.head_configs.{self.model_type}.{key}"
f".anchor_part {head_config[key]['anchor_part']!r} is not "
f"a node in the skeleton {skeleton_node_names!r}."
)
logger.error(message)
raise ValueError(message)
if "edges" in head_config[key].keys():
if head_config[key]["edges"] is None:
if skeleton is None:
message = (
f"model_config.head_configs.{self.model_type}.{key}"
".edges is null and the labels carry no skeleton, so "
"the edges cannot be inferred. Provide labels with a "
"skeleton, or set edges explicitly."
)
logger.error(message)
raise ValueError(message)
edges = [
(x.source.name, x.destination.name) for x in skeleton.edges
]
self.config.model_config.head_configs[self.model_type][key][
"edges"
] = edges
else:
self._validate_explicit_edges(key, head_config[key]["edges"])
self._validate_nonempty_edges(key)
if "classes" in head_config[key].keys():
if head_config[key]["classes"] is None:
tracks = []
for train_label in self.train_labels:
tracks.extend(
[x.name for x in train_label.tracks if x is not None]
)
# Sorted, not `list(set(...))`: set order varies between
# processes (string hash randomization), so the class ->
# channel assignment would differ run to run.
classes = sorted(set(tracks))
if not len(classes):
message = (
f"No tracks found. ID models need tracks to be defined."
)
logger.error(message)
raise Exception(message)
self.config.model_config.head_configs[self.model_type][key][
"classes"
] = classes
self._validate_classes(key)
# NOTE: no per-class identity UUID is minted. The simplified sleap-io
# `Identity` (name + metadata, sleap-io #535) matches by NAME across files
# and retrains, so the class name IS the canonical cross-file identity key;
# the old train→inference uuid bridge is obsolete. A `class_output ==
# "identity"` multi_class model emits `sio.Identity(name=<class name>)` at
# inference (see `predictor._multiclass_identities`).
if self.model_type in BBOX_MODEL_TYPES:
self._setup_bbox_head_config()
def _validate_explicit_edges(self, key: str, edges) -> None:
"""Check an explicitly set PAF `edges` list against the skeleton's edges.
PAF targets are always built from the skeleton's own edges
(`BottomUpDataset` reads `skeletons[0].edge_inds`): one x/y channel pair
per edge, in skeleton order and direction. The head's channel count is
`2 * len(edges)`. So an explicit list of a different length crashed deep
in the loss with a tensor-size mismatch, and one of the same length but
with different, reordered or reversed edges trained without any error
against mislabeled channels. Same contract as `part_names`: it must match
the skeleton exactly.
"""
field = f"model_config.head_configs.{self.model_type}.{key}.edges"
skeleton = self.skeletons[0] if self.skeletons else None
if skeleton is None:
message = (
f"{field} is set, but the labels carry no skeleton. PAF targets "
"are built from the skeleton's edges, so labels with a skeleton "
"are required."
)
logger.error(message)
raise ValueError(message)
try:
given = [tuple(str(node) for node in edge) for edge in edges]
except TypeError:
given = None
if given is None or any(len(edge) != 2 for edge in given):
message = (
f"{field} must be a list of [source, destination] node-name "
f"pairs. Got {edges!r}."
)
logger.error(message)
raise ValueError(message)
skeleton_edges = [(e.source.name, e.destination.name) for e in skeleton.edges]
if given != skeleton_edges:
message = (
f"{field} must exactly match the skeleton's edges (same order and "
f"source -> destination direction) -- subsets, reordered or "
f"reversed edges are not supported. Got {given!r}, skeleton has "
f"{skeleton_edges!r}. Set edges to null to use the skeleton's "
f"edges automatically."
)
logger.error(message)
raise ValueError(message)
def _validate_nonempty_edges(self, key: str) -> None:
"""Fail fast when a PAF head would have no edges to predict (#567).
A single-node skeleton, or one with nodes but no edges, gives a PAF head
with zero output channels; training then failed later with a low-level
shape error in edge-map generation instead of saying why.
"""
edges = self.config.model_config.head_configs[self.model_type][key]["edges"]
if edges is not None and len(edges) == 0:
message = (
f"model_config.head_configs.{self.model_type}.{key}.edges is empty: "
f"the {self.model_type} model groups nodes into instances with "
"part affinity fields, which need at least one skeleton edge. Add "
"edges to the skeleton, or use a model type that doesn't need them "
"(e.g. centroid + centered_instance)."
)
logger.error(message)
raise ValueError(message)
def _validate_classes(self, key: str) -> None:
"""Check the resolved `classes` against the track names in the labels.
The ID datasets map each instance to its class with
`classes.index(instance.track.name)`, so a track name missing from
`classes` crashed on the first sample load with an unhelpful error (inside
a worker traceback when `num_workers > 0`). This covers explicit lists and
inferred ones alike: classes are inferred from the training labels only,
so a track that appears only in the validation labels is also missing.
Duplicates would give one name two channels. Classes no instance uses are
allowed (a head can be sized for more identities than are labeled yet),
but warned about.
"""
field = f"model_config.head_configs.{self.model_type}.{key}.classes"
classes = list(
self.config.model_config.head_configs[self.model_type][key]["classes"]
)
duplicates = sorted({c for c in classes if classes.count(c) > 1})
if duplicates:
message = (
f"{field} lists {duplicates!r} more than once; each class needs "
f"its own output channel. Got {classes!r}."
)
logger.error(message)
raise ValueError(message)
used = set()
for labels in [*self.train_labels, *self.val_labels]:
for lf in labels:
for instance in lf.instances:
if instance.track is not None:
used.add(instance.track.name)
missing = sorted(used - set(classes))
if missing:
message = (
f"{field} has no class for the track name(s) {missing!r}, which "
f"instances in the training/validation labels use. Got "
f"{classes!r}. Every track name used in the labels needs a class "
f"(classes set to null are inferred from the training labels "
f"only, so a track used only in the validation labels is "
f"missing too)."
)
logger.error(message)
raise ValueError(message)
unused = sorted(set(classes) - used)
if unused and used:
logger.warning(
f"{field} includes {unused!r}, which no instance in the "
f"training/validation labels uses; those output channels will "
f"never see a positive target."
)
def _setup_ckpt_path(self):
"""Setup checkpoint path."""
# if run_name is None, assign a new dir name
ckpt_dir = self.config.trainer_config.ckpt_dir
if ckpt_dir is None or ckpt_dir == "" or ckpt_dir == "None":
ckpt_dir = "."
self.config.trainer_config.ckpt_dir = ckpt_dir
run_name = self.config.trainer_config.run_name
run_name_is_empty = run_name is None or run_name == "" or run_name == "None"
# Validate: multi-GPU + disk cache requires explicit run_name
if run_name_is_empty:
is_disk_caching = (
self.config.data_config.data_pipeline_fw
== "torch_dataset_cache_img_disk"
)
num_devices = self._get_trainer_devices()
if is_disk_caching and num_devices > 1:
raise ValueError(
f"Multi-GPU training with disk caching requires an explicit `run_name`.\n\n"
f"Detected {num_devices} device(s) with "
f"`data_pipeline_fw='torch_dataset_cache_img_disk'`.\n"
f"Without an explicit run_name, each GPU worker generates a different "
f"timestamp-based directory, causing cache synchronization failures.\n\n"
f"Please provide a run_name using one of these methods:\n"
f" - CLI: sleap-nn train config.yaml trainer_config.run_name=my_experiment\n"
f" - Config file: Set `trainer_config.run_name: my_experiment`\n"
f" - Python API: train(..., run_name='my_experiment')"
)
# Auto-generate timestamp-based run_name (safe for single GPU or non-disk-cache)
sum_train_lfs = sum([len(train_label) for train_label in self.train_labels])
sum_val_lfs = sum([len(val_label) for val_label in self.val_labels])
run_name = (
datetime.now().strftime("%y%m%d_%H%M%S")
+ f".{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
)
# If checkpoint path already exists, add suffix to prevent overwriting
if (Path(ckpt_dir) / run_name).exists() and (
Path(ckpt_dir) / run_name / "best.ckpt"
).exists():
logger.info(
f"Checkpoint path already exists: {Path(ckpt_dir) / run_name}... adding suffix to prevent overwriting."
)
for i in count(1):
new_run_name = f"{run_name}-{i}"
if not (Path(ckpt_dir) / new_run_name).exists():
run_name = new_run_name
break
self.config.trainer_config.run_name = run_name
# set output dir for cache img
if self.config.data_config.data_pipeline_fw == "torch_dataset_cache_img_disk":
if self.config.data_config.cache_img_path is None:
self.config.data_config.cache_img_path = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
)
def _disable_pretrained_normalize_for_embedding(self):
"""Turn off the pretrained backbone's input normalization for `embedding`.
`PretrainedBackbone` applies model-specific (ImageNet) mean/std inside its
forward, and documents its input contract as ``[0, 1]`` -- which is what
every other model type feeds it. The `embedding` model type breaks that
contract on purpose: `EmbeddingLightningModule._build_input` per-crop
standardizes to ~N(0, 1) before the backbone sees the crop, so applying the
ImageNet shift on top of it hands the stem mean -1.99 / std 4.43 instead of
the ~N(0, 1) it was pretrained on.
This is silent -- it trains, it just trains badly. Measured on the gerbil
re-ID set (DINOv2-with-registers, 3 epochs, seed 0, paired runs): val rank-1
**0.363** with normalization left on vs **0.920** with it off.
Rather than let the default quietly cost 2.5x rank-1, force it off and say
so. Only the pretrained backbone has this knob, so this runs only on that
branch of :meth:`_verify_model_input_channels`. Note this is
``backbone_config.pretrained.normalize`` (input normalization), NOT
``head_configs.embedding.embedding.normalize`` (L2-normalizing the output
vector), which is unrelated and stays on.
"""
pretrained_cfg = self.config.model_config.backbone_config.pretrained
if not OmegaConf.select(pretrained_cfg, "normalize", default=True):
return
pretrained_cfg.normalize = False
logger.info(
"Disabling `model_config.backbone_config.pretrained.normalize` for the "
"`embedding` model type: the embedding pipeline already per-crop "
"standardizes to ~N(0, 1), so the backbone's ImageNet mean/std would "
"double-normalize the input (measured cost: 2.5x val rank-1). Set it "
"back to `true` only if you also disable the per-crop standardize."
)
def _verify_model_input_channels(self):
"""Verify input channels in model_config based on input image and pretrained model weights."""
# check in channels, verify with img channels / ensure_rgb/ ensure_grayscale
if self.train_labels[0] is not None:
img_channels = self.train_labels[0][0].image.shape[-1]
if self.config.data_config.preprocessing.ensure_rgb:
img_channels = 3
if self.config.data_config.preprocessing.ensure_grayscale:
img_channels = 1
model_in_channels = self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels
if model_in_channels != img_channels:
target_format = "grayscale" if model_in_channels == 1 else "rgb"
logger.warning(
f"Image has {img_channels} channel(s) but model has "
f"{model_in_channels} input channel(s). Images will be "
f"converted to {target_format} to fit the model architecture."
)
self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels = img_channels
logger.info(
f"Updating backbone in_channels to {img_channels} based on the input image channels."
)
# verify input img channels with pretrained model ckpts (if any)
if (
self.backbone_type == "convnext" or self.backbone_type == "swint"
) and self.config.model_config.backbone_config[
f"{self.backbone_type}"
].pre_trained_weights is not None:
current_in_channels = self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels
if current_in_channels != 3:
logger.warning(
f"Image has {current_in_channels} channel(s) but the "
f"pretrained {self.backbone_type} backbone requires 3 "
"(ImageNet RGB) input channels. Images will be converted "
"to rgb to fit the model architecture."
)
self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels = 3
# The ImageNet stem needs a 3-channel input, but the `embedding`
# model type keeps its (default grayscale) DATA channels: a 1-channel
# crop is repeated to 3 in `Model.forward`. Flipping ensure_rgb here
# would silently turn an explicitly-grayscale re-ID model into an RGB
# one. For every other model type, sync the data to RGB as before.
if self.model_type != "embedding":
self.config.data_config.preprocessing.ensure_rgb = True
self.config.data_config.preprocessing.ensure_grayscale = False
logger.info(
f"Updating backbone in_channels to 3 based on the pretrained model weights."
)
# External pretrained (HuggingFace) backbones have 3-channel stems and
# expect RGB (grayscale is replicated). Force in_channels=3 + ensure_rgb so
# the data pipeline feeds 3-channel images and the pretrained stem loads.
elif self.backbone_type == "pretrained":
if self.config.model_config.backbone_config.pretrained.in_channels != 3:
self.config.model_config.backbone_config.pretrained.in_channels = 3
logger.info("Updating pretrained backbone in_channels to 3 (RGB stem).")
if self.model_type == "embedding":
# As in the two branches above, the `embedding` model type keeps its
# own DATA channels: a 1-channel crop is repeated to 3 in
# `Model.forward`, so the stem is fed correctly either way, and
# flipping `ensure_rgb` here would silently turn an
# explicitly-grayscale re-ID model into an RGB one.
self._disable_pretrained_normalize_for_embedding()
else:
self.config.data_config.preprocessing.ensure_rgb = True
self.config.data_config.preprocessing.ensure_grayscale = False
elif (
self.backbone_type == "unet"
and self.config.model_config.pretrained_backbone_weights is not None
):
if self.config.model_config.pretrained_backbone_weights.endswith(".ckpt"):
pretrained_backbone_ckpt = torch.load(
self.config.model_config.pretrained_backbone_weights,
map_location="cpu", # this will be loaded on cpu as it's just used to get the input channels
weights_only=False,
)
input_channels = list(pretrained_backbone_ckpt["state_dict"].values())[
0
].shape[
-3
] # get input channels from first layer
if (
self.config.model_config.backbone_config.unet.in_channels
!= input_channels
):
self.config.model_config.backbone_config.unet.in_channels = (
input_channels
)
logger.info(
f"Updating backbone in_channels to {input_channels} based on the pretrained model weights."
)
if input_channels == 1:
self.config.data_config.preprocessing.ensure_grayscale = True
self.config.data_config.preprocessing.ensure_rgb = False
logger.info(
f"Updating data preprocessing to ensure_grayscale to True based on the pretrained model weights."
)
elif input_channels == 3 and self.model_type != "embedding":
# Embedding keeps its (default grayscale) DATA channels even with a
# 3-channel pretrained stem (gray is repeated in Model.forward).
self.config.data_config.preprocessing.ensure_rgb = True
self.config.data_config.preprocessing.ensure_grayscale = False
logger.info(
f"Updating data preprocessing to ensure_rgb to True based on the pretrained model weights."
)
elif self.config.model_config.pretrained_backbone_weights.endswith(".h5"):
input_channels = get_keras_first_layer_channels(
self.config.model_config.pretrained_backbone_weights
)
if (
self.config.model_config.backbone_config.unet.in_channels
!= input_channels
):
self.config.model_config.backbone_config.unet.in_channels = (
input_channels
)
logger.info(
f"Updating backbone in_channels to {input_channels} based on the pretrained model weights."
)
if input_channels == 1:
self.config.data_config.preprocessing.ensure_grayscale = True
self.config.data_config.preprocessing.ensure_rgb = False
logger.info(
f"Updating data preprocessing to ensure_grayscale to True based on the pretrained model weights."
)
elif input_channels == 3 and self.model_type != "embedding":
# Embedding keeps grayscale data with a 3-channel pretrained stem.
self.config.data_config.preprocessing.ensure_rgb = True
self.config.data_config.preprocessing.ensure_grayscale = False
logger.info(
f"Updating data preprocessing to ensure_rgb to True based on the pretrained model weights."
)
def _verify_accelerator_config(self):
"""Verify the configured `trainer_accelerator` is available on this machine.
A saved training config may have been created on a different machine (e.g.
`trainer_accelerator: mps` from a Mac, reloaded on a Linux/CUDA box). Passing an
unavailable accelerator straight to `lightning.Trainer` raises deep inside
`train()`, after dataset/ckpt setup has already run. This check runs early
(from `setup_config()`) and falls back to `"auto"` on a mismatch instead,
letting the existing device-count resolution logic pick the right backend.
"""
accelerator = self.config.trainer_config.trainer_accelerator
if accelerator in ("auto", "cpu"):
return
if accelerator in ("gpu", "cuda"):
available = torch.cuda.is_available()
elif accelerator == "mps":
available = (
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
)
else:
logger.info(
f"Configured accelerator '{accelerator}' is not a recognized option; "
"changing it to 'auto' (Lightning will select the best available device)."
)
self.config.trainer_config.trainer_accelerator = "auto"
return
if available:
logger.info(
f"Configured accelerator '{accelerator}' is available on this machine "
"and will be used."
)
else:
logger.info(
f"Configured accelerator '{accelerator}' is not available on this "
"machine; changing it to 'auto' (Lightning will select the best "
"available device)."
)
self.config.trainer_config.trainer_accelerator = "auto"
def setup_config(self):
"""Compute config parameters."""
logger.info("Setting up config...")
# Normalize empty strings to None for optional wandb fields
if self.config.trainer_config.wandb.prv_runid == "":
self.config.trainer_config.wandb.prv_runid = None
# The carrier (pose / mask) an embedding model trains on; the crop sizing
# below depends on it.
if self.model_type == "embedding":
self._setup_embedding_carrier()
# compute preprocessing parameters from the labels objects and fill in the config
self._setup_preprocessing_config()
# save skeleton to config
skeleton_yaml = yaml.safe_load(SkeletonYAMLEncoder().encode(self.skeletons))
skeleton_names = skeleton_yaml.keys()
self.config["data_config"]["skeletons"] = []
for skeleton_name in skeleton_names:
skl = skeleton_yaml[skeleton_name]
skl["name"] = skeleton_name
self.config["data_config"]["skeletons"].append(skl)
# setup head config - partnames, edges and class names
self._setup_head_config()
# Identity-equality gates (SPEC §4.4): the embedding objective's pos/neg
# sources silently assert "same / different animal"; validate them against the
# declared `data_config.identity` semantics (error for global_id without global
# names; warn for unproofread tracklets / non-deduplicated same-frame negatives).
if self.model_type == "embedding":
# Real `sio.Identity` annotations ground `global_id` grouping directly (no
# `track_names_are_global` promise needed).
has_identities = any(
bool(getattr(label, "identities", None) or [])
for label in self.train_labels
)
objective = OmegaConf.select(
self.config,
"model_config.head_configs.embedding.embedding.objective",
default=None,
)
# The training crops' groups, from the labels alone (before any image
# is cached), so an objective that can never draw a negative fails now.
batch_groups = embedding_training_groups(
self.train_labels,
embedding_membership_from_config(
self.config, self.train_labels + self.val_labels
),
self.config.model_config.head_configs.embedding.embedding.detection_mode,
)
validate_embedding_identity(
objective=objective,
identity=OmegaConf.select(
self.config, "data_config.identity", default=None
),
has_identities=has_identities,
batch_groups=batch_groups,
)
validate_embedding_views(
objective,
use_augmentations_train=bool(
self.config.data_config.use_augmentations_train
),
augmentation_config=self.config.data_config.augmentation_config,
)
# set max stride for the backbone: convnext and swint
if self.backbone_type == "convnext":
self.config.model_config.backbone_config.convnext.max_stride = (
self.config.model_config.backbone_config.convnext.stem_patch_stride
* (2**3)
* 2
)
elif self.backbone_type == "swint":
self.config.model_config.backbone_config.swint.max_stride = (
self.config.model_config.backbone_config.swint.stem_patch_stride
* (2**3)
* 2
)
# set output stride for backbone from head config and verify max stride
self.config = check_output_strides(self.config)
# fail fast on a contradictory / unknown centroid method (#586)
self.config = check_centroid_methods(self.config)
# auto-size + validate tiling geometry (no-op unless tiling.enabled)
self._setup_tiling_config()
self.config = check_tiling(self.config)
# verify the configured accelerator is available on this machine
self._verify_accelerator_config()
# resolve `num_workers: "auto"` on the dataloaders into a concrete count.
# Must run before `_setup_datasets()`, which reads `num_workers` to size
# the in-memory cache budget, and before the dataloaders are built. The
# labels are passed so `auto` can bound itself by the memory each worker
# would add to the image cache, using the same estimate and buffer that
# `_setup_datasets` checks against.
self.config = check_num_workers(
self.config,
train_labels=self.train_labels,
val_labels=self.val_labels,
memory_buffer=MEMORY_BUFFER,
)
# if trainer_devices is None, set it to "auto"
if self.config.trainer_config.trainer_devices is None:
self.config.trainer_config.trainer_devices = (
"auto"
if OmegaConf.select(
self.config, "trainer_config.trainer_device_indices", default=None
)
is None
else len(
OmegaConf.select(
self.config,
"trainer_config.trainer_device_indices",
default=None,
)
)
)
# setup checkpoint path (generates run_name if not specified)
self._setup_ckpt_path()
# Default wandb run name to trainer run_name if not specified
# Note: This must come after _setup_ckpt_path() which generates run_name
if self.config.trainer_config.wandb.name is None:
self.config.trainer_config.wandb.name = self.config.trainer_config.run_name
# verify input_channels in model_config based on input image and pretrained model weights
self._verify_model_input_channels()
def _setup_model_ckpt_dir(self):
"""Create the model ckpt folder and save ground truth labels."""
ckpt_path = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix()
logger.info(f"Setting up model ckpt dir: `{ckpt_path}`...")
# Only rank 0 (or non-distributed) should create directories and save files
if RANK in [0, -1]:
if not Path(ckpt_path).exists():
try:
Path(ckpt_path).mkdir(parents=True, exist_ok=True)
except OSError as e:
message = f"Cannot create a new folder in {ckpt_path}.\n {e}"
logger.error(message)
raise OSError(message)
# Check if we should filter to user-labeled frames only
user_instances_only = OmegaConf.select(
self.config, "data_config.user_instances_only", default=True
)
# Save train and val ground truth labels. The two lists can differ in
# length (e.g. two train files and one val file), so save each in full.
for split, label_list in (
("train", self.train_labels),
("val", self.val_labels),
):
for idx, split_labels in enumerate(label_list):
# Filter to user-labeled frames if needed (for evaluation)
if user_instances_only:
split_labels = self._filter_to_user_labeled(split_labels)
split_labels.save(
Path(ckpt_path) / f"labels_gt.{split}.{idx}.slp",
restore_original_videos=False,
)
# Save test ground truth labels if test paths are provided
test_file_path = OmegaConf.select(
self.config, "data_config.test_file_path", default=None
)
if test_file_path is not None:
# Normalize to list of strings
if isinstance(test_file_path, str):
test_paths = [test_file_path]
else:
test_paths = list(test_file_path)
for idx, test_path in enumerate(test_paths):
# Only save if it's a .slp file (not a video file)
if test_path.endswith(".slp") or test_path.endswith(".pkg.slp"):
try:
test_labels = sio.load_slp(test_path)
# Same GT boxes as training/eval derive (no-op for
# non-bbox models), or the bbox filter drops every
# frame of a pose/mask test file.
self._apply_bbox_source([test_labels])
if user_instances_only:
test_filtered = self._filter_to_user_labeled(
test_labels
)
else:
test_filtered = test_labels
test_filtered.save(
Path(ckpt_path) / f"labels_gt.test.{idx}.slp",
restore_original_videos=False,
)
except Exception as e:
logger.warning(
f"Could not save test ground truth for {test_path}: {e}"
)
def _setup_viz_datasets(self):
"""Setup dataloaders."""
data_viz_config = self.config.copy()
data_viz_config.data_config.data_pipeline_fw = "torch_dataset"
return get_train_val_datasets(
train_labels=self.train_labels,
val_labels=self.val_labels,
config=data_viz_config,
rank=-1,
)
def _setup_datasets(self):
"""Setup dataloaders."""
base_cache_img_path = None
if self.config.data_config.data_pipeline_fw == "torch_dataset_cache_img_memory":
# check available memory. If insufficient memory, default to disk caching.
# Account for DataLoader worker memory overhead
train_num_workers = self.config.trainer_config.train_data_loader.num_workers
val_num_workers = self.config.trainer_config.val_data_loader.num_workers
max_num_workers = max(train_num_workers, val_num_workers)
mem_available = check_cache_memory(
self.train_labels,
self.val_labels,
memory_buffer=MEMORY_BUFFER,
num_workers=max_num_workers,
)
if not mem_available and self.model_type == "embedding":
# The embedding (appearance / re-ID) model must NOT use the lossy JPEG
# disk cache (it silently degrades the model — see the disk-cache gate
# in get_train_val_datasets). Fall back to an uncached, lossless
# pipeline instead of disk.
self.config.data_config.data_pipeline_fw = "torch_dataset"
base_cache_img_path = None
logger.info(
"Insufficient memory for in-memory caching of the `embedding` "
"model; falling back to an uncached pipeline (`torch_dataset`) to "
"avoid the lossy JPEG disk cache."
)
elif not mem_available:
# Validate: multi-GPU + auto-generated run_name + fallback to disk cache
original_run_name = self._initial_config.trainer_config.run_name
run_name_was_auto = (
original_run_name is None
or original_run_name == ""
or original_run_name == "None"
)
if run_name_was_auto and self.trainer.num_devices > 1:
raise ValueError(
f"Memory caching failed and disk caching fallback requires an "
f"explicit `run_name` for multi-GPU training.\n\n"
f"Detected {self.trainer.num_devices} device(s) with insufficient "
f"memory for in-memory caching.\n"
f"Without an explicit run_name, each GPU worker generates a different "
f"timestamp-based directory, causing cache synchronization failures.\n\n"
f"Please provide a run_name using one of these methods:\n"
f" - CLI: sleap-nn train config.yaml trainer_config.run_name=my_experiment\n"
f" - Config file: Set `trainer_config.run_name: my_experiment`\n"
f" - Python API: train(..., run_name='my_experiment')\n\n"
f"Alternatively, use `data_pipeline_fw='torch_dataset'` to disable caching."
)
self.config.data_config.data_pipeline_fw = (
"torch_dataset_cache_img_disk"
)
base_cache_img_path = (
Path(self.config.data_config.cache_img_path)
if self.config.data_config.cache_img_path is not None
else Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
)
logger.info(
f"Insufficient memory for in-memory caching. `jpg` files will be created for disk-caching."
)
self.config.data_config.cache_img_path = base_cache_img_path
elif self.config.data_config.data_pipeline_fw == "torch_dataset_cache_img_disk":
# Get cache img path
base_cache_img_path = (
Path(self.config.data_config.cache_img_path)
if self.config.data_config.cache_img_path is not None
else Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
)
if self.config.data_config.cache_img_path is None:
self.config.data_config.cache_img_path = base_cache_img_path
return get_train_val_datasets(
train_labels=self.train_labels,
val_labels=self.val_labels,
config=self.config,
rank=self.trainer.global_rank,
)
def _setup_loggers_callbacks(self, viz_train_dataset, viz_val_dataset):
"""Create loggers and callbacks."""
logger.info("Setting up callbacks and loggers...")
loggers = []
callbacks = []
# Checkpoint/early-stop SELECTION metric. Contrastive (embedding) objectives
# are NOT well-selected by val/loss, so select on a retrieval metric (SPEC §8).
if self.model_type == "embedding":
# (The objective itself is validated earlier, in `validate_embedding_identity`
# at config setup, before any dataset is built or cached.)
emb_select = OmegaConf.select(
self.config, "trainer_config.eval.select_metric", default="rank1"
)
# Per-metric selection mode: retrieval / verification-AUC / kNN-accuracy are
# higher-better; EER is lower-better. Validate the name so a typo cannot
# silently monitor a key that is never logged (ModelCheckpoint would then
# never save "best" and a strict EarlyStopping would raise).
emb_metric_modes = {
"rank1": "max",
"mAP": "max",
"auc": "max",
"knn_acc": "max",
"eer": "min",
}
if emb_select not in emb_metric_modes:
raise ValueError(
f"trainer_config.eval.select_metric='{emb_select}' is not a valid "
f"embedding selection metric; choose one of "
f"{'|'.join(emb_metric_modes)}."
)
ckpt_monitor = f"eval/val/{emb_select}"
ckpt_mode = emb_metric_modes[emb_select]
# The selection metric only exists on eval epochs, so let the
# checkpointer consider "best" only on those epochs. Same formula the
# callback uses ((epoch + 1) % frequency), so they cannot drift, and it
# keeps a saved "best" tied to the epoch its metric was measured on
# instead of a later epoch carrying a stale value forward.
ckpt_every_n_epochs = int(
OmegaConf.select(
self.config, "trainer_config.eval.frequency", default=1
)
)
elif self.model_type in BBOX_MODEL_TYPES:
bbox_select = self._bbox_select_metric()
if bbox_select is None:
ckpt_monitor, ckpt_mode = "val/loss", "min"
ckpt_every_n_epochs = 1
else:
# Box quality is measured every `eval.frequency` epochs; select and
# early-stop on those epochs only (as for the embedding metric).
ckpt_monitor, ckpt_mode = f"eval/val/{bbox_select}", "max"
ckpt_every_n_epochs = int(
OmegaConf.select(
self.config, "trainer_config.eval.frequency", default=1
)
)
else:
ckpt_monitor, ckpt_mode = "val/loss", "min"
ckpt_every_n_epochs = 1
if self.config.trainer_config.save_ckpt:
ckpt_dirpath = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix()
# checkpoint callback
checkpoint_callback = ModelCheckpoint(
save_top_k=self.config.trainer_config.model_ckpt.save_top_k,
# `last.ckpt` is written by its own every-epoch callback below: here it
# would only be refreshed when a new "best" is saved.
save_last=False,
dirpath=ckpt_dirpath,
filename="best",
# Config-driven monitor/mode (main #690/#692) wins when the user set
# it away from the "val/loss" schema default (e.g. a seg quality
# metric). Left at the default, fall back to the model-type-aware
# default computed above (embedding -> retrieval metric, else val/loss)
# since contrastive objectives are not well-selected by val/loss.
monitor=(
self.config.trainer_config.model_ckpt.monitor
if self.config.trainer_config.model_ckpt.monitor != "val/loss"
else ckpt_monitor
),
mode=(
self.config.trainer_config.model_ckpt.mode
if self.config.trainer_config.model_ckpt.monitor != "val/loss"
else ckpt_mode
),
# Only meaningful (and only non-1) when the monitored metric is the
# eval-gated embedding one; an explicitly configured monitor keeps
# the every-epoch default.
every_n_epochs=(
1
if self.config.trainer_config.model_ckpt.monitor != "val/loss"
else ckpt_every_n_epochs
),
)
callbacks.append(checkpoint_callback)
if self.config.trainer_config.model_ckpt.save_last:
callbacks.append(LastEpochCheckpoint(dirpath=ckpt_dirpath))
# csv log callback
csv_log_keys = [
"epoch",
"train/loss",
"val/loss",
"learning_rate",
"train/time",
"val/time",
]
# Negative-frame split metrics (train + val), only for model types
# that support frame-level negatives and only when the feature is
# enabled. Gating avoids empty columns in the common no-negatives case.
use_negative_frames = OmegaConf.select(
self.config, "data_config.use_negative_frames", default=False
)
if use_negative_frames and self.model_type in [
"single_instance",
"centroid",
"bottomup",
"multi_class_bottomup",
]:
# Aggregate split metrics (all four supported model types).
csv_log_keys.extend(
[
"train/n_positive",
"train/n_negative",
"train/loss_positive",
"train/loss_negative",
"train/loss_positive_unweighted",
"train/loss_negative_unweighted",
"val/n_positive",
"val/n_negative",
"val/loss_positive",
"val/loss_negative",
"val/loss_positive_unweighted",
"val/loss_negative_unweighted",
]
)
# Per-head split metrics (two-head models only).
if self.model_type == "bottomup":
csv_log_keys.extend(
[
"train/confmaps_loss_positive",
"train/confmaps_loss_negative",
"train/paf_loss_positive",
"train/paf_loss_negative",
"val/confmaps_loss_positive",
"val/confmaps_loss_negative",
"val/paf_loss_positive",
"val/paf_loss_negative",
]
)
elif self.model_type == "multi_class_bottomup":
csv_log_keys.extend(
[
"train/confmaps_loss_positive",
"train/confmaps_loss_negative",
"train/classmap_loss_positive",
"train/classmap_loss_negative",
"val/confmaps_loss_positive",
"val/confmaps_loss_negative",
"val/classmap_loss_positive",
"val/classmap_loss_negative",
]
)
# Add model-specific keys for wandb parity
if self.model_type in [
"single_instance",
"centered_instance",
"multi_class_topdown",
]:
csv_log_keys.extend(
[f"train/confmaps/{name}" for name in self.skeletons[0].node_names]
)
if self.model_type == "centroid":
# Foreground/background confmap MSE split -- see
# `LightningModel._log_confmap_fg_bg_loss`. Always computed for
# this model type (not gated on `centroid_focal_loss_alpha`),
# since it's also the diagnostic that motivates whether a
# focal-style loss would help.
csv_log_keys.extend(
[
"train/confmap_loss_fg",
"train/confmap_loss_bg",
"train/confmap_fg_frac",
"val/confmap_loss_fg",
"val/confmap_loss_bg",
"val/confmap_fg_frac",
]
)
if self.model_type == "bottomup":
csv_log_keys.extend(
[
"train/confmaps_loss",
"train/paf_loss",
"val/confmaps_loss",
"val/paf_loss",
]
)
if self.model_type == "multi_class_bottomup":
csv_log_keys.extend(
[
"train/confmaps_loss",
"train/classmap_loss",
"train/class_accuracy",
"val/confmaps_loss",
"val/classmap_loss",
"val/class_accuracy",
]
)
if self.model_type == "multi_class_topdown":
csv_log_keys.extend(
[
"train/confmaps_loss",
"train/classvector_loss",
"train/class_accuracy",
"val/confmaps_loss",
"val/classvector_loss",
"val/class_accuracy",
]
)
if self.model_type == "bottomup_segmentation":
csv_log_keys.extend(
[
"train/fg_loss",
"train/center_loss",
"train/offset_loss",
"val/fg_loss",
"val/center_loss",
"val/offset_loss",
"val/fg_iou",
]
)
if self.model_type == "centered_instance_segmentation":
csv_log_keys.extend(
[
"train/fg_loss",
"val/fg_loss",
"val/fg_iou",
]
)
if self.model_type == "semantic_segmentation":
csv_log_keys.extend(
[
"train/fg_loss",
"val/fg_loss",
"val/fg_iou",
]
)
if self.model_type == "bbox_centernet":
csv_log_keys.extend(
[
f"{stage}/{term}"
for stage in ("train", "val")
for term in ("center_loss", "size_loss", "offset_loss")
]
)
if self.model_type == "bbox_ssd":
csv_log_keys.extend(
[
f"{stage}/{term}"
for stage in ("train", "val")
for term in ("cls_loss", "box_loss")
]
)
# The embedding model's retrieval callback is NOT gated on eval.enabled
# (ModelCheckpoint selects on its metric, so it always runs -- see the
# callback branching below). Its keys therefore must NOT sit inside the
# eval.enabled gate, or the columns would vanish while the callback kept
# logging them. `train/pos_per_anchor` is a per-epoch training metric.
if self.model_type == "embedding":
csv_log_keys.extend(
[
"train/pos_per_anchor",
"eval/val/rank1",
"eval/val/mAP",
"eval/val/auc",
"eval/val/eer",
"eval/val/knn_acc",
]
)
# Eval-callback keys (only when trainer_config.eval.enabled, mirroring
# the callback branching below). These are only computed every
# eval.frequency epochs; CSVLoggerCallback NaN-resets them at the
# start of each validation epoch so non-eval epochs show NaN instead
# of silently repeating the last-computed eval value.
elif self.config.trainer_config.eval.enabled:
if self.model_type == "centroid":
csv_log_keys.extend(
[
"eval/val/centroid_dist_avg",
"eval/val/centroid_dist_median",
"eval/val/centroid_dist_p90",
"eval/val/centroid_dist_p95",
"eval/val/centroid_dist_max",
"eval/val/centroid_precision",
"eval/val/centroid_recall",
"eval/val/centroid_f1",
"eval/val/centroid_n_tp",
"eval/val/centroid_n_fp",
"eval/val/centroid_n_fn",
]
)
elif self.model_type in BBOX_MODEL_TYPES:
csv_log_keys.extend([f"eval/val/{k}" for k in BBOX_EVAL_LOG_KEYS])
elif self.model_type == "semantic_segmentation":
csv_log_keys.extend(
[
"eval/val/fg_mean_iou",
"eval/val/fg_mean_cldice",
"eval/val/fg_mean_boundary_iou",
"eval/val/fg_frame_recall",
]
)
elif self.model_type in (
"bottomup_segmentation",
"centered_instance_segmentation",
):
csv_log_keys.extend(
[
"eval/val/mask_mean_iou",
"eval/val/mask_mean_iou_all_gt",
"eval/val/mask_mean_cldice",
"eval/val/mask_precision",
"eval/val/mask_recall",
"eval/val/mask_f1",
"eval/val/mask_n_tp",
"eval/val/mask_n_fp",
"eval/val/mask_n_fn",
]
)
else:
csv_log_keys.extend(
[
"eval/val/mOKS",
"eval/val/oks_voc_mAP",
"eval/val/oks_voc_mAR",
"eval/val/distance/avg",
"eval/val/distance/p50",
"eval/val/distance/p95",
"eval/val/distance/p99",
"eval/val/mPCK",
"eval/val/PCK_5",
"eval/val/PCK_10",
"eval/val/visibility_precision",
"eval/val/visibility_recall",
]
)
csv_logger = CSVLoggerCallback(
filepath=Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
/ "training_log.csv",
keys=csv_log_keys,
)
callbacks.append(csv_logger)
if self.config.trainer_config.early_stopping.stop_training_on_plateau:
# early stopping callback
es_kwargs = dict(
monitor=ckpt_monitor,
mode=ckpt_mode,
verbose=False,
min_delta=self.config.trainer_config.early_stopping.min_delta,
patience=self.config.trainer_config.early_stopping.patience,
)
if self.model_type == "embedding" or ckpt_monitor.startswith(
"eval/val/bbox_"
):
# The verification metrics (auc/eer) can be NaN on a degenerate val
# set/shard; don't let that abort the run — pose/seg models keep the
# strict defaults. Box metrics likewise exist on eval epochs only.
es_kwargs["strict"] = False
es_kwargs["check_finite"] = False
# The retrieval metric is measured every `eval.frequency` epochs and
# only carried forward in between; check (and count patience) on
# eval epochs only, the same cadence the checkpointer selects on.
callbacks.append(
EvalCadenceEarlyStopping(
eval_frequency=ckpt_every_n_epochs, **es_kwargs
)
)
else:
callbacks.append(EarlyStopping(**es_kwargs))
if self.config.trainer_config.use_wandb:
# wandb logger
wandb_config = self.config.trainer_config.wandb
if wandb_config.wandb_mode == "offline":
os.environ["WANDB_MODE"] = "offline"
else:
if RANK in [0, -1]:
wandb.login(key=self.config.trainer_config.wandb.api_key)
wandb_logger = WandbLogger(
entity=wandb_config.entity,
project=wandb_config.project,
name=wandb_config.name,
save_dir=(
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix(),
id=self.config.trainer_config.wandb.prv_runid,
group=self.config.trainer_config.wandb.group,
)
loggers.append(wandb_logger)
# Log message about wandb local logs cleanup
should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
wandb_config.delete_local_logs is None
and wandb_config.wandb_mode != "offline"
)
if should_delete_wandb_logs:
logger.info(
"WandB local logs will be deleted after training completes. "
"To keep logs, set trainer_config.wandb.delete_local_logs=false"
)
# save the configs as yaml in the checkpoint dir
# Mask API key in both configs to prevent saving to disk
self.config.trainer_config.wandb.api_key = ""
if self._initial_config is not None:
self._initial_config.trainer_config.wandb.api_key = ""
# zmq callbacks
if self.config.trainer_config.zmq.controller_port is not None:
controller_address = "tcp://127.0.0.1:" + str(
self.config.trainer_config.zmq.controller_port
)
# `controller_polling_timeout` is documented in microseconds, while
# pyzmq's `poll` takes milliseconds: convert, rounding down, so the
# default 10 us polls without waiting instead of waiting 10 ms per batch.
poll_timeout_ms = (
self.config.trainer_config.zmq.controller_polling_timeout // 1000
)
callbacks.append(
TrainingControllerZMQ(
address=controller_address, poll_timeout=poll_timeout_ms
)
)
if self.config.trainer_config.zmq.publish_port is not None:
publish_address = "tcp://127.0.0.1:" + str(
self.config.trainer_config.zmq.publish_port
)
callbacks.append(ProgressReporterZMQ(address=publish_address))
# viz callbacks - use unified callback for all visualization outputs
if self.config.trainer_config.visualize_preds_during_training:
viz_dir = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
/ "viz"
)
if not Path(viz_dir).exists():
if RANK in [0, -1]:
Path(viz_dir).mkdir(parents=True, exist_ok=True)
# Get wandb viz config options
log_wandb = self.config.trainer_config.use_wandb and OmegaConf.select(
self.config, "trainer_config.wandb.save_viz_imgs_wandb", default=False
)
wandb_modes = []
if log_wandb:
if OmegaConf.select(
self.config, "trainer_config.wandb.viz_enabled", default=True
):
wandb_modes.append("direct")
if OmegaConf.select(
self.config, "trainer_config.wandb.viz_boxes", default=False
):
wandb_modes.append("boxes")
if OmegaConf.select(
self.config, "trainer_config.wandb.viz_masks", default=False
):
wandb_modes.append("masks")
# Single unified callback handles all visualization outputs
callbacks.append(
UnifiedVizCallback(
model_trainer=self,
train_dataset=viz_train_dataset,
val_dataset=viz_val_dataset,
model_type=self.model_type,
save_local=self.config.trainer_config.save_ckpt,
local_save_dir=viz_dir,
log_wandb=log_wandb,
wandb_modes=wandb_modes if wandb_modes else ["direct"],
wandb_box_size=OmegaConf.select(
self.config, "trainer_config.wandb.viz_box_size", default=5.0
),
wandb_confmap_threshold=OmegaConf.select(
self.config,
"trainer_config.wandb.viz_confmap_threshold",
default=0.1,
),
log_wandb_table=OmegaConf.select(
self.config, "trainer_config.wandb.log_viz_table", default=False
),
img_format=OmegaConf.select(
self.config,
"trainer_config.viz_img_format",
default="png",
),
)
)
# Add custom progress bar with better metric formatting
if self.config.trainer_config.enable_progress_bar:
callbacks.append(SleapProgressBar())
# Add epoch-end evaluation callback if enabled
# Embedding models are selected on a retrieval metric, so the retrieval
# callback is REQUIRED (it logs eval/val/<metric> that ModelCheckpoint reads),
# not gated on eval.enabled. Embedding models are skeleton-less so they must
# never reach the keypoint EpochEndEvaluationCallback below.
if self.model_type == "embedding":
emb_freq = (
OmegaConf.select(
self.config, "trainer_config.eval.frequency", default=1
)
or 1
)
# It logs every selection metric; `trainer_config.eval.select_metric`
# picks which one ModelCheckpoint / EarlyStopping monitor (see above).
callbacks.append(EmbeddingEvaluationCallback(eval_frequency=emb_freq))
elif self.config.trainer_config.eval.enabled:
if self.model_type == "centroid":
# Use centroid-specific evaluation with distance-based metrics
callbacks.append(
CentroidEvaluationCallback(
videos=self.val_labels[0].videos,
eval_frequency=self.config.trainer_config.eval.frequency,
match_threshold=self.config.trainer_config.eval.match_threshold,
)
)
elif self.model_type in BBOX_MODEL_TYPES:
# Detectors have no keypoints for OKS/PCK: COCO-style box AP over
# the val set, decoded with the same decoder as inference.
callbacks.append(
BBoxEvaluationCallback(
eval_frequency=self.config.trainer_config.eval.frequency,
)
)
elif self.model_type in (
"bottomup_segmentation",
"centered_instance_segmentation",
"semantic_segmentation",
):
# Segmentation has no keypoint predictions for OKS/PCK. Instance seg
# (bottomup / centered-instance) recovers per-instance masks on the
# shared preprocessed grid and reports instance-level mask-IoU mAP /
# precision / recall (grouping-sensitive). Semantic (whole-frame
# foreground) has no instances, so it runs the matching-free
# foreground IoU / clDice / boundary-IoU variant. Both complement the
# coarse val/fg_iou logged in validation_step.
callbacks.append(
SegmentationEvaluationCallback(
eval_frequency=self.config.trainer_config.eval.frequency,
match_threshold=self.config.trainer_config.eval.match_threshold,
foreground=(self.model_type == "semantic_segmentation"),
)
)
else:
# Use standard OKS/PCK evaluation for pose models
callbacks.append(
EpochEndEvaluationCallback(
skeleton=self.skeletons[0],
videos=self.val_labels[0].videos,
eval_frequency=self.config.trainer_config.eval.frequency,
oks_stddev=self.config.trainer_config.eval.oks_stddev,
oks_scale=self.config.trainer_config.eval.oks_scale,
)
)
# Sync the tiling sampler + shared epoch tensor each epoch (tiling only).
tiling = OmegaConf.select(
self.config, "data_config.preprocessing.tiling", default=None
)
if tiling is not None and tiling.enabled:
callbacks.append(TilingEpochCallback())
return loggers, callbacks
def _delete_cache_imgs(self):
"""Delete cache images in disk."""
base_cache_img_path = Path(self.config.data_config.cache_img_path)
train_cache_img_path = Path(base_cache_img_path) / "train_imgs"
val_cache_img_path = Path(base_cache_img_path) / "val_imgs"
if (train_cache_img_path).exists():
logger.info(f"Deleting cache imgs from `{train_cache_img_path}`...")
shutil.rmtree(
(train_cache_img_path).as_posix(),
ignore_errors=True,
)
if (val_cache_img_path).exists():
logger.info(f"Deleting cache imgs from `{val_cache_img_path}`...")
shutil.rmtree(
(val_cache_img_path).as_posix(),
ignore_errors=True,
)
def train(self, dry_run: bool = False):
"""Train the lightning model.
Args:
dry_run: If True, check the run in seconds instead of training: every
setup step (labels, config resolution and validation, model,
dataloaders) with the config's own settings, including its data
pipeline and worker count, then a single training batch and a
single validation batch (Lightning's `fast_dev_run=1`) on one
batch's worth of frames (see `_setup_dry_run_datasets`). It also
checks the run can write where it will: the configs, any image
cache and a checkpoint go to a hidden `.dryrun-*` folder next to
the run folder, removed afterwards with any folder the dry run had
to create, so nothing is left behind. No logs or wandb run.
"""
if not dry_run:
return self._train(dry_run=False)
self._dry_run_dirs = []
try:
return self._train(dry_run=True)
finally:
self._remove_dry_run_dirs()
def _make_dry_run_dir(self, base: Path) -> Path:
"""Create a dry run's scratch folder in `base`, as the real run would.
Creating it (and `base`, if missing) is the write check: a folder the run
couldn't write to fails here instead of hours into training. The folder
name can't collide with a run folder, and a dry run that was killed
before cleanup leaves only an obviously-safe-to-delete `.dryrun-*`.
"""
base = Path(base)
# The outermost folder this creates, so cleanup removes no more than that.
created_root = None
for parent in [base, *base.parents]:
if parent.exists():
break
created_root = parent
run_name = self.config.trainer_config.run_name
path = base / f".dryrun-{run_name}-{os.getpid()}"
try:
path.mkdir(parents=True)
except OSError as e:
# One line, so the cause is the error's last line (what a GUI shows).
message = f"Cannot create a new folder in {base}: {e}"
logger.error(message)
raise OSError(message)
self._dry_run_dirs.append((path, base, created_root))
return path
def _remove_dry_run_dirs(self):
"""Delete what a dry run wrote, and the folders it had to create.
Newest first: a later scratch folder can sit inside a folder an earlier
one created, which can only be removed once it's empty.
"""
for path, base, created_root in reversed(self._dry_run_dirs):
shutil.rmtree(path, ignore_errors=True)
if created_root is None:
continue
folder = base
while True:
try:
folder.rmdir() # only if empty: never anything of the user's
except OSError:
break
if folder == created_root:
break
folder = folder.parent
self._dry_run_dirs = []
def _setup_dry_run_datasets(self):
"""Datasets for a dry run: the config's pipeline on one batch of frames.
Returns `(train_dataset, val_dataset, full_train_dataset,
full_val_dataset)`. The full datasets are uncached, which only indexes
the samples and reads no image, and give the counts a real run would
have (samples, steps per epoch). Then the labels are cut, in memory, to
one batch's worth of labeled frames each (`Labels.extract(copy=False)`
shares the videos, embedded or not), and the datasets the batches come
from are built from those with the configured data pipeline and workers:
the real caching path and worker processes, without caching everything.
Embedding models keep their full labels, since their P x K sampler needs
several identities.
"""
self._dry_run_cache_fits = None
counting_config = deepcopy(self.config)
counting_config.data_config.data_pipeline_fw = "torch_dataset"
full_train_dataset, full_val_dataset = get_train_val_datasets(
train_labels=self.train_labels,
val_labels=self.val_labels,
config=counting_config,
rank=self.trainer.global_rank,
)
if self.config.data_config.data_pipeline_fw == "torch_dataset_cache_img_memory":
# The subset always fits, so check the real cache here (and report it).
workers = max(
self.config.trainer_config.train_data_loader.num_workers,
self.config.trainer_config.val_data_loader.num_workers,
)
self._dry_run_cache_fits = check_cache_memory(
self.train_labels,
self.val_labels,
memory_buffer=MEMORY_BUFFER,
num_workers=workers,
)
self._dry_run_frame_counts = (
sum(len(x) for x in self.train_labels),
sum(len(x) for x in self.val_labels),
)
if self.model_type != "embedding":
def first_frames(labels_list, n):
subsets = []
for labels in labels_list:
inds = [i for i, lf in enumerate(labels) if lf.has_user_instances]
subsets.append(
labels.extract(
(inds or list(range(len(labels))))[:n], copy=False
)
)
return subsets
self.train_labels = first_frames(
self.train_labels,
self.config.trainer_config.train_data_loader.batch_size,
)
self.val_labels = first_frames(
self.val_labels, self.config.trainer_config.val_data_loader.batch_size
)
# A disk cache (configured, or the memory cache's fallback) goes to a
# scratch folder where the real one would be, removed afterwards.
# A cache in the run folder (no `cache_img_path`, or the disk pipeline's
# default, which setup sets to the run folder) goes in the run's scratch
# folder, the first one `_train` made: a scratch folder of its own would
# create a folder with the run's name. Elsewhere it gets its own.
cache_base = self.config.data_config.cache_img_path
run_dir = Path(self.config.trainer_config.ckpt_dir) / (
self.config.trainer_config.run_name
)
scratch = (
self._dry_run_dirs[0][0]
if cache_base is None or Path(cache_base) == run_dir
else self._make_dry_run_dir(cache_base)
)
self.config.data_config.cache_img_path = (scratch / "cache").as_posix()
train_dataset, val_dataset = self._setup_datasets()
return train_dataset, val_dataset, full_train_dataset, full_val_dataset
def _train(self, dry_run: bool):
"""Train the lightning model (see `train`)."""
logger.info(f"Setting up for {'a dry run' if dry_run else 'training'}...")
start_setup_time = time.time()
# initialize the labels object and update config.
if not len(self.train_labels) or not len(self.val_labels):
# Reached when the bare constructor was used (`ModelTrainer(config=...)`)
# instead of `get_model_trainer_from_config`. This guard used to pass the
# CONFIG where `_setup_train_val_labels` expects a `List[sio.Labels]`, so
# it dereferenced `labels[0].skeletons` and died with
# `ConfigKeyError: Missing key 0` -- it could only ever raise. Both entry
# points now run the SAME initializer, so this path cannot drift from the
# factory again (it had already lost `_set_seed` and the video check).
self._initialize_from_config()
if dry_run:
# The configs and checkpoint a real run writes go here instead.
dry_run_dir = self._make_dry_run_dir(self.config.trainer_config.ckpt_dir)
else:
# create the ckpt dir.
self._setup_model_ckpt_dir()
# create the train and val datasets for visualization.
viz_train_dataset = None
viz_val_dataset = None
if self.config.trainer_config.visualize_preds_during_training and not dry_run:
logger.info(f"Setting up visualization train and val datasets...")
viz_train_dataset, viz_val_dataset = self._setup_viz_datasets()
# setup loggers and callbacks for Trainer.
logger.info(f"Setting up Trainer...")
if dry_run:
# Every logger and callback writes into the run folder (or publishes
# over ZMQ), so a dry run uses none.
loggers, callbacks = False, []
else:
loggers, callbacks = self._setup_loggers_callbacks(
viz_train_dataset=viz_train_dataset, viz_val_dataset=viz_val_dataset
)
# set up the strategy (for multi-gpu training)
strategy = OmegaConf.select(
self.config, "trainer_config.trainer_strategy", default="auto"
)
# set up profilers
cfg_profiler = self.config.trainer_config.profiler
profiler = None
if cfg_profiler is not None:
if cfg_profiler in self._profilers:
profiler = self._profilers[cfg_profiler]
else:
message = f"{cfg_profiler} is not a valid option. Please choose one of {list(self._profilers.keys())}"
logger.error(message)
raise ValueError(message)
devices = (
OmegaConf.select(
self.config, "trainer_config.trainer_device_indices", default=None
)
if OmegaConf.select(
self.config, "trainer_config.trainer_device_indices", default=None
)
is not None
else self.config.trainer_config.trainer_devices
)
logger.info(f"Trainer devices: {devices}")
# if trainer devices is set to less than the number of available GPUs, use the least used GPUs
if (
torch.cuda.is_available()
and self.config.trainer_config.trainer_accelerator != "cpu"
and isinstance(self.config.trainer_config.trainer_devices, int)
and self.config.trainer_config.trainer_devices < torch.cuda.device_count()
and self.config.trainer_config.trainer_device_indices is None
):
devices = [
int(x)
for x in np.argsort(get_gpu_memory())[::-1][
: self.config.trainer_config.trainer_devices
]
]
# Sort device indices in ascending order for NCCL compatibility.
# NCCL expects devices in consistent ascending order across ranks
# to properly set up communication rings. Without sorting, DDP may
# assign multiple ranks to the same GPU, causing "Duplicate GPU detected" errors.
devices.sort()
logger.info(f"Using GPUs with most available memory: {devices}")
# create lightning.Trainer instance.
# The embedding loader uses a custom group-aware ``batch_sampler`` that shards
# itself per rank (seed + rank). Lightning's default sampler replacement cannot
# inject a DistributedSampler into a custom batch_sampler (it raises), so disable
# replacement for embedding; every other model type passes an explicit
# DistributedSampler under DDP, so the default (True) is correct for them.
use_distributed_sampler = self.model_type != "embedding"
self.trainer = L.Trainer(
callbacks=callbacks,
logger=loggers,
enable_checkpointing=self.config.trainer_config.save_ckpt and not dry_run,
devices=devices,
max_epochs=self.config.trainer_config.max_epochs,
accelerator=self.config.trainer_config.trainer_accelerator,
enable_progress_bar=self.config.trainer_config.enable_progress_bar,
strategy=strategy,
profiler=profiler,
log_every_n_steps=1,
use_distributed_sampler=use_distributed_sampler,
# One train batch + one val batch, then stop.
fast_dev_run=1 if dry_run else False,
)
self.trainer.strategy.barrier()
# setup datasets. Steps per epoch come from the full datasets; on a dry run
# the batches come from a subset of them.
if dry_run:
train_dataset, val_dataset, full_train_dataset, full_val_dataset = (
self._setup_dry_run_datasets()
)
else:
train_dataset, val_dataset = self._setup_datasets()
full_train_dataset, full_val_dataset = train_dataset, val_dataset
# Barrier after dataset creation to ensure all workers wait for disk caching
# (rank 0 caches to disk, others must wait before reading cached files)
self.trainer.strategy.barrier()
# set-up steps per epoch
train_steps_per_epoch = self.config.trainer_config.train_steps_per_epoch
tiling = OmegaConf.select(
self.config, "data_config.preprocessing.tiling", default=None
)
if train_steps_per_epoch is None:
if (
tiling is not None
and tiling.enabled
and tiling.steps_per_epoch is not None
):
# TRAIN decouple: the tiling knob overrides the tile-count length.
train_steps_per_epoch = tiling.steps_per_epoch
logger.info(
f"train_steps_per_epoch not set; using tiling.steps_per_epoch={train_steps_per_epoch}"
)
else:
# One pass over the data. `embedding` draws P x K crops per step, not
# `train_data_loader.batch_size`, so the step size comes from the
# dataset-aware helper the dataloader factory also uses.
samples_per_step = get_train_samples_per_step(
full_train_dataset, self.config
)
train_steps_per_epoch = get_steps_per_epoch(
dataset=full_train_dataset, batch_size=samples_per_step
)
logger.info(
f"train_steps_per_epoch not set; computed {train_steps_per_epoch} "
f"from training dataset ({len(full_train_dataset)} samples, "
f"{samples_per_step} per step)"
)
else:
logger.info(
f"Using configured train_steps_per_epoch={train_steps_per_epoch}"
)
min_train_steps_per_epoch = self.config.trainer_config.min_train_steps_per_epoch
if min_train_steps_per_epoch > train_steps_per_epoch:
logger.info(
f"train_steps_per_epoch={train_steps_per_epoch} is below "
f"min_train_steps_per_epoch={min_train_steps_per_epoch}; using the minimum"
)
train_steps_per_epoch = min_train_steps_per_epoch
self.config.trainer_config.train_steps_per_epoch = train_steps_per_epoch
logger.info(f"Final train_steps_per_epoch={train_steps_per_epoch}")
# VAL: always full-coverage (every grid tile visited once), NOT decoupled.
val_steps_per_epoch = get_steps_per_epoch(
dataset=full_val_dataset,
batch_size=self.config.trainer_config.val_data_loader.batch_size,
)
logger.info(f"Training on {self.trainer.num_devices} device(s)")
logger.info(f"Training on {self.trainer.strategy.root_device} accelerator")
# initialize the lightning model.
# need to initialize after Trainer is initialized (for trainer accelerator)
logger.info(f"Setting up lightning module for {self.model_type} model...")
self.lightning_model = LightningModel.get_lightning_model_from_config(
config=self.config,
)
logger.info(f"Backbone model: {self.lightning_model.model.backbone}")
logger.info(f"Head model: {self.lightning_model.model.head_layers}")
total_params = sum(p.numel() for p in self.lightning_model.parameters())
logger.info(f"Total model parameters: {total_params:,}")
self.config.model_config.total_params = total_params
# setup dataloaders
# need to set up dataloaders after Trainer is initialized (for ddp). DistributedSampler depends on the rank
logger.info(
f"Input image shape: {train_dataset[0]['image'].shape if 'image' in train_dataset[0] else train_dataset[0]['instance_image'].shape}"
)
train_dataloader, val_dataloader = get_train_val_dataloaders(
train_dataset=train_dataset,
val_dataset=val_dataset,
config=self.config,
rank=self.trainer.global_rank,
train_steps_per_epoch=self.config.trainer_config.train_steps_per_epoch,
val_steps_per_epoch=val_steps_per_epoch,
trainer_devices=self.trainer.num_devices,
)
if dry_run:
# The run's two config files, written where the real run writes them.
OmegaConf.save(
self._initial_config, (dry_run_dir / "initial_config.yaml").as_posix()
)
OmegaConf.save(
self.config, (dry_run_dir / "training_config.yaml").as_posix()
)
# save config only in rank 0 process
if self.trainer.global_rank == 0 and not dry_run:
ckpt_path = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix()
# Overwrite version with current sleap-nn version
self._initial_config.sleap_nn_version = sleap_nn.__version__
self.config.sleap_nn_version = sleap_nn.__version__
OmegaConf.save(
self._initial_config,
(Path(ckpt_path) / "initial_config.yaml").as_posix(),
)
if self.config.trainer_config.use_wandb:
if wandb.run is None:
wandb.init(
dir=(
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix(),
project=self.config.trainer_config.wandb.project,
entity=self.config.trainer_config.wandb.entity,
name=self.config.trainer_config.wandb.name,
id=self.config.trainer_config.wandb.prv_runid,
group=self.config.trainer_config.wandb.group,
)
# Define custom x-axes for wandb metrics
# Epoch-level metrics use epoch as x-axis, step-level use default global_step
wandb.define_metric("epoch")
# Training metrics (train/ prefix for grouping) - all use epoch x-axis
wandb.define_metric("train/*", step_metric="epoch")
wandb.define_metric("train/confmaps/*", step_metric="epoch")
# Validation metrics (val/ prefix for grouping)
wandb.define_metric("val/*", step_metric="epoch")
# Evaluation metrics (eval/ prefix for grouping)
wandb.define_metric("eval/*", step_metric="epoch")
# Visualization images (need explicit nested paths)
wandb.define_metric("viz/*", step_metric="epoch")
wandb.define_metric("viz/train/*", step_metric="epoch")
wandb.define_metric("viz/val/*", step_metric="epoch")
self.config.trainer_config.wandb.current_run_id = wandb.run.id
wandb.config["run_name"] = self.config.trainer_config.wandb.name
wandb.config["run_config"] = OmegaConf.to_container(
self.config, resolve=True
)
OmegaConf.save(
self.config,
(
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
/ "training_config.yaml"
).as_posix(),
)
self.trainer.strategy.barrier()
resume_ckpt_path = self.config.trainer_config.resume_ckpt_path
if dry_run and resume_ckpt_path is not None:
# `fast_dev_run` caps the run at one step and one epoch, so restoring the
# checkpoint's trainer state would leave no batch to run (and Lightning
# refuses a checkpoint past epoch 1 outright). Load its weights strictly
# instead: that checks the checkpoint fits this model, and the one
# training batch then runs from them.
logger.info(f"Dry run: loading weights from {resume_ckpt_path}")
ckpt = torch.load(resume_ckpt_path, map_location="cpu", weights_only=False)
self.lightning_model.load_state_dict(ckpt["state_dict"])
resume_ckpt_path = None
# Flag to track if training was interrupted (not completed normally)
training_interrupted = False
try:
logger.info(
f"Finished trainer set up. [{time.time() - start_setup_time:.1f}s]"
)
logger.info(f"Starting training loop...")
start_train_time = time.time()
self.trainer.fit(
self.lightning_model,
train_dataloader,
val_dataloader,
ckpt_path=resume_ckpt_path,
)
except KeyboardInterrupt:
logger.info("Stopping training...")
training_interrupted = True
finally:
logger.info(
f"Finished training loop. [{(time.time() - start_train_time) / 60:.1f} min]"
)
# Note: wandb.finish() is called in train.py after post-training evaluation
# delete image disk caching
if (
self.config.data_config.data_pipeline_fw
== "torch_dataset_cache_img_disk"
and self.config.data_config.delete_cache_imgs_after_training
):
if self.trainer.global_rank == 0:
self._delete_cache_imgs()
# delete viz folder if requested
if (
self.config.trainer_config.visualize_preds_during_training
and not self.config.trainer_config.keep_viz
and not dry_run
):
if self.trainer.global_rank == 0:
viz_dir = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
/ "viz"
)
if viz_dir.exists():
logger.info(f"Deleting viz folder at {viz_dir}...")
shutil.rmtree(viz_dir, ignore_errors=True)
# Clean up entire run folder if training was interrupted (KeyboardInterrupt)
# (Never on a dry run: it created no run folder, so one at this path
# belongs to someone else.)
if training_interrupted and self.trainer.global_rank == 0 and not dry_run:
run_dir = (
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
)
if run_dir.exists():
logger.info(
f"Training canceled - cleaning up run folder at {run_dir}..."
)
shutil.rmtree(run_dir, ignore_errors=True)
if dry_run and not training_interrupted:
if self.config.trainer_config.save_ckpt:
# Checkpoint serialization and the write, where the run's go.
self.trainer.save_checkpoint((dry_run_dir / "best.ckpt").as_posix())
self._log_dry_run_summary(
full_train_dataset, full_val_dataset, val_steps_per_epoch
)
def _log_dry_run_summary(self, train_dataset, val_dataset, val_steps_per_epoch):
"""Log what a real run with this config would do, after a passing dry run."""
cfg = self.config
run_dir = Path(cfg.trainer_config.ckpt_dir) / cfg.trainer_config.run_name
n_train, n_val = self._dry_run_frame_counts
workers = cfg.trainer_config.train_data_loader.num_workers
lines = [
"Dry run passed: setup, one training batch and one validation batch "
"all ran, with this config's settings on one batch of frames.",
f" Model type: {self.model_type}",
f" Labeled frames: {n_train} train, {n_val} val",
f" Samples: {len(train_dataset)} train, "
f"{len(val_dataset)} val",
f" Steps per epoch: {cfg.trainer_config.train_steps_per_epoch} "
f"train, {val_steps_per_epoch} val",
f" Max epochs: {cfg.trainer_config.max_epochs}",
f" Model parameters: {cfg.model_config.total_params:,}",
f" Devices: {self.trainer.num_devices} x "
f"{self.trainer.strategy.root_device}",
f" Run folder: {run_dir.as_posix() if cfg.trainer_config.save_ckpt else '(none: save_ckpt is false)'}",
f" Data pipeline: {cfg.data_config.data_pipeline_fw}, "
f"{workers} worker{'' if workers == 1 else 's'}",
]
cache_fits = self._dry_run_cache_fits
if cache_fits is not None:
lines.append(
" Image cache: "
+ (
"fits in memory"
if cache_fits
else "does not fit in memory; a real run falls back to a "
"slower pipeline"
)
)
written = (
"configs and a checkpoint" if cfg.trainer_config.save_ckpt else "configs"
)
lines.append(
f" Write check: wrote {written} to "
f"{Path(cfg.trainer_config.ckpt_dir).as_posix()}, then removed them"
)
if cfg.trainer_config.resume_ckpt_path is not None:
lines.append(
f" Resumes from: {cfg.trainer_config.resume_ckpt_path}"
)
logger.info("\n".join(lines))