@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
@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 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)
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_centroid_labels(
self, label: sio.Labels, val_fraction: float, seed: Optional[int]
):
"""Train/val split for the centroid model that keeps centroid-only frames.
`sio.Labels.make_training_splits` returns only frames with user
instances, so it would drop pure-centroid frames. This reimplements a
deterministic fractional split over frames with a centroid training
target (user instances OR user centroids).
"""
frames = [lf for lf in label if self._is_training_frame(lf)]
for lf in frames:
lf.instances = lf.user_instances
def mk(selected):
return sio.Labels(
labeled_frames=selected,
videos=label.videos,
skeletons=label.skeletons,
tracks=label.tracks,
suggestions=label.suggestions,
provenance=label.provenance,
)
n = len(frames)
if n <= 1:
# Too few to hold out a val frame; use the same frame for both so
# training still runs (mirrors small-dataset behavior elsewhere).
return mk(list(frames)), mk(list(frames))
rng = np.random.default_rng(seed)
order = rng.permutation(n)
n_val = max(1, int(round(n * val_fraction)))
val_sel = [frames[i] for i in order[:n_val]]
train_sel = [frames[i] for i in order[n_val:]]
return mk(train_sel), mk(val_sel)
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)
# 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` file shave same skeleton structure (if multiple slp file paths are provided)
# Mask-only labels (e.g. instance segmentation) may carry no skeleton.
if self.skeletons:
skeleton = self.skeletons[0]
for index, train_label in enumerate(labels):
if not train_label.skeletons:
continue
skel_temp = train_label.skeletons[0]
skeletons_equal = skeleton.matches(skel_temp)
if not skeletons_equal:
message = f"The skeletons in the training labels: {index + 1} do not match the skeleton in the first training label file."
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
# 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":
# Centroid-aware split keeps pure-centroid frames that
# make_training_splits would drop (no user instances).
train_split, val_split = self._split_centroid_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 _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):
"""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.
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
bbox_size = find_max_instance_bbox_size(
train_label,
max_hw=max_hw,
user_instances_only=self.config.data_config.user_instances_only,
)
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,
)
(
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,
)
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 _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,
)
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
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]
)
classes = list(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
# 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`).
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
# 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
)
validate_embedding_identity(
objective=OmegaConf.select(
self.config,
"model_config.head_configs.embedding.embedding.objective",
default=None,
),
identity=OmegaConf.select(
self.config, "data_config.identity", default=None
),
has_identities=has_identities,
)
# 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()
# 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
for idx, (train, val) in enumerate(zip(self.train_labels, self.val_labels)):
# Filter to user-labeled frames if needed (for evaluation)
if user_instances_only:
train_filtered = self._filter_to_user_labeled(train)
val_filtered = self._filter_to_user_labeled(val)
else:
train_filtered = train
val_filtered = val
train_filtered.save(
Path(ckpt_path) / f"labels_gt.train.{idx}.slp",
restore_original_videos=False,
)
val_filtered.save(
Path(ckpt_path) / f"labels_gt.val.{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)
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":
# Touch the objective/sampler so a missing config fails loud here rather
# than deep in the dataloader.
OmegaConf.select(
self.config,
"model_config.head_configs.embedding.embedding.objective.sampler.kind",
default=None,
)
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
)
)
else:
ckpt_monitor, ckpt_mode = "val/loss", "min"
ckpt_every_n_epochs = 1
if self.config.trainer_config.save_ckpt:
# checkpoint callback
checkpoint_callback = ModelCheckpoint(
save_top_k=self.config.trainer_config.model_ckpt.save_top_k,
save_last=self.config.trainer_config.model_ckpt.save_last,
dirpath=(
Path(self.config.trainer_config.ckpt_dir)
/ self.config.trainer_config.run_name
).as_posix(),
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)
# 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",
]
)
# 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 == "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":
# The retrieval selection metric is logged only on eval epochs (so it is
# absent when eval.frequency > 1) and the verification metrics (auc/eer)
# can be NaN on a degenerate val set/shard. Don't let an absent/NaN value
# abort the run — pose/seg models keep the strict defaults.
es_kwargs["strict"] = False
es_kwargs["check_finite"] = False
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
)
callbacks.append(TrainingControllerZMQ(address=controller_address))
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
)
emb_select = OmegaConf.select(
self.config, "trainer_config.eval.select_metric", default="rank1"
)
callbacks.append(
EmbeddingEvaluationCallback(
eval_frequency=emb_freq, select_metric=emb_select
)
)
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 (
"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):
"""Train the lightning model."""
logger.info(f"Setting up for 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()
# 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:
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...")
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,
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,
)
self.trainer.strategy.barrier()
# setup datasets
train_dataset, val_dataset = self._setup_datasets()
# 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:
train_steps_per_epoch = get_steps_per_epoch(
dataset=train_dataset,
batch_size=self.config.trainer_config.train_data_loader.batch_size,
)
logger.info(
f"train_steps_per_epoch not set; computed {train_steps_per_epoch} from training dataset"
)
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=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 self.trainer.global_rank == 0: # save config only in rank 0 process
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()
# 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=self.config.trainer_config.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
):
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)
if training_interrupted and self.trainer.global_rank == 0:
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)