import os import json import glob import math import time import wandb import torch import random import auraloss import itertools import torchaudio import numpy as np import pandas as pd import torch.nn as nn import matplotlib.pyplot as plt import torch.distributed as dist import torch.nn.functional as F from tqdm import tqdm from torch.cuda.amp import autocast from typing import Tuple, Dict, List, Set, Any from torch.utils.data import DistributedSampler from torch.nn.parallel import DistributedDataParallel from torch.optim.lr_scheduler import LinearLR, ChainedScheduler from sklearn.metrics import f1_score, roc_auc_score, average_precision_score def seed_worker(worker_id): """Function to be called by each DataLoader worker.""" # Get base seed from worker_info worker_info = torch.utils.data.get_worker_info() base_seed = worker_info.seed np_seed = int(base_seed) % (2**32 - 1) # Set seeds for each worker using its unique base_seed random.seed(base_seed) torch.manual_seed(base_seed) np.random.seed(np_seed) def setup_seeds(): """Set up seeds for distributed training.""" # Get rank for this process # rank = dist.get_rank() # Get local rank (GPU id for this process) local_rank = int(os.environ["LOCAL_RANK"]) # Get world size (total number of processes) world_size = dist.get_world_size() # Create a base seed using rank base_seed = 42 # Your chosen base seed process_seed = int(base_seed + local_rank) # ensure seed is in range 0-2^32-1 process_seed = process_seed % (2**32 - 1) print(f"Setting seed for process {local_rank} to {process_seed}") # Set seeds for this process random.seed(process_seed) torch.manual_seed(process_seed) np.random.seed(process_seed) # If using CUDA, set its seeds too if torch.cuda.is_available(): torch.cuda.manual_seed(process_seed) torch.cuda.manual_seed_all(process_seed) return process_seed def corrupt_audio( audio, sample_rate, highpass_prob=0.1, lowpass_prob=0.1, noise_prob=0.1, tanh_prob=0.1, clip_prob=0.1, preemphasis_prob=0.05, deemphasis_prob=0.05, bass_boost_prob=0.1, treble_boost_prob=0.1, mp3_prob=0.90, ): if np.random.uniform() < highpass_prob: # highpass # sample on a log scale freq_hz = 10 ** np.random.uniform(np.log10(20), np.log10(20000)) audio = torchaudio.functional.highpass_biquad(audio, sample_rate, freq_hz) if np.random.uniform() < lowpass_prob: # lowpass freq_hz = 10 ** np.random.uniform(np.log10(20), np.log10(20000)) audio = torchaudio.functional.lowpass_biquad(audio, sample_rate, freq_hz) if np.random.uniform() < bass_boost_prob: # bass boost gain_db = np.random.uniform(6, 12) freq_hz = np.random.uniform(20, 240) audio = torchaudio.functional.bass_biquad(audio, sample_rate, gain_db, freq_hz) if np.random.uniform() < treble_boost_prob: # treble boost gain_db = np.random.uniform(6, 12) freq_hz = np.random.uniform(1000, 10000) audio = torchaudio.functional.treble_biquad( audio, sample_rate, gain_db, freq_hz ) if np.random.uniform() < noise_prob: # noise noise_gain_db = np.random.uniform(-48, -24) audio = audio + torch.randn_like(audio) * 10 ** (noise_gain_db / 20.0) if np.random.uniform() < tanh_prob: # tanh gain_db = np.random.uniform(12, 24) audio = torch.tanh(audio * 10 ** (gain_db / 20.0)) if np.random.uniform() < clip_prob: # clip gain_db = np.random.uniform(12, 24) audio = torch.clamp(audio * 10 ** (gain_db / 20.0), -1, 1) if np.random.uniform() < preemphasis_prob: # preemphasis coeff = np.random.uniform(0.75, 1.0) audio = torchaudio.functional.preemphasis(audio, coeff) if np.random.uniform() < deemphasis_prob: # deemphasis coeff = np.random.uniform(0.75, 1.0) audio = torchaudio.functional.deemphasis(audio, coeff) if np.random.uniform() < mp3_prob: # mp3 bit_rate = np.random.choice( [ 8000, 16000, 24000, 32000, 48000, 64000, 96000, 112000, 128000, ] ) effector = torchaudio.io.AudioEffector( format="mp3", codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate), ) audio = effector.apply(audio.T, sample_rate).T return torch.clamp(audio, -1.0, 1.0) def pad_semantic(semantic, chunk_size, pad_token, extra=0): if semantic.shape[0] < chunk_size: pad_len = chunk_size - semantic.shape[0] + extra semantic = torch.cat( [semantic, torch.full((pad_len,), pad_token, dtype=torch.long)], ) return semantic class SemanticPairsDataset(torch.utils.data.Dataset): def __init__( self, data_dir: str, item_ids: List[str], chunk_size: int = 750, # 30s at 25hz pad_token: int = 4000, mask_prob: float = 0.0, ): self.data_dir = data_dir self.item_ids = item_ids self.chunk_size = chunk_size self.pad_token = pad_token self.mask_prob = mask_prob assert len(self.item_ids) > 0, "No item ids found" # load all the data into memory self.examples = [] for item_id in tqdm(self.item_ids): orig_sem_path = os.path.join( self.data_dir, item_id, f"{item_id}_original_semantic.npz" ) gen_sem_path = os.path.join( self.data_dir, item_id, f"{item_id}_generated_semantic.npz" ) orig_sem = torch.from_numpy(np.load(orig_sem_path)["semantic_codes"]).long() gen_sem = torch.from_numpy(np.load(gen_sem_path)["semantic_codes"]).long() # if the generated semantic is shorter than the chunk_size, pad it with pad token if gen_sem.shape[0] <= self.chunk_size: continue if orig_sem.shape[0] <= self.chunk_size: continue # store the original and generated semantic codes self.examples.append((orig_sem, gen_sem)) print(f"Loaded {len(self.examples)} examples") def __len__(self): return len(self.examples) def __getitem__(self, idx): orig_sem, gen_sem = self.examples[idx] # Ensure both sequences are at least chunk_size + 1 to avoid issues with random sampling min_len = min(orig_sem.shape[0], gen_sem.shape[0]) # now we need to construct a pair of chunks # we have a few assumptions to construct these chunks # first, we assume that original is always better than generated # we can sample random chunks from the original and generated # second, we assume that early chunks are better than later chunks in generations dice_roll = np.random.rand() if dice_roll < 0.5 and min_len >= self.chunk_size * 2: # compare early and late chunks from generated max_pos_start = min_len - 2 * self.chunk_size pos_start_idx = np.random.randint(0, max_pos_start + 1) pos_end_idx = pos_start_idx + self.chunk_size min_neg_start = pos_end_idx max_neg_start = min_len - self.chunk_size neg_start_idx = np.random.randint(min_neg_start, max_neg_start + 1) neg_end_idx = neg_start_idx + self.chunk_size pos_chunk = gen_sem[pos_start_idx:pos_end_idx] neg_chunk = gen_sem[neg_start_idx:neg_end_idx] else: # compare chunks from original and generated # sample a random chunk from the original start_idx = np.random.randint(0, min_len - self.chunk_size) end_idx = start_idx + self.chunk_size dice_roll_2 = np.random.rand() if dice_roll_2 < 0.5: # use the same indices pos_chunk = orig_sem[start_idx:end_idx] neg_chunk = gen_sem[start_idx:end_idx] else: # use different indices start_idx_2 = np.random.randint(0, min_len - self.chunk_size) end_idx_2 = start_idx_2 + self.chunk_size pos_chunk = orig_sem[start_idx:end_idx] neg_chunk = gen_sem[start_idx_2:end_idx_2] # as an augmentation, mask out randomly a chunk of tokens from each if np.random.rand() < self.mask_prob: # mask out a chunk of tokens from the negative # ensure mask size is at most 30% of the chunk size max_mask_size = int(self.chunk_size * 0.3) mask_size = np.random.randint(1, max_mask_size + 1) mask_start = np.random.randint(0, neg_chunk.shape[0] - mask_size + 1) mask_end = mask_start + mask_size neg_chunk[mask_start:mask_end] = self.pad_token if np.random.rand() < self.mask_prob: # mask out a chunk of tokens from the positive # ensure mask size is at most 30% of the chunk size max_mask_size = int(self.chunk_size * 0.3) mask_size = np.random.randint(1, max_mask_size + 1) mask_start = np.random.randint(0, pos_chunk.shape[0] - mask_size + 1) mask_end = mask_start + mask_size pos_chunk[mask_start:mask_end] = self.pad_token return neg_chunk, pos_chunk, torch.tensor(1) class SinusoidalPositionalEncoding(nn.Module): def __init__(self, hidden_dim): super().__init__() position = torch.arange(10000).unsqueeze(1) div_term = torch.exp( torch.arange(0, hidden_dim, 2) * -(math.log(10000.0) / hidden_dim) ) pe = torch.zeros(10000, hidden_dim) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer("pe", pe) def forward(self, x): # x: (batch, num_patches, hidden_dim) return x + self.pe[: x.size(1)] class Permute(nn.Module): def __init__(self, *dims): super().__init__() self.dims = dims def forward(self, x): return x.permute(*self.dims) class AudioQualityModel(nn.Module): def __init__( self, num_labels: int = 1, hidden_dim=1024, latent_dim=128, num_heads=8, conv_layer_strides=[2, 3, 5, 8, 8], num_transformer_layers=12, dropout=0.1, objective: str = "score", input_type: str = "audio", # audio, semantic semantic_vocab_size: int = None, ): super().__init__() self.num_labels = num_labels self.hidden_dim = hidden_dim self.latent_dim = latent_dim self.objective = objective self.input_type = input_type self.semantic_vocab_size = semantic_vocab_size # wav2vec2.0 feature encoder if input_type == "audio": self.conv_layers = nn.Sequential( nn.Conv1d( 2, hidden_dim, kernel_size=7, stride=conv_layer_strides[0], padding=3, ), # nn.GroupNorm(num_groups=hidden_dim, num_channels=hidden_dim), nn.ReLU(), *[ nn.Sequential( nn.Conv1d( hidden_dim, hidden_dim, kernel_size=3, stride=stride, padding=1, ), nn.GroupNorm(num_groups=hidden_dim, num_channels=hidden_dim), nn.ReLU(), ) for stride in conv_layer_strides[1:] ], ) elif input_type == "semantic": assert semantic_vocab_size is not None # must be provided # use a learnable embedding to embed the semantic tokens self.conv_layers = nn.Sequential( nn.Embedding( semantic_vocab_size + 1, hidden_dim ), # extra for pad token Permute(0, 2, 1), ) else: raise ValueError(f"Invalid input type: {input_type}") # Position embedding self.pos_embed = SinusoidalPositionalEncoding(hidden_dim) # Transformer encoder encoder_layer = nn.TransformerEncoderLayer( d_model=hidden_dim, nhead=num_heads, dim_feedforward=hidden_dim * 4, dropout=dropout, ) self.transformer = nn.TransformerEncoder(encoder_layer, num_transformer_layers) # Output head self.mlp_head = nn.Sequential( # nn.LayerNorm(hidden_dim), nn.Linear(hidden_dim, latent_dim), # Single quality score output ) self.proj_head = nn.Sequential( # nn.LayerNorm(latent_dim), nn.Linear(latent_dim, 512), nn.GELU(), nn.Linear(512, 512), nn.GELU(), nn.Linear(512, 512), nn.GELU(), nn.Linear(512, 1), ) def get_embeddings(self, x): # x shape: (batch_size, 2, time) x = self.conv_layers(x) # (batch, dim, seq) x = x.transpose(1, 2) # (batch, seq, dim) # Add position embeddings x = self.pos_embed(x) # Maintains (batch, seq, dim) # Transformer expects (seq, batch, dim) x = x.transpose(0, 1) # (seq, batch, dim) x = self.transformer(x) x = x.transpose(0, 1) # (batch, seq, dim) x = x.mean(dim=1) # (batch, dim) x = self.mlp_head(x) # (batch, latent_dim) # l2 normalize # x = x / x.norm(dim=1, keepdim=True) return x def get_score(self, audio): """Get the score for a single audio tensor. Args: audio: torch.Tensor, shape (2, seq_len) """ assert audio.shape[0] == 2 # chunk into 5s chunks with torch.no_grad(): chunk_size = 48000 * 5 chunks = audio.unfold(1, chunk_size, chunk_size) # move the chunk dim to the batch dim chunks = chunks.permute(1, 0, 2) # iterate over chunks scores = [] for chunk in chunks: score = self(chunk.unsqueeze(0)) scores.append(score) # take the mean score across chunks mean_score = torch.stack(scores).mean(dim=0) return scores, mean_score def get_score_batch(self, audio): """Get the score for a batch of audio tensors. Args: audio: torch.Tensor, shape (batch_size, 2, seq_len) Returns: torch.Tensor: shape (batch_size,) containing quality scores for each audio """ assert audio.shape[1] == 2 # chunk into 5s chunks with torch.no_grad(): chunk_size = 48000 * 5 # unfold each audio in batch along time dimension chunks = audio.unfold( 2, chunk_size, chunk_size ) # (batch, 2, num_chunks, chunk_size) # reshape to (batch * num_chunks, 2, chunk_size) chunks = chunks.permute(0, 2, 1, 3) batch_size, num_chunks = chunks.shape[0], chunks.shape[1] chunks = chunks.reshape(-1, 2, chunk_size) # get scores for all chunks scores = self(chunks) # reshape back to (batch, num_chunks) scores = scores.reshape(batch_size, num_chunks) # take mean across chunks for each audio mean_scores = scores.mean(dim=1) return mean_scores def forward(self, audio_a, audio_b=None): # Get embeddings in one pass embeds_a = self.get_embeddings(audio_a) if audio_b is not None: embeds_b = self.get_embeddings(audio_b) embeds = torch.cat([embeds_a, embeds_b], dim=1) compare_score = self.compare_head(embeds) # audio_a_pred = self.proj_head(embeds_a) # audio_b_pred = self.proj_head(embeds_b) return compare_score # , audio_a_pred, audio_b_pred else: return self.proj_head(embeds_a) def setup_test_pairs(data_dir: str, item_ids: List[str], chunk_size: int, mode: str): print(f"Loading test set from {data_dir}...") # load all the data into memory examples = [] for item_id in tqdm(item_ids): orig_sem_path = os.path.join( data_dir, item_id, f"{item_id}_original_semantic.npz" ) gen_sem_path = os.path.join( data_dir, item_id, f"{item_id}_generated_semantic.npz" ) orig_sem = torch.from_numpy(np.load(orig_sem_path)["semantic_codes"]).long() gen_sem = torch.from_numpy(np.load(gen_sem_path)["semantic_codes"]).long() # if the generated semantic is shorter than the chunk_size, pad it with pad token if gen_sem.shape[0] <= chunk_size: continue if orig_sem.shape[0] <= chunk_size: continue if mode == "orig_gen": positive = orig_sem[:chunk_size] # original negative = gen_sem[:chunk_size] # generated elif mode == "early_late": positive = gen_sem[:chunk_size] # first chunk negative = gen_sem[-chunk_size:] # last chunk else: raise ValueError(f"Invalid mode: {mode}") # store the original and generated semantic codes examples.append((positive, negative)) return examples def run_test(model, test_pairs): model.eval() scores = [] scores_delta = [] assert len(test_pairs) > 0, "No test pairs found" for orig_sem_chunk, gen_sem_chunk in tqdm(test_pairs): orig_sem_chunk = orig_sem_chunk.cuda().unsqueeze(0) gen_sem_chunk = gen_sem_chunk.cuda().unsqueeze(0) # get the score for both the original and the corrupted audio score_orig = model(orig_sem_chunk) score_gen = model(gen_sem_chunk) score_orig = score_orig.item() score_gen = score_gen.item() # orig is always the positive # gen is always the negative scores.append(score_orig > score_gen) scores_delta.append(score_orig - score_gen) return scores, scores_delta def validate(model, val_loader): model.eval() total_loss = 0 correct = 0 total = 0 with torch.no_grad(): for batch in tqdm(val_loader): audio_a, audio_b, labels = batch audio_a = audio_a.cuda() audio_b = audio_b.cuda() labels = labels.cuda() r_i = model(audio_a) r_j = model(audio_b) loss = bradley_terry_loss(r_i, r_j, labels) total_loss += loss.item() # Calculate accuracy predictions = (r_i < r_j).float() correct += (predictions == labels).sum().item() total += labels.size(0) if dist.get_rank() == 0: # preds, scores, labels = gathered_tensors metrics = { "loss": total_loss / len(val_loader), "accuracy": correct / total if total > 0 else 0, } return metrics return None def save_checkpoint( model, optimizer, run_config, global_step, checkpoint_dir, ): if int(os.environ["LOCAL_RANK"]) == 0: checkpoint = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "run_config": run_config, "global_step": global_step, } checkpoint_path = os.path.join(checkpoint_dir, "last_ckpt.pt") torch.save(checkpoint, checkpoint_path) def bradley_terry_loss( r_i: torch.Tensor, r_j: torch.Tensor, labels: torch.Tensor ) -> torch.Tensor: """ Compute Bradley-Terry loss for paired comparisons. Args: r_i: Logits/scores for first options in pairs, shape (batch_size,) r_j: Logits/scores for second options in pairs, shape (batch_size,) labels: Binary tensor indicating whether first option (0) or second option (1) was preferred, shape (batch_size,) Returns: Mean loss value as a torch.Tensor """ # Compute negative log likelihood using logsigmoid for numerical stability loss = -( (labels) * F.logsigmoid(r_j - r_i) + (1 - labels) * F.logsigmoid(r_i - r_j) ) return loss.mean() def setup_datasets(data_dir: str, val_size: int = 100): item_dirs = [ d for d in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, d)) ] # valid item dirs are those that contain both original and generated semantic files item_ids = [ os.path.basename(d) for d in item_dirs if os.path.exists(os.path.join(data_dir, d, f"{d}_original_semantic.npz")) and os.path.exists(os.path.join(data_dir, d, f"{d}_generated_semantic.npz")) ] # sort the item_ids by name item_ids.sort() # split into train and val train_item_ids = item_ids[:-val_size] val_item_ids = item_ids[-val_size:] print(f"Found {len(item_ids)} item ids in {data_dir}") print(f"Splitting into {len(train_item_ids)} train and {len(val_item_ids)} val") return train_item_ids, val_item_ids if __name__ == "__main__": run_start_time = time.strftime("%Y-%m-%d_%H-%M-%S") checkpoint_dir = f"/app/suno/christian/checkpoints/ear-v2/{run_start_time}_s{random.randint(0, 9999)}" os.makedirs(checkpoint_dir, exist_ok=False) torch.set_float32_matmul_precision("medium") # Initialize distributed process group local_rank = int(os.environ.get("LOCAL_RANK", 0)) dist.init_process_group(backend="nccl") torch.cuda.set_device(local_rank) # Set up seeds for this process process_seed = setup_seeds() # set the seed differently for each process # torch.manual_seed(local_rank) run_config = { "training": { "max_steps": 50_000, "run_name": "base-semantic-30s", "project_name": "ear-v2", "lr": 1e-5, "grad_clip_norm": 1.0, "preload_ckpt": None, "preload_optimizer": False, "warmup_steps": 1000, }, "model": { "hidden_dim": 1024, "latent_dim": 128, "num_heads": 8, "conv_layer_strides": [2, 3, 5, 8, 8], # 25hz "num_transformer_layers": 12, "dropout": 0.1, "objective": "bradley_terry", "input_type": "semantic", "semantic_vocab_size": 4000, }, "dataset": { "data_dir": "/mnt/localdisk/christian/", "batch_size": 56, "num_workers": 1, "chunk_size": 750, # number of tokens "sample_rate": 48_000, "val_size": 1000, "mask_prob": 0.2, }, } # Initialize wandb (only one process should do this) if local_rank == 0: wandb.init( project=run_config["training"]["project_name"], name=run_config["training"]["run_name"], ) wandb.config.update( {"checkpoint_dir": checkpoint_dir, "run_config": run_config} ) # setup dataset train_item_ids, val_item_ids = setup_datasets( run_config["dataset"]["data_dir"], val_size=run_config["dataset"]["val_size"] ) train_dataset = SemanticPairsDataset( run_config["dataset"]["data_dir"], item_ids=train_item_ids, chunk_size=run_config["dataset"]["chunk_size"], mask_prob=run_config["dataset"]["mask_prob"], ) train_sampler = DistributedSampler( train_dataset, rank=local_rank, shuffle=True, seed=42 ) # Generate seed sequence for workers g_tr = torch.Generator() g_tr.manual_seed(process_seed) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=run_config["dataset"]["batch_size"], sampler=train_sampler, num_workers=run_config["dataset"]["num_workers"], persistent_workers=True, # this is necessary for the buffer to work generator=g_tr, worker_init_fn=seed_worker, ) val_dataset = SemanticPairsDataset( run_config["dataset"]["data_dir"], item_ids=val_item_ids, chunk_size=run_config["dataset"]["chunk_size"], mask_prob=0.0, ) val_sampler = DistributedSampler( val_dataset, rank=local_rank, shuffle=False, seed=42 ) g_val = torch.Generator() g_val.manual_seed(process_seed) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=run_config["dataset"]["batch_size"], sampler=val_sampler, num_workers=run_config["dataset"]["num_workers"], persistent_workers=True, generator=g_val, worker_init_fn=seed_worker, ) # setup test audio pairs # only do this on the first process if local_rank == 0: test_pairs_orig_gen = setup_test_pairs( run_config["dataset"]["data_dir"], val_item_ids, run_config["dataset"]["chunk_size"], mode="orig_gen", ) test_pairs_early_late = setup_test_pairs( run_config["dataset"]["data_dir"], val_item_ids, run_config["dataset"]["chunk_size"], mode="early_late", ) global_step = 0 # setup model model = AudioQualityModel(**run_config["model"]) num_params = sum(p.numel() for p in model.parameters()) print(f"Number of parameters: {num_params/1e6:0.1f}M") if run_config["training"]["preload_ckpt"] is not None: print(f"Preloading checkpoint from {run_config['training']['preload_ckpt']}...") checkpoint = torch.load( run_config["training"]["preload_ckpt"], map_location="cpu" ) new_state_dict = {} for k, v in checkpoint["model"].items(): if k.startswith("module."): new_state_dict[k[7:]] = v else: new_state_dict[k] = v model.load_state_dict(new_state_dict) print("Done loading checkpoint.") model.cuda() model = DistributedDataParallel(model, device_ids=[local_rank]) # model = torch.compile(model) # Add dynamo compilation optimizer = torch.optim.AdamW( model.parameters(), lr=run_config["training"]["lr"], weight_decay=1e-4 ) if ( run_config["training"]["preload_optimizer"] and run_config["training"]["preload_ckpt"] is not None ): print("Loading optimizer state dict") optimizer.load_state_dict(checkpoint["optimizer"]) global_step = checkpoint["global_step"] warmup_scheduler = LinearLR( optimizer, start_factor=0.001, end_factor=1.0, total_iters=run_config["training"]["warmup_steps"], ) cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, run_config["training"]["max_steps"] - run_config["training"]["warmup_steps"], ) scheduler = ChainedScheduler([warmup_scheduler, cosine_scheduler]) while global_step < run_config["training"]["max_steps"]: pbar = tqdm(train_loader, total=len(train_loader)) for batch in pbar: optimizer.zero_grad() # audio_a = audio[:, 0] # Shape: [bs, 2, 2, 480000] # audio_b = audio[:, 1] # Shape: [bs, 2, 2, 480000] audio_a, audio_b, label_tensor = batch audio_a = audio_a.cuda() audio_b = audio_b.cuda() label_tensor = label_tensor.cuda() # preds = model(audio_a, audio_b).squeeze(1) r_i = model(audio_a) r_j = model(audio_b) loss = bradley_terry_loss(r_i, r_j, label_tensor) loss.backward() torch.nn.utils.clip_grad_norm_( model.parameters(), run_config["training"]["grad_clip_norm"] ) optimizer.step() scheduler.step() loss = loss.mean() grad_norm = torch.norm( torch.stack( [ torch.norm(p.grad) for p in model.parameters() if p.grad is not None ] ) ) # compute the accuracy for bradley terry with torch.no_grad(): # For bradley terry, accuracy is whether the model correctly predicted # which sample was preferred based on the relative scores pred_prefs = (r_j > r_i).float() accuracy = (pred_prefs == label_tensor.float()).float().mean() # also reduce the loss if dist.is_initialized(): dist.all_reduce(loss) loss = loss / dist.get_world_size() # also reduce the accuracy if dist.is_initialized(): dist.all_reduce(accuracy) accuracy = accuracy / dist.get_world_size() if local_rank == 0: pbar.set_postfix({"loss": loss.item(), "accuracy": accuracy.item()}) wandb.log( { "train/loss": loss.item(), "train/grad_norm": grad_norm.item(), "trainer/lr": optimizer.param_groups[0]["lr"], "trainer/global_step": global_step, "train/accuracy": accuracy.item(), } ) global_step += 1 val_dict = validate(model, val_loader) if val_dict is not None and local_rank == 0: # run test on the first gpu print("Running test...") orig_gen_scores, orig_gen_scores_delta = run_test( model, test_pairs_orig_gen ) orig_gen_accuracy = sum(orig_gen_scores) / len(orig_gen_scores) orig_gen_scores_delta = sum(orig_gen_scores_delta) / len( orig_gen_scores_delta ) early_late_scores, early_late_scores_delta = run_test( model, test_pairs_early_late ) early_late_accuracy = sum(early_late_scores) / len(early_late_scores) early_late_scores_delta = sum(early_late_scores_delta) / len( early_late_scores_delta ) print("Test done.") metrics_to_log = { "val/loss": val_dict["loss"], "val/accuracy": val_dict["accuracy"], "test/orig_vs_gen_accuracy": orig_gen_accuracy, "test/orig_vs_gen_scores_delta": orig_gen_scores_delta, "test/early_vs_late_accuracy": early_late_accuracy, "test/early_vs_late_scores_delta": early_late_scores_delta, } wandb.log(metrics_to_log) save_checkpoint(model, optimizer, run_config, global_step, checkpoint_dir) print("Done!")