defaults:
  - _self_
  - model: null
  - effects: all

seed: 12345
train: True
sample_rate: 48000
chunk_size: 262144 # 5.5s
logs_dir: "./logs"
render_files: True
render_root: "./data"
accelerator: null
log_audio: True

# Effects
num_kept_effects: [2,2] # [min, max]
num_removed_effects: [2,2] # [min, max]
shuffle_kept_effects: True
shuffle_removed_effects: False
num_classes: 5
effects_to_keep:
  - reverb
  - chorus
  - delay
effects_to_remove:
  - compressor
  - distortion

callbacks:
  model_checkpoint:
    _target_: pytorch_lightning.callbacks.ModelCheckpoint
    monitor: "valid_loss"   # name of the logged metric which determines when model is improving
    save_top_k: 1           # save k best models (determined by above metric)
    save_last: True         # additionaly always save model from last epoch
    mode: "min"             # can be "max" or "min"
    verbose: False
    dirpath: ${logs_dir}/ckpts/${now:%Y-%m-%d-%H-%M-%S}
    filename: '{epoch:02d}-{valid_loss:.3f}'
  learning_rate_monitor:
    _target_: pytorch_lightning.callbacks.LearningRateMonitor
    logging_interval: "step"
  audio_logging:
    _target_: remfx.callbacks.AudioCallback
    sample_rate: ${sample_rate}
    log_audio: ${log_audio}

datamodule:
  _target_: remfx.datasets.EffectDatamodule
  train_dataset:
    _target_: remfx.datasets.EffectDataset
    total_chunks: 8000
    sample_rate: ${sample_rate}
    root: ${oc.env:DATASET_ROOT}
    chunk_size: ${chunk_size}
    mode: "train"
    effect_modules: ${effects}
    effects_to_keep: ${effects_to_keep}
    effects_to_remove: ${effects_to_remove}
    num_kept_effects: ${num_kept_effects}
    num_removed_effects: ${num_removed_effects}
    shuffle_kept_effects: ${shuffle_kept_effects}
    shuffle_removed_effects: ${shuffle_removed_effects}
    render_files: ${render_files}
    render_root: ${render_root}
    parallel: True
  val_dataset:
    _target_: remfx.datasets.EffectDataset
    total_chunks: 1000
    sample_rate: ${sample_rate}
    root: ${oc.env:DATASET_ROOT}
    chunk_size: ${chunk_size}
    mode: "val"
    effect_modules: ${effects}
    effects_to_keep: ${effects_to_keep}
    effects_to_remove: ${effects_to_remove}
    num_kept_effects: ${num_kept_effects}
    num_removed_effects: ${num_removed_effects}
    shuffle_kept_effects: ${shuffle_kept_effects}
    shuffle_removed_effects: ${shuffle_removed_effects}
    render_files: ${render_files}
    render_root: ${render_root}
  test_dataset:
    _target_: remfx.datasets.EffectDataset
    total_chunks: 1000
    sample_rate: ${sample_rate}
    root: ${oc.env:DATASET_ROOT}
    chunk_size: ${chunk_size}
    mode: "test"
    effect_modules: ${effects}
    effects_to_keep: ${effects_to_keep}
    effects_to_remove: ${effects_to_remove}
    num_kept_effects: ${num_kept_effects}
    num_removed_effects: ${num_removed_effects}
    shuffle_kept_effects: ${shuffle_kept_effects}
    shuffle_removed_effects: ${shuffle_removed_effects}
    render_files: ${render_files}
    render_root: ${render_root}

  train_batch_size: 16
  test_batch_size: 1
  num_workers: 8
  pin_memory: True
  persistent_workers: True

# logger:
#   _target_: pytorch_lightning.loggers.WandbLogger
#   project: ${oc.env:WANDB_PROJECT}
#   entity: ${oc.env:WANDB_ENTITY}
#   # offline: False  # set True to store all logs only locally
#   job_type: "train"
#   group: ""
#   save_dir: "."
#   log_model: True
logger:
  _target_: pytorch_lightning.loggers.CSVLogger
  save_dir: "."

trainer:
  _target_: pytorch_lightning.Trainer
  precision: 32 # Precision used for tensors, default `32`
  min_epochs: 0
  max_epochs: -1
  log_every_n_steps: 1 # Logs metrics every N batches
  accumulate_grad_batches: 1
  accelerator: ${accelerator}
  devices: 1
  gradient_clip_val: 10.0
  max_steps: 50000