Model Configuration

Model architectures in DeepLabCut PyTorch are defined using configuration files written in YAML format. These configuration files specify the model architecture, training hyperparameters, data augmentation settings, and more.

Configuration File Structure

The primary configuration file is named pytorch_cfg.yaml and is stored in the model's training directory. This file is automatically generated during the standard DeepLabCut workflow, but can also be created manually for custom projects.

Creating Configuration Files

The deeplabcut.pose_estimation_pytorch.config module provides functions to create and manipulate configuration files.

Basic Configuration Creation

Use make_pytorch_pose_config to generate a model configuration:

from pathlib import Path
import deeplabcut.pose_estimation_pytorch as dlc_torch

# Configuration for a DeepLabCut project
project_cfg = {
    "bodyparts": ["nose", "left_ear", "right_ear", "tail_base"],
    "individuals": ["mouse1", "mouse2"],
    # ... other project settings
}

pose_config_path = Path("/path/to/model/train")
model_cfg = dlc_torch.config.make_pytorch_pose_config(
    project_config=project_cfg,
    pose_config_path=pose_config_path,
    net_type="hrnet_w32",
    top_down=True,
    save=True,  # Save the configuration to disk
)

Configuration for COCO Datasets

For COCO-format datasets without a DeepLabCut project, use make_basic_project_config:

from pathlib import Path
import deeplabcut.pose_estimation_pytorch as dlc_torch

# Create a minimal project configuration
project_cfg = dlc_torch.config.make_basic_project_config(
    dataset_path="/path/to/COCOProject",
    bodyparts=["nose", "left_eye", "right_eye", "left_ear", "right_ear"],
    max_individuals=2,
    multi_animal=True,
)

# Generate model configuration
model_cfg = dlc_torch.config.make_pytorch_pose_config(
    project_config=project_cfg,
    pose_config_path=Path("/path/to/experiment/train"),
    net_type="dlcrnet_ms5",
    top_down=False,
    save=True,
)

Configuration File Components

A complete pytorch_cfg.yaml file contains the following sections.

Model Architecture

Specifies the backbone, optional neck, and head:

model:
  backbone:
    type: HRNet
    variant: w32
  neck: null          # omit or set to null for no neck
  head:
    type: HeatmapHead
    weight_init: normal
    predictor:
      type: HeatmapPredictor
      location_refinement: true
      locref_std: 7.2801
    target_generator:
      type: HeatmapGaussianGenerator
      num_heatmaps: "num_bodyparts"
      pos_dist_thresh: 17
      generate_locref: true
    criterion:
      heatmap:
        type: WeightedMSECriterion
        weight: 1.0
      locref:
        type: WeightedHuberCriterion
        weight: 0.05

Data Configuration

Controls data loading, augmentation, and preprocessing. Augmentations under train are applied only during training; inference augmentations are applied during evaluation and video analysis:

data:
  colormode: RGB        # RGB or GRAY
  bbox_margin: 20       # pixels added around bounding boxes (top-down only)
  train:
    normalize_images: true
    crop_sampling:
      width: 448
      height: 448
      max_shift: 0.1
      method: hybrid
    affine:
      p: 0.5
      rotation: 30
      scaling: [0.5, 1.25]
      translation: 0
    gaussian_noise: 12.75
    motion_blur: true
    hflip: true
  inference:
    normalize_images: true

Training Settings

Controls the training loop — batch size, number of epochs, data loading, and random seed:

train_settings:
  batch_size: 8
  epochs: 200
  seed: 42
  dataloader_workers: 4
  dataloader_pin_memory: true
  display_iters: 500

Runner Configuration

The runner manages the training loop, optimisation, checkpointing, and evaluation. Use any optimizer from torch.optim and any scheduler from torch.optim.lr_scheduler:

runner:
  type: PoseTrainingRunner
  gpus: null              # null = use device setting; list of ints for multi-GPU
  key_metric: "test.mAP"
  key_metric_asc: true    # true if higher is better
  eval_interval: 10       # evaluate every N epochs

  optimizer:
    type: AdamW
    params:
      lr: 0.0001
      weight_decay: 0.01

  scheduler:
    type: LRListScheduler
    params:
      milestones: [160, 190]
      lr_list: [[1e-5], [1e-6]]

  snapshots:
    max_snapshots: 5        # keep only the N most recent snapshots
    save_epochs: 25         # save a snapshot every N epochs
    save_optimizer_state: false

  logger:
    type: WandbLogger       # omit or set to null for local-only logging
    project_name: my-project
    tags: ["model=hrnet_w32"]

Resuming Training

Resume from a specific snapshot by setting:

resume_training_from: /path/to/model/train/snapshot-010.pt

Inference Configuration

Controls inference-specific behaviour, set independently of training augmentations:

method: td          # bu = bottom-up, td = top-down, ctd = conditional top-down
device: auto        # auto, cpu, cuda, or cuda:N

Top-Down Detector Configuration

Top-down models require a separate detector. The detector pytorch_cfg.yaml mirrors the pose model structure but uses DetectorTrainingRunner:

runner:
  type: DetectorTrainingRunner
  key_metric: "test.mAP@50:95"
  key_metric_asc: true
  eval_interval: 10
  optimizer:
    type: AdamW
    params:
      lr: 1e-4
  scheduler:
    type: LRListScheduler
    params:
      milestones: [160]
      lr_list: [[1e-5]]
  snapshots:
    max_snapshots: 5
    save_epochs: 25
    save_optimizer_state: false

train_settings:
  batch_size: 1
  epochs: 250
  dataloader_workers: 0
  display_iters: 500