config.yaml

defaults:
  - _self_
  - dset: musdb44
  - svd: default
  - variant: default
  - override hydra/hydra_logging: colorlog
  - override hydra/job_logging: colorlog

dummy:
dset:
  musdb: /checkpoint/defossez/datasets/musdbhq
  musdb_samplerate: 44100
  wav:  # path to custom wav dataset
  segment: 11
  shift: 1
  train_valid: false
  full_cv: true
  samplerate: 44100
  channels: 2
  normalize: true
  metadata: ./metadata
  sources: ['drums', 'bass', 'other', 'vocals']

test:
  save: False
  best: True
  workers: 2
  every: 20
  split: true
  shifts: 1
  overlap: 0.25
  sdr: true
  metric: 'loss'  # metric used for best model selection on the valid set, can also be nsdr
  nonhq:   # path to non hq MusDB for evaluation

epochs: 360
batch_size: 64
max_batches:  # limit the number of batches per epoch, useful for debugging
              # or if your dataset is gigantic.
optim:
  lr: 3e-4
  momentum: 0.9
  beta2: 0.999
  loss: l1    # l1 or mse
  optim: adam
  weight_decay: 0
  clip_grad: 0

seed: 42
debug: false
valid_apply: true
flag:
save_every:
weights: [1., 1., 1., 1.]  # weights over each source for the training/valid loss.

augment:
  shift_same: false
  repitch:
    proba: 0.2
    max_tempo: 12
  remix:
    proba: 1
    group_size: 4
  scale:
    proba: 1
    min: 0.25
    max: 1.25
  flip: true

continue_from:  # continue from other XP, give the XP Dora signature.
continue_pretrained:   # signature of a pretrained XP, this cannot be a bag of models.
pretrained_repo:   # repo for pretrained model (default is official AWS)
continue_best: true
continue_opt: false

misc:
  num_workers: 10
  num_prints: 4
  show: false
  verbose: false

# List of decay for EMA at batch or epoch level, e.g. 0.999.
# Batch level EMA are kept on GPU for speed.
ema:
  epoch: []
  batch: []

model_segment:  # override the segment parameter for the model, usually 4 times the training segment.
model: demucs
demucs:  # see demucs/demucs.py for a detailed description
  # Channels
  channels: 64
  growth: 2
  # Main structure
  depth: 6
  rewrite: true
  lstm_layers: 0
  # Convolutions
  kernel_size: 8
  stride: 4
  context: 1
  # Activations
  gelu: true
  glu: true
  # Normalization
  norm_groups: 4
  norm_starts: 4
  # DConv residual branch
  dconv_depth: 2
  dconv_mode: 1  # 1 = branch in encoder, 2 = in decoder, 3 = in both.
  dconv_comp: 4
  dconv_attn: 4
  dconv_lstm: 4
  dconv_init: 1e-4
  # Pre/post treatment
  resample: true
  normalize: false
  # Weight init
  rescale: 0.1

hdemucs:  # see demucs/hdemucs.py for a detailed description
  # Channels
  channels: 48
  channels_time:
  growth: 2
  # STFT
  nfft: 4096
  wiener_iters: 0
  end_iters: 0
  wiener_residual: false
  cac: true
  # Main structure
  depth: 6
  rewrite: true
  hybrid: true
  hybrid_old: false
  # Frequency Branch
  multi_freqs: []
  multi_freqs_depth: 3
  freq_emb: 0.2
  emb_scale: 10
  emb_smooth: true
  # Convolutions
  kernel_size: 8
  stride: 4
  time_stride: 2
  context: 1
  context_enc: 0
  # normalization
  norm_starts: 4
  norm_groups: 4
  # DConv residual branch
  dconv_mode: 1
  dconv_depth: 2
  dconv_comp: 4
  dconv_attn: 4
  dconv_lstm: 4
  dconv_init: 1e-3
  # Weight init
  rescale: 0.1


svd:  # see svd.py for documentation
  penalty: 0
  min_size: 0.1
  dim: 1
  niters: 2
  powm: false
  proba: 1
  conv_only: false
  convtr: false
  bs: 1

quant:  # quantization hyper params
  diffq:    # diffq penalty, typically 1e-4 or 3e-4
  qat:      # use QAT with a fixed number of bits (not as good as diffq)
  min_size: 0.2
  group_size: 8

dora:
  dir: outputs
  exclude: ["misc.*", "slurm.*", 'test.reval', 'flag']

slurm:
  time: 4320
  constraint: volta32gb
  setup: ['module load cuda/11.0 cudnn/v8.0.3.33-cuda.11.0 NCCL/2.8.3-1-cuda.11.0']

# Hydra config
hydra:
  job_logging:
    formatters:
      colorlog:
        datefmt: "%m-%d %H:%M:%S"