base.yaml

model:
  ckpt_path: null
  encoder_cfg: 
    model_name: "resnet18"
    backend: "timm"
    in_chans: 3
  d_model: 1024
  nhead: 8
  num_encoder_layers: 1
  num_decoder_layers: 1
  dropout: 0.1
  activation: "relu"
  return_intermediate_dec: False
  norm: False
  num_layers_attn_head: 2
  dropout_attn_head: 0.1
  embedding_meta: 
    pos:
        mode: "fixed" # supports fixed, learned, rope
        normalize: true
    temp:
        mode: "fixed" # supports fixed, learned, rope
    embedding_agg_method: "stack" # supports stack, average, concatenate
  return_embedding: False
  decoder_self_attn: False

loss:
  neg_unmatched: false
  epsilon: 1e-4
  asso_weight: 1.0

#currently assumes adam. TODO adapt logic for other optimizers like sgd
optimizer:
  name: "Adam"
  lr: 0.001
  betas: [0.9, 0.999]
  eps: 1e-8
  weight_decay: 0.01

#currently assumes reduce lr on plateau
scheduler:
  name: "ReduceLROnPlateau"
  mode: "min"
  factor: 0.5
  patience: 10
  threshold: 1e-4
  threshold_mode: "rel"

tracker:
  window_size: 8
  use_vis_feats: true
  overlap_thresh: 0.01
  mult_thresh: true
  decay_time: null
  iou: null
  max_center_dist: null

runner:
  metrics:
      train: ['num_switches']
      val: ['num_switches']
      test: ['num_switches']
  persistent_tracking:
      train: false
      val: true
      test: true
      
dataset:
  train_dataset:
    slp_files: ["../../tests/data/sleap/two_flies.slp"]
    video_files: ["../../tests/data/sleap/two_flies.mp4"]
    padding: 5
    crop_size: 128
    chunk: true
    clip_length: 32
    mode: "train"

  val_dataset:
    slp_files: ["../../tests/data/sleap/two_flies.slp"]
    video_files: ["../../tests/data/sleap/two_flies.mp4"]
    padding: 5
    crop_size: 128 
    chunk: True
    clip_length: 32
    mode: "val"

  test_dataset:
    slp_files: ["../../tests/data/sleap/two_flies.slp"]
    video_files: ["../../tests/data/sleap/two_flies.mp4"]
    padding: 5
    crop_size: 128 
    chunk: True
    clip_length: 32
    mode: "test"

dataloader:
  train_dataloader:
    shuffle: true
    num_workers: 0
  val_dataloader:
    shuffle: false
    num_workers: 0
  test_dataloader: 
    shuffle: false
    num_workers: 0

logging:
  logger_type: null
  name: "example_train"
  entity: null
  job_type: "train"
  notes: "Example train job"
  dir: "./logs"
  group: "example"
  save_dir: './logs'
  project: "GTR"
  log_model: null

early_stopping:
  monitor: "val_loss"
  min_delta: 0.1
  patience: 10
  mode: "min"
  check_finite: true
  stopping_threshold: 1e-8
  divergence_threshold: 30

checkpointing:
  monitor: ["val_loss","val_num_switches"]
  verbose: true
  save_last: true
  dirpath: null
  auto_insert_metric_name: true
  every_n_epochs: 1

trainer:
  # only use this for local apple silicon runs; change for cluster runs
  # accelerator: "mps"
  # devices: 1
  check_val_every_n_epoch: 1
  enable_checkpointing: true
  gradient_clip_val: null
  limit_train_batches: 1.0
  limit_test_batches: 1.0
  limit_val_batches: 1.0
  log_every_n_steps: 1
  max_epochs: 1
  min_epochs: 1

view_batch:
  enable: False
  num_frames: 0
  no_train: False