train_with_wav2vec2.yaml

# ############################################################################
# Model: Direct SLU
# Encoder: Pre-trained wav2vec2 -> LSTM
# Decoder: GRU + beamsearch
# Tokens: BPE with unigram
# losses: NLL
# Training: SLURP
# Authors:  Loren Lugosch, Mirco Ravanelli 2020 Boumadane Abdelmoumene, AbdelWahab Heba, Yingzhi Wang 2021
# For more wav2vec2/HuBERT results, please see https://arxiv.org/pdf/2111.02735.pdf
# ############################################################################

# Seed needs to be set at top of yaml, before objects with parameters are made
seed: 1986
__set_seed: !apply:speechbrain.utils.seed_everything [!ref <seed>]
output_folder: !ref results/<seed>
save_folder: !ref <output_folder>/save
train_log: !ref <output_folder>/train_log.txt
test_wer_file: !ref <output_folder>/wer_test_real.txt

# Data files
# The SLURP dataset will be automatically downloaded in the specified data_folder
data_folder: !<PLACEHOLDER> # e.g, /localscratch/SLURP
# data_folder_rirs: !ref <data_folder>
train_splits: ["train_synthetic", "train_real"]
csv_train: !ref <output_folder>/train-type=direct.csv
csv_valid: !ref <output_folder>/devel-type=direct.csv
csv_test: !ref <output_folder>/test-type=direct.csv
tokenizer_file: speechbrain/SLU-direct-SLURP-hubert-enc
skip_prep: False


# URL for the wav2vec2 model, you can change to benchmark different models
wav2vec2_hub: "facebook/hubert-base-ls960"

####################### Training Parameters ####################################
number_of_epochs: 35
batch_size: 6
lr: 0.0003
lr_wav2vec2: 0.00001
# token_type: unigram # ["unigram", "bpe", "char"]
sorting: random
ckpt_interval_minutes: 15 # save checkpoint every N min

#freeze all wav2vec2
freeze_wav2vec2: False

#set to true to freeze the CONV part of the wav2vec2 model
freeze_wav2vec2_conv: True

####################### Model Parameters #######################################
sample_rate: 16000
emb_size: 128
dec_neurons: 512
output_neurons: 58 # index(eos/bos) = 0
encoder_dim: 768

# Decoding parameters
bos_index: 0
eos_index: 0
min_decode_ratio: 0.0
max_decode_ratio: 10.0
slu_beam_size: 80
eos_threshold: 1.5
temperature: 1.25

dataloader_opts:
    batch_size: !ref <batch_size>
    shuffle: True

epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter
    limit: !ref <number_of_epochs>

# Models
wav2vec2: !new:speechbrain.lobes.models.huggingface_transformers.hubert.HuBERT
    source: !ref <wav2vec2_hub>
    output_norm: True
    freeze: !ref <freeze_wav2vec2>
    save_path: !ref <save_folder>/wav2vec2_checkpoint

output_emb: !new:speechbrain.nnet.embedding.Embedding
    num_embeddings: !ref <output_neurons>
    embedding_dim: !ref <emb_size>

dec: !new:speechbrain.nnet.RNN.AttentionalRNNDecoder
    enc_dim: !ref <encoder_dim>
    input_size: !ref <emb_size>
    rnn_type: lstm
    attn_type: content
    hidden_size: !ref <dec_neurons>
    attn_dim: 512
    num_layers: 3
    scaling: 1.0
    dropout: 0.0

seq_lin: !new:speechbrain.nnet.linear.Linear
    input_size: !ref <dec_neurons>
    n_neurons: !ref <output_neurons>

############################## Augmentations ###################################

# Speed perturbation
speed_perturb: !new:speechbrain.augment.time_domain.SpeedPerturb
    orig_freq: !ref <sample_rate>
    speeds: [95, 100, 105]

# Frequency drop: randomly drops a number of frequency bands to zero.
drop_freq: !new:speechbrain.augment.time_domain.DropFreq
    drop_freq_low: 0
    drop_freq_high: 1
    drop_freq_count_low: 1
    drop_freq_count_high: 3
    drop_freq_width: 0.05

# Time drop: randomly drops a number of temporal chunks.
drop_chunk: !new:speechbrain.augment.time_domain.DropChunk
    drop_length_low: 1000
    drop_length_high: 2000
    drop_count_low: 1
    drop_count_high: 5

# Augmenter: Combines previously defined augmentations to perform data augmentation
wav_augment: !new:speechbrain.augment.augmenter.Augmenter
    concat_original: True
    min_augmentations: 4
    max_augmentations: 4
    augment_prob: 1.0
    augmentations: [
        !ref <speed_perturb>,
        !ref <drop_freq>,
        !ref <drop_chunk>]

modules:
    wav2vec2: !ref <wav2vec2>
    output_emb: !ref <output_emb>
    dec: !ref <dec>
    seq_lin: !ref <seq_lin>

model: !new:torch.nn.ModuleList
    - [!ref <output_emb>,
       !ref <dec>, !ref <seq_lin>]

tokenizer: !new:sentencepiece.SentencePieceProcessor

pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
    collect_in: !ref <save_folder>/SLURP_tokenizer
    loadables:
        tokenizer: !ref <tokenizer>
    paths:
        tokenizer: !ref <tokenizer_file>/tokenizer_58_unigram.model

beam_searcher: !new:speechbrain.decoders.S2SRNNBeamSearcher
    embedding: !ref <output_emb>
    decoder: !ref <dec>
    linear: !ref <seq_lin>
    bos_index: !ref <bos_index>
    eos_index: !ref <eos_index>
    min_decode_ratio: !ref <min_decode_ratio>
    max_decode_ratio: !ref <max_decode_ratio>
    beam_size: !ref <slu_beam_size>
    eos_threshold: !ref <eos_threshold>
    temperature: !ref <temperature>
    using_max_attn_shift: False
    max_attn_shift: 30

opt_class: !name:torch.optim.Adam
    lr: !ref <lr>

wav2vec2_opt_class: !name:torch.optim.Adam
    lr: !ref <lr_wav2vec2>

lr_annealing: !new:speechbrain.nnet.schedulers.NewBobScheduler
    initial_value: !ref <lr>
    improvement_threshold: 0.0025
    annealing_factor: 0.8
    patient: 0

lr_annealing_wav2vec2: !new:speechbrain.nnet.schedulers.NewBobScheduler
    initial_value: !ref <lr_wav2vec2>
    improvement_threshold: 0.0025
    annealing_factor: 0.9

checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer
    checkpoints_dir: !ref <save_folder>
    recoverables:
        model: !ref <model>
        wav2vec2: !ref <wav2vec2>
        scheduler: !ref <lr_annealing>
        lr_annealing_wav2vec2: !ref <lr_annealing_wav2vec2>
        counter: !ref <epoch_counter>

log_softmax: !new:speechbrain.nnet.activations.Softmax
    apply_log: True

seq_cost: !name:speechbrain.nnet.losses.nll_loss
    label_smoothing: 0.1

train_logger: !new:speechbrain.utils.train_logger.FileTrainLogger
    save_file: !ref <train_log>

error_rate_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats

cer_computer: !name:speechbrain.utils.metric_stats.ErrorRateStats
    split_tokens: True