@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."""
# Verify config structure.
config = verify_training_cfg(config)
model_trainer = cls(config=config)
model_trainer.model_type = get_model_type_from_cfg(model_trainer.config)
model_trainer.backbone_type = get_backbone_type_from_cfg(model_trainer.config)
if model_trainer.config.trainer_config.seed is not None:
model_trainer._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 model_trainer.config.data_config.train_labels_path
]
val_labels = (
[
sio.load_slp(path)
for path in model_trainer.config.data_config.val_labels_path
]
if model_trainer.config.data_config.val_labels_path is not None
else None
)
model_trainer._setup_train_val_labels(
labels=train_labels, val_labels=val_labels
)
else:
model_trainer._setup_train_val_labels(
labels=train_labels, val_labels=val_labels
)
model_trainer._initial_config = model_trainer.config.copy()
# update config parameters
model_trainer.setup_config()
# Check if all videos exist across all labels
all_videos_exist = all(
video.exists(check_all=True)
for labels in [*model_trainer.train_labels, *model_trainer.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."
)
return model_trainer
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
def _is_training_frame(self, lf: "sio.LabeledFrame") -> bool:
"""Whether a frame carries a usable training target.
Normally that means user instances. The centroid model can additionally
train on frames that carry 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).
"""
if lf.has_user_instances:
return True
if self.model_type == "centroid":
return any(not c.is_predicted for c in lf.centroids)
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
# 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
)
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 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}")
# 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 _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 _setup_preprocessing_config(self):
"""Setup preprocessing config."""
# compute max_heigt, max_width, and crop_size (if not provided in the config)
max_height = self.config.data_config.preprocessing.max_height
max_width = self.config.data_config.preprocessing.max_width
if self.model_type in (
"centered_instance",
"multi_class_topdown",
"centered_instance_segmentation",
):
crop_size = self.config.data_config.preprocessing.crop_size
max_h, max_w = 0, 0
max_crop_size = 0
for train_label in self.train_labels:
# compute max h and w from slp file if not provided
if max_height is None or max_width is None:
current_max_h, current_max_w = get_max_height_width(train_label)
if current_max_h > max_h:
max_h = current_max_h
if current_max_w > max_w:
max_w = current_max_w
if self.model_type in (
"centered_instance",
"multi_class_topdown",
"centered_instance_segmentation",
):
# compute crop size if not provided in config
if crop_size is None:
# Get padding from config or auto-compute from augmentation settings
padding = self.config.data_config.preprocessing.crop_padding
if padding is None:
# Auto-compute padding based on augmentation settings
aug_config = self.config.data_config.augmentation_config
if (
self.config.data_config.use_augmentations_train
and aug_config is not None
and aug_config.geometric is not None
):
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 rotation_enabled or scale_enabled:
# First find the actual max bbox size from labels
bbox_size = find_max_instance_bbox_size(train_label)
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
padding = compute_augmentation_padding(
bbox_size=bbox_size,
rotation_max=rotation_max,
scale_max=scale_max,
)
else:
padding = 0
else:
padding = 0
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,
)
if crop_sz > max_crop_size:
max_crop_size = crop_sz
# if preprocessing params were None, replace with computed params
if max_height is None or max_width is None:
self.config.data_config.preprocessing.max_height = max_h
self.config.data_config.preprocessing.max_width = max_w
if (
self.model_type
in (
"centered_instance",
"multi_class_topdown",
"centered_instance_segmentation",
)
and crop_size is None
):
self.config.data_config.preprocessing.crop_size = max_crop_size
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_node_names = list(self.skeletons[0].node_names)
for key in head_config:
if "part_names" in head_config[key].keys():
if head_config[key]["part_names"] is None:
self.config.model_config.head_configs[self.model_type][key][
"part_names"
] = self.skeletons[0].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:
edges = [
(x.source.name, x.destination.name)
for x in self.skeletons[0].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
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 _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
if (
self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels
!= img_channels
):
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:
if (
self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels
!= 3
):
self.config.model_config.backbone_config[
f"{self.backbone_type}"
].in_channels = 3
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).")
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:
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:
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()
# 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)
# 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:
# 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("./")
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 = []
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",
monitor=self.config.trainer_config.model_ckpt.monitor,
mode=self.config.trainer_config.model_ckpt.mode,
)
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 == "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",
]
)
# 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.
if 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
callbacks.append(
EarlyStopping(
monitor="val/loss",
mode="min",
verbose=False,
min_delta=self.config.trainer_config.early_stopping.min_delta,
patience=self.config.trainer_config.early_stopping.patience,
)
)
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
if 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):
self._setup_train_val_labels(self.config)
self.setup_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.
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,
)
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)