config.yaml

seed_everything: 3407

data:
  class_path: unicodec.decoder.dataset.VocosDataModule
  init_args:
    train_params:
      filelist_path: ./data/train/finetune_data
      sampling_rate: 24000
      num_samples: 240000
      batch_size: 10  #18
      num_workers: 8

    val_params:
      filelist_path: ./data/infer/large_data_domain
      sampling_rate: 24000
      num_samples: 240000
      batch_size: 5   # 10
      num_workers: 8

model:
  class_path: unicodec.decoder.experiment.VocosEncodecExp
  init_args:
    sample_rate: 24000
    initial_learning_rate: 5e-5
    mel_loss_coeff: 450
    mrd_loss_coeff: 1.0
    # ctr_loss_coeff: 0.001
    num_warmup_steps: 5000 # Optimizers warmup steps
    pretrain_mel_steps: 0  # 0 means GAN objective from the first iteration
    use_ema: false

    # automatic evaluation
    evaluate_utmos: true
    evaluate_pesq: true
    evaluate_periodicty: true

    resume: true
    resume_config: 
    resume_model: 
    feature_extractor:
      class_path: unicodec.decoder.feature_extractors.EncodecFeatures
      init_args:
        encodec_model: encodec_24khz
        bandwidths: [6.6, 6.6, 6.6, 6.6]
        train_codebooks: true
        num_quantizers: 1
        dowmsamples: [8, 5, 4, 2]
        vq_bins: 16384
        vq_kmeans: 200
        use_transformer: true
        mask: false

    backbone:
      class_path: unicodec.decoder.models.VocosBackbone
      init_args:
        input_channels: 512
        dim: 768
        intermediate_dim: 2304
        num_layers: 12
        adanorm_num_embeddings: 4  # len(bandwidths)

    head:
      class_path: unicodec.decoder.heads.ISTFTHead
      init_args:
        dim: 768
        n_fft: 1280  #4*hop_length
        hop_length: 320  # 8*5*4*2
        padding: same

trainer:
  logger:
    class_path: pytorch_lightning.loggers.TensorBoardLogger
    init_args:
      save_dir: /debug/
  callbacks:
    - class_path: pytorch_lightning.callbacks.LearningRateMonitor
    - class_path: pytorch_lightning.callbacks.ModelSummary
      init_args:
        max_depth: 2
    - class_path: pytorch_lightning.callbacks.ModelCheckpoint
      init_args:
        monitor: val_loss
        filename: vocos_checkpoint_{epoch}_{step}_{val_loss:.4f}
        save_top_k: 100
        save_last: true
        # every_n_train_steps: 5000
    - class_path: unicodec.decoder.helpers.GradNormCallback

  # Lightning calculates max_steps across all optimizer steps (rather than number of batches)
  # This equals to 1M steps per generator and 1M per discriminator
  max_steps: 20000000
  # You might want to limit val batches when evaluating all the metrics, as they are time-consuming
  limit_val_batches: 100
  accelerator: gpu
  strategy: ddp
  devices: [0,1,2,3,4,5,6,7]
  num_nodes: 4
  log_every_n_steps: 200
  # val_check_interval: 5000