pretrain_wav2vec2.py

import os
import shutil
import numpy as np
import argparse
import torch
from torch.utils.data.dataloader import DataLoader
from accelerate import Accelerator
from tqdm import tqdm
from transformers import get_scheduler, set_seed
from dataset import LibriSpeechDataset, Wav2Vec2CollateFunctionForPreTraining
from model import Wav2Vec2ForPreTraining
from utils import Wav2Vec2Config

def parse_args():
    parser = argparse.ArgumentParser(description="Wav2Vec2 Pretraining Arguments on Librispeech")
    parser.add_argument(
        "--experiment_name", 
        required=True, 
        type=str
    )
    parser.add_argument(
        "--working_directory", 
        required=True, 
        type=str
    )
    parser.add_argument(
        "--path_to_data_root", 
        help="Path to data root directory",
        required=True,
        type=str
    )
    parser.add_argument(
        "--train_splits", 
        help="Select Librispeech Training Splits (960 Hours of data Split into 100, 360 and 500 hour sections)",
        default=["train-clean-100"],
        choices=("train-clean-100", "train-clean-360", "train-other-500", "dev-clean", "test-clean"),
        nargs='+',
        type=str
    )
    parser.add_argument(
        "--test_splits", 
        help="Select Librispeech Testing Splits (Select from limited Validation + Testing Datasets)",
        default=["test-clean"],
        choices=("train-clean-100", "train-clean-360", "train-other-500", "dev-clean", "test-clean"),
        nargs="+",
        type=str
    )
    parser.add_argument(
        "--minimum_audio_duration",
        help="Filter out any audio sample less than these many seconds", 
        default=2.0, 
        type=float
    )
    parser.add_argument(
        "--maximum_audio_duration",
        help="Filter out any audio samples greater than these many seconds",
        default=20.0, 
        type=float
    )
    parser.add_argument(
        "--sampling_rate",
        help="Sampling rate to sample all audio to before passing to model",
        default=16000,
        type=int
    )
    parser.add_argument(
        "--audio_input_channels",
        help="Number of input channels from audio",
        default=1, 
        type=int
    )
    parser.add_argument(
        "--masking_probability",
        help="The probability for each token to be the start of a span mask",
        default=0.065, 
        type=float
    )
    parser.add_argument(
        "--masking_span_length",
        help="Number of consecutive tokens to mask in every span",
        default=10, 
        type=int
    )
    parser.add_argument(
        "--minimum_spans",
        help="Minimum number of span masks to have in every sequence",
        default=2, 
        type=int
    )
    parser.add_argument(
        "--num_negatives",
        help="For every masked token, how many negatives do we want to sample for contrastive loss?",
        default=100, 
        type=int
    )
    parser.add_argument(
        "--num_workers",
        help="Number of workers for dataloading",
        default=16, 
        type=int
    )
    parser.add_argument(
        "--conv_dim",
        help="Sequence of channels dims in encoder convolutions",
        default=(512, 512, 512, 512, 512, 512, 512),
        nargs="+",
        type=int
    )
    parser.add_argument(
        "--conv_kernel",
        help="Kernel size for each convolution in encoder convolutions",
        default=(10, 3, 3, 3, 3, 2, 2),
        nargs="+",
        type=int
    )
    parser.add_argument(
        "--conv_stride",
        help="Strides for each convolution in encoder convolutions",
        default=(5, 2, 2, 2, 2, 2, 2),
        nargs="+",
        type=int
    )
    parser.add_argument(
        "--disable_conv_bias",
        help="Do you want to have bias on the convolutional encoder",
        action=argparse.BooleanOptionalAction,
        default=False
    )
    parser.add_argument(
        "--feature_proj_dropout_p",
        help="Dropout on feature projection from convolution to transformer embedding dims",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--conv_positional_emb_drop_p",
        help="Positional Embedding Dropout probability",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--conv_positional_emb_groups",
        help="Number of groups in convolution for positional encodings",
        default=16,
        type=int
    )
    parser.add_argument(
        "--conv_positional_emb_kernel_size",
        help="Kernel size for convolutional positional encodings",
        default=128,
        type=int
    )
    parser.add_argument(
        "--num_transformer_layers",
        help="Number of transformer blocks in model",
        default=12,
        type=int
    )
    parser.add_argument(
        "--num_attention_heads",
        help="Number of heads of attention",
        default=12,
        type=int
    )
    parser.add_argument(
        "--embedding_dimension",
        help="Transformer embedding dimension",
        default=768,
        type=int
    )
    parser.add_argument(
        "--mlp_ratio",
        help="Hidden layer expansion factor for feed forward layers",
        default=4,
        type=int
    )
    parser.add_argument(
        "--mlp_dropout_p",
        help="Dropout probability in feedforward layers",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--attention_dropout_p",
        help="Dropout probability on attention matrix",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--transformer_encoder_dropout_p",
        help="Post transformer block dropout probability",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--layer_dropout",
        help="Entire transformer layerblock dropout (https://paperswithcode.com/method/layerdrop). \
            though im not sure how to implement this with DDP, as it throws unused parameters error",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--initializer_range",
        help="Standard deviation of linear layers initialized as normal distribution",
        default=0.02,
        type=float
    )
    parser.add_argument(
        "--num_codevector_groups",
        help="Number of codebooks in our quantizer",
        default=2,
        type=int
    )
    parser.add_argument(
        "--num_codevectors_per_group",
        help="Number of codevectors per group",
        default=320,
        type=int
    )
    parser.add_argument(
        "--codevector_dim",
        help="Dimension of codevectors in vector quantization",
        default=256,
        type=int
    )
    parser.add_argument(
        "--pre_quantizer_dropout_p",
        help="Dropout before quantization of tokens",
        default=0.0,
        type=float
    )
    parser.add_argument(
        "--max_gumbel_temperature",
        type=float,
        default=2.0,
        help="Maximum temperature for gumbel softmax.",
    )
    parser.add_argument(
        "--min_gumbel_temperature",
        type=float,
        default=0.5,
        help="Minimum temperature for gumbel softmax.",
    )
    parser.add_argument(
        "--gumbel_temperature_decay", 
        type=float, 
        default=0.999995, 
        help="Decay of gumbel temperature during training."
    )
    parser.add_argument(
        "--contrastive_logits_temperature",
        help="Temperature to scale cosine similarity before softmax",
        default=0.1, 
        type=float
    )
    parser.add_argument(
        "--diversity_loss_weight",
        help="Weight to scale diversity loss",
        default=0.1, 
        type=float
    )
    parser.add_argument(
        "--per_gpu_batch_size",
        help="Overall batch size per gpu during training",
        default=64, 
        type=int
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        help="Splits per_gpu_batch_size by gradient_accumulation_steps",
        default=8, 
        type=int
    )
    parser.add_argument(
        "--num_training_steps", 
        help="Number of training steps to take",
        default=200000,
        type=int
    )
    parser.add_argument(
        "--num_warmup_steps", 
        type=int, 
        default=32000, 
        help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument(
        "--lr_scheduler_type",
        type=str,
        default="polynomial",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )
    parser.add_argument(
        "--logging_steps", 
        help="Number of iterations for every log of metrics to wandb",
        default=1,
        type=int
    )
    parser.add_argument(
        "--evaluation_interval", 
        help="Number of iterations for every evaluation and plotting",
        default=1000, 
        type=int
    )
    parser.add_argument(
        "--checkpoint_interval",
        help="Number of iterations for checkpointing",
        default=1000,
        type=int
    )
    parser.add_argument(
        "--learning_rate", 
        help="Max learning rate for all Learning Rate Schedulers", 
        default=0.001, 
        type=float
    )
    parser.add_argument(
        "--bias_weight_decay", 
        help="Apply weight decay to bias",
        default=False, 
        action=argparse.BooleanOptionalAction
    )
    parser.add_argument(
        "--norm_weight_decay",
        help="Apply weight decay to normalization weight and bias",
        default=False,
        action=argparse.BooleanOptionalAction
    )
    parser.add_argument(
        "--weight_decay",
        help="Weight decay constant for AdamW optimizer", 
        default=0.01, 
        type=float
    )
    parser.add_argument(
        "--adam_beta1",
        type=float,
        default=0.9,
        help="Beta1 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_beta2",
        type=float,
        default=0.98,
        help="Beta2 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_epsilon",
        type=float,
        default=1e-6,
        help="Epsilon for AdamW optimizer",
    )
    parser.add_argument(
        "--num_keep_checkpoints",
        help="Number of Checkpoints to Keep, if None, all checkpoints will be saved",
        default=None, 
        type=int
    )
    parser.add_argument(
        "--seed", 
        help="Set seed in model for reproducible training",
        default=None,
        type=int
    )
    parser.add_argument(
        "--resume_from_checkpoint", 
        help="Checkpoint folder for model to resume training from, inside the experiment folder", 
        default=None, 
        type=str
    )   
    parser.add_argument(
        "--log_wandb", 
        help="Flag to enable logging to wandb",
        default=False, 
        action=argparse.BooleanOptionalAction
    )
    args = parser.parse_args()
    return args

def multiply_gradients(params,constant):
    for param in params:
        if param.grad is not None:
            param.grad.data.mul_(constant)

def compute_gradient_norms(params,scale=1):
    total_norm=0
    for p in params:
        if p.grad is not None:
            param_norm=(p.grad.detach().data/scale).norm(2)
            total_norm+=param_norm.item()**2
    total_norm=total_norm**0.5
    return total_norm

def compute_batch_duration(attention_mask,sampling_rate):
    total_duration_seconds=torch.sum(attention_mask.sum(axis=-1)/sampling_rate)
    total_duration_hours=total_duration_seconds/3600
    return total_duration_hours

args=parse_args()

if args.seed is not None:
    set_seed(args.seed)

path_to_experiment=os.path.join(args.working_directory,args.experiment_name)
accelerator=Accelerator(project_dir=path_to_experiment,log_with="wandb" if args.log_wandb else None)
if args.log_wandb:
    accelerator.init_trackers(args.experiment_name)

config = Wav2Vec2Config(
    conv_dim=tuple(args.conv_dim), 
    conv_stride=tuple(args.conv_stride),
    conv_kernel=tuple(args.conv_kernel),
    conv_bias=not args.disable_conv_bias, 
    feature_projection_dropout_p=args.feature_proj_dropout_p,
    conv_positional_emb_drop_p=args.conv_positional_emb_drop_p, 
    conv_positional_emb_groups=args.conv_positional_emb_groups,
    conv_positional_emb_kernel_size=args.conv_positional_emb_kernel_size,
    num_transformer_layers=args.num_transformer_layers,
    num_attention_heads=args.num_attention_heads,
    mlp_ratio=args.mlp_ratio,
    mlp_dropout_p=args.mlp_dropout_p, 
    attention_dropout_p=args.attention_dropout_p,
    transformer_encoder_dropout=args.transformer_encoder_dropout_p,
    layer_dropout=args.layer_dropout if accelerator.num_processes == 1 else 0.0,
    initializer_range=args.initializer_range,
    num_codevector_groups=args.num_codevector_groups,
    num_codevectors_per_group=args.num_codevectors_per_group,
    codevector_dim=args.codevector_dim, 
    pre_quantizer_dropout=args.pre_quantizer_dropout_p, 
    masking_probability=args.masking_probability, 
    masking_span_length=args.masking_span_length, 
    minimum_spans=args.minimum_spans, 
    contrastive_logits_temperature=args.contrastive_logits_temperature, 
    diversity_loss_weight=args.diversity_loss_weight,
    num_negatives=args.num_negatives
)

model = Wav2Vec2ForPreTraining(config)
model_parameters = filter(lambda p: p.requires_grad, model.parameters())
params = sum([np.prod(p.size()) for p in model_parameters])
accelerator.print("Number of Parameters:", params)
train_set = LibriSpeechDataset(path_to_data_root=args.path_to_data_root, 
                               include_splits=args.train_splits,
                               max_audio_duration=args.maximum_audio_duration, 
                               min_audio_duration=args.minimum_audio_duration,
                               sampling_rate=args.sampling_rate,
                               return_transcripts=False)
test_set = LibriSpeechDataset(path_to_data_root=args.path_to_data_root, 
                              include_splits=args.test_splits,
                              max_audio_duration=args.maximum_audio_duration, 
                              min_audio_duration=args.minimum_audio_duration,
                              sampling_rate=args.sampling_rate,
                              return_transcripts=False)
data_collator = Wav2Vec2CollateFunctionForPreTraining(config)
minibatch_size = args.per_gpu_batch_size // args.gradient_accumulation_steps
train_dataloader = DataLoader(train_set, 
                              batch_size=minibatch_size, 
                              shuffle=False, 
                              num_workers=8, 
                              collate_fn=data_collator)
eval_dataloader = DataLoader(test_set, 
                             batch_size=minibatch_size, 
                             shuffle=False, 
                             num_workers=8, 
                             collate_fn=data_collator)
if (not args.bias_weight_decay) or (not args.norm_weight_decay):
    accelerator.print("Disabling Weight Decay on Some Parameters")
    weight_decay_params = []
    no_weight_decay_params = []
    for name, param in model.named_parameters():

        if param.requires_grad:
            if "bias" in name and not args.bias_weight_decay:
                no_weight_decay_params.append(param)
            elif "groupnorm" in name and not args.norm_weight_decay:
                no_weight_decay_params.append(param)
            else:
                weight_decay_params.append(param)
    optimizer_group = [
        {"params": weight_decay_params, "weight_decay": args.weight_decay},
        {"params": no_weight_decay_params, "weight_decay": 0.0}
    ]
    optimizer = torch.optim.AdamW(optimizer_group, 
                                  lr=args.learning_rate,
                                  betas=[args.adam_beta1, args.adam_beta2],
                                  eps=args.adam_epsilon)
else:
    optimizer = torch.optim.AdamW(model.parameters(), 
                                lr=args.learning_rate,
                                betas=[args.adam_beta1, args.adam_beta2],
                                eps=args.adam_epsilon)

scheduler = get_scheduler(
    name=args.lr_scheduler_type,
    optimizer=optimizer,
    num_warmup_steps=args.num_warmup_steps,
    num_training_steps=args.num_training_steps,
)

model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)
accelerator.register_for_checkpointing(scheduler)

if args.resume_from_checkpoint is not None:

    path_to_checkpoint = os.path.join(path_to_experiment, args.resume_from_checkpoint)
    
    with accelerator.main_process_first():
        accelerator.load_state(path_to_checkpoint)
    
    completed_steps = int(args.resume_from_checkpoint.split("_")[-1])
    accelerator.print(f"Resuming from Iteration: {completed_steps}")
else:
    completed_steps = 0

train = True
progress_bar = tqdm(range(completed_steps, args.num_training_steps), disable=not accelerator.is_local_main_process)

while train:

    accumulate_steps = 0

    accumulated_hours_per_batch = 0

    accumulated_percent_masked = 0
    
    for batch in train_dataloader:        

        num_losses = batch["mask_time_indices"].sum()
        
        hours_of_audio = compute_batch_duration(batch["attention_mask"], args.sampling_rate)
        accumulated_hours_per_batch += hours_of_audio

        percent_masked = num_losses / batch["sub_attention_mask"].sum()
        accumulated_percent_masked += percent_masked / args.gradient_accumulation_steps

        batch = {k:v.to(accelerator.device) for (k,v) in batch.items()}

        outputs = model(**batch)

        loss = outputs.loss / args.gradient_accumulation_steps

        accelerator.backward(loss)

        if accelerator.state.num_processes > 1:
            num_losses = accelerator.gather_for_metrics(num_losses).sum()
            gradient_multiplier = accelerator.state.num_processes / num_losses
            multiply_gradients(model.module.parameters(), gradient_multiplier)
        else:
            multiply_gradients(model.parameters(), 1 / num_losses)

        accumulate_steps += 1

        if accumulate_steps % args.gradient_accumulation_steps == 0:

            if hasattr(accelerator, "scaler") and accelerator.scaler is not None:
                scale = accelerator.scaler._scale.item()
            else:
                scale = 1

            if accelerator.state.num_processes > 1:
                grad_norm = compute_gradient_norms(model.module.parameters(), scale)
            else:
                grad_norm = compute_gradient_norms(model.parameters(), scale)

            optimizer.step()
            optimizer.zero_grad()
            scheduler.step()

            gumbel_temperature = max(args.max_gumbel_temperature * args.gumbel_temperature_decay**completed_steps,
                                     args.min_gumbel_temperature)
            
            if hasattr(model, "module"):
                model.module.set_gumbel_temperature(gumbel_temperature)
            else:
                model.set_gumbel_temperature(gumbel_temperature)

            if completed_steps % args.logging_steps == 0:
                
                loss = outputs.loss.detach()
                contrastive_loss = outputs.contrastive_loss.detach()
                diversity_loss = outputs.diversity_loss.detach()
                perplexity = outputs.codevector_perplexity.detach()
                hours_per_batch = accumulated_hours_per_batch.detach()
                percent_masked = accumulated_percent_masked.detach()

                if accelerator.state.num_processes > 1:

                    loss = torch.sum(accelerator.gather_for_metrics(loss)) / num_losses
                    contrastive_loss = torch.sum(accelerator.gather_for_metrics(contrastive_loss)) / num_losses
                    diversity_loss = torch.sum(accelerator.gather_for_metrics(diversity_loss)) / num_losses
                    perplexity = torch.sum(accelerator.gather_for_metrics(perplexity)) / num_losses
                    hours_per_batch = torch.sum(accelerator.gather_for_metrics(hours_per_batch))
                    percent_masked = torch.mean(accelerator.gather_for_metrics(percent_masked))
                
                else:
                    loss = loss / num_losses
                    contrastive_loss = contrastive_loss / num_losses
                    diversity_loss = diversity_loss / num_losses
                    perplexity = perplexity / num_losses
                    hours_per_batch = hours_per_batch
                    percent_masked = percent_masked

                log = {"train_loss": loss,
                       "train_contrast_loss": contrastive_loss,
                       "train_div_loss": diversity_loss,
                       "pct_masked": percent_masked,
                       "batch_hours": hours_per_batch, 
                       "perplexity": perplexity,
                       "lr": scheduler.get_last_lr()[0],
                       "temp": gumbel_temperature,
                       "grad_norm": grad_norm}

                logging_string = ""
                for k, v in log.items():
                    logging_string += f"|{k[6:] if 'train_' in k else k}: {round(v.item() if torch.is_tensor(v) else v, 3)}"
                
                if accelerator.is_main_process:
                    progress_bar.write(logging_string)
                
                if args.log_wandb:
                    accelerator.log(log, step=completed_steps)

            if completed_steps % args.evaluation_interval == 0:
                if accelerator.is_main_process:
                    progress_bar.write("Evaluating Model!!")
                
                model.eval()

                log = {"val_loss": 0, 
                       "val_contrast_loss": 0, 
                       "val_div_loss": 0}
                all_num_losses = 0

                for batch in tqdm(eval_dataloader, disable=not accelerator.is_main_process):
                    
                    num_losses = batch["mask_time_indices"].sum()

                    batch = {k:v.to(accelerator.device) for (k,v) in batch.items()}

                    with torch.inference_mode():
                        output = model(**batch)


                    loss = output.loss
                    contrastive_loss = output.contrastive_loss
                    diversity_loss = outputs.diversity_loss

                    if accelerator.num_processes > 1:
                        loss = torch.sum(accelerator.gather_for_metrics(loss))
                        contrastive_loss = torch.sum(accelerator.gather_for_metrics(contrastive_loss))
                        diversity_loss = torch.sum(accelerator.gather_for_metrics(diversity_loss))
                        num_losses = torch.sum(accelerator.gather_for_metrics(num_losses))

                    log["val_loss"] += loss
                    log["val_contrast_loss"] += contrastive_loss
                    log["val_div_loss"] += diversity_loss 
                    all_num_losses += num_losses
                
                log = {k: v / all_num_losses for (k,v) in log.items()}

                logging_string = ""
                for k, v in log.items():
                    logging_string += f"|{k[4:]}: {round(v.item() if torch.is_tensor(v) else v, 3)}"

                if accelerator.is_main_process:
                    progress_bar.write(logging_string)
                
                if args.log_wandb:
                    accelerator.log(log, step=completed_steps)

                model.train()

            if (completed_steps % args.checkpoint_interval == 0):
                
                path_to_checkpoint = os.path.join(path_to_experiment, f"checkpoint_{completed_steps}")

                if accelerator.is_main_process:
                    progress_bar.write(f"Saving Checkpoint to {path_to_checkpoint}")

                accelerator.wait_for_everyone()

                if accelerator.is_main_process:
                    accelerator.save_state(output_dir=path_to_checkpoint)

                if args.num_keep_checkpoints is not None:
                    if accelerator.is_main_process:
                        all_checkpoints = os.listdir(path_to_experiment)
                        all_checkpoints = sorted(all_checkpoints, key=lambda x: int(x.split(".")[0].split("_")[-1]))
                        
                        if len(all_checkpoints) > args.num_keep_checkpoints:
                            checkpoints_to_delete = all_checkpoints[:-args.num_keep_checkpoints]

                            for checkpoint_to_delete in checkpoints_to_delete:
                                path_to_checkpoint_to_delete = os.path.join(path_to_experiment, checkpoint_to_delete)
                                if os.path.isdir(path_to_checkpoint_to_delete):
                                    shutil.rmtree(path_to_checkpoint_to_delete)
                                    
                accelerator.wait_for_everyone()
                
            if completed_steps >= args.num_training_steps:
                train = False
                if accelerator.is_main_process:
                    progress_bar.write("Completed Training!!")
                break

            completed_steps += 1
            progress_bar.update(1)
            accumulate_steps = accumulated_hours_per_batch = accumulated_percent_masked = 0

accelerator.end_training()