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 torch.nn as nn import matplotlib.pyplot as plt import torch.distributed as dist 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 from suno_utils.utils.text import read_jsonl class PreferenceAudioDataset(torch.utils.data.Dataset): def __init__( self, metas_filepath: str, audio_dir: str, sample_rate: int, num_workers: int = 1, chunk_size_s: float = 5.0, buffer_size: int = 10_000, max_chunks_per_file: int = 100, random_crop: bool = False, # change this to probability of random crop ): self.metas_filepath = metas_filepath self.audio_dir = audio_dir self.sample_rate = sample_rate self.chunk_size_s = chunk_size_s self.buffer_size = buffer_size self.chunk_size_samples = int(chunk_size_s * sample_rate) self.num_workers = num_workers self.max_chunks_per_file = max_chunks_per_file self.random_crop = random_crop self.preprocess_chunk_size_samples = self.chunk_size_samples # if random_crop, then we adjust chunks to be larger if self.random_crop: self.preprocess_chunk_size_samples = self.chunk_size_samples * 1.5 self.preprocess_chunk_size_samples = int(self.preprocess_chunk_size_samples) self.items_since_last_reload = buffer_size # force a reload self.buffer = [] # load manifest metas = read_jsonl(metas_filepath) print(f"Loaded {len(metas)} metas from {metas_filepath}") self.metas = metas def _reload_buffer(self): self.buffer = [] self.items_since_last_reload = 0 # load new metas # Create a list of indices and shuffle them indices = list(range(len(self.metas))) np.random.shuffle(indices) # Iterate through shuffled indices until buffer is full pbar = tqdm(indices) for meta_idx in pbar: meta = self.metas[meta_idx] positive_id = meta["positive_id"] negative_id = meta["negative_id"] pbar.set_postfix({"buffer_size": len(self.buffer)}) try: pos_audio, sr = torchaudio.load( os.path.join(self.audio_dir, positive_id + ".mp3") ) neg_audio, sr = torchaudio.load( os.path.join(self.audio_dir, negative_id + ".mp3") ) except Exception as e: # print(f"Error loading audio for {positive_id} and {negative_id}: {e}") continue # skip if audio is too short if ( pos_audio.shape[1] < self.preprocess_chunk_size_samples or neg_audio.shape[1] < self.preprocess_chunk_size_samples ): continue # chunk into chunks of chunk_size_samples pos_chunks = pos_audio.unfold( 1, self.preprocess_chunk_size_samples, self.preprocess_chunk_size_samples, ) neg_chunks = neg_audio.unfold( 1, self.preprocess_chunk_size_samples, self.preprocess_chunk_size_samples, ) # torch.Size([2, 48, 240000]) torch.Size([2, 48, 240000]) # print(pos_chunks.shape, neg_chunks.shape) # sys.exit() # add to buffer as separate items for chunk_idx in range(pos_chunks.shape[1]): self.buffer.append((pos_chunks[:, chunk_idx], neg_chunks[:, chunk_idx])) # exit if buffer is full if len(self.buffer) >= self.buffer_size: break def __len__(self): return self.buffer_size * 1000 def __getitem__(self, _): if self.items_since_last_reload >= self.buffer_size: self._reload_buffer() # meta = self.metas[idx] # positive_id = meta["positive_id"] # negative_id = meta["negative_id"] # pos_audio, sr = torchaudio.load(self.audio_dir, positive_id + ".mp3") # neg_audio, sr = torchaudio.load(self.audio_dir, negative_id + ".mp3") # make sure buffer is loaded audios = [] # get a random item from buffer pos_audio, neg_audio = self.buffer[np.random.randint(0, len(self.buffer) - 1)] for audio in [pos_audio, neg_audio]: # random crop to chunk size if self.random_crop and np.random.uniform() < 0.5: start_idx = np.random.randint( 0, audio.shape[-1] - self.chunk_size_samples ) end_idx = start_idx + self.chunk_size_samples audio = audio[..., start_idx:end_idx] else: audio = audio[..., : self.chunk_size_samples] # if np.random.uniform() < 0.5: # peak normalize # audio = audio / audio.abs().max().clamp(1e-8) # if np.random.uniform() < 0.5: # apply gain reduction # gain_reduction_db = np.random.uniform(-12, 0) # audio *= 10 ** (gain_reduction_db / 20.0) audios.append(audio) label_tensor = torch.tensor([0.0]) # in this case always return 0 # increment items since last reload self.items_since_last_reload += 1 return torch.stack(audios), label_tensor 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 AudioQualityModel(nn.Module): def __init__( self, num_labels: int = 1, # not used for reward model 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", ): super().__init__() self.num_labels = num_labels self.hidden_dim = hidden_dim self.latent_dim = latent_dim self.objective = objective # wav2vec2.0 feature encoder 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:] ], ) # 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): # always peak normalize # x = x / x.abs().max().clamp(1e-8) # x shape: (batch_size, 2, time) x = self.conv_layers(x) # (batch, 512, seq) x = x.transpose(1, 2) # (batch, seq, 512) # Add position embeddings x = self.pos_embed(x) # Maintains (batch, seq, 512) # Transformer expects (seq, batch, dim) x = x.transpose(0, 1) # (seq, batch, 512) x = self.transformer(x) x = x.transpose(0, 1) # (batch, seq, 512) x = x.mean(dim=1) # (batch, 512) 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) # get the score for each chunk scores = self(chunks) # take the mean score across chunks mean_score = scores.mean(dim=0) return 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 run_test(model, audio_pairs): model.eval() scores = [] scores_delta = [] assert len(audio_pairs) > 0, "No audio pairs found" for audio_a, audio_b, label in tqdm(audio_pairs): audio_a = audio_a.cuda() audio_b = audio_b.cuda() # get the score for both the original and the corrupted audio score_a = model.module.get_score(audio_a).item() score_b = model.module.get_score(audio_b).item() # we want the original to be lower than then corrupted score # and we want the delta to be as high as possible if label == "a" and score_a > score_b: scores.append(1) scores_delta.append(score_b - score_a) elif label == "b" and score_b > score_a: scores.append(1) scores_delta.append(score_a - score_b) else: scores.append(0) scores_delta.append(score_a - score_b) return scores, scores_delta def process_audio_pairs(base_dir: str): """ Process audio pairs from different test directories and return paired data. Args: base_dir: Base directory containing the test folders Returns: List of tuples containing (audio_pairs, test_type) """ # Test types and their corresponding directories test_types = { "codec_cycled": "codec_cycle_test", "diff_cycled": "diffusion_cycle_test", "gen_pref": "dpo_diffusion_test_set", } def load_and_process_audio(filepath: str): """Load and preprocess audio file.""" try: x, sr = torchaudio.load(filepath) if sr != 48000: x = torchaudio.functional.resample(x, sr, 48000) return x[:, : 48000 * 120] # Trim to 120 seconds except Exception as e: print(f"Error processing {filepath}: {e}") return None def find_pairs(filepaths): """Find and process audio pairs from a list of file paths.""" pairs = [] processed_files = set() for filepath in tqdm(filepaths): if filepath in processed_files: continue if "a_input" in filepath: base_name = filepath.split("a_input_")[0] else: base_name = filepath.split("_input_")[0].split("_cycled_")[0] pref_true = "pref=True" in filepath pref_false = "pref=False" in filepath # Load first audio file x = load_and_process_audio(filepath) if x is None: continue # Find matching pair for other_path in filepaths: if base_name in other_path and other_path != filepath: y = load_and_process_audio(other_path) if y is None: continue # Add pair with correct label based on preference if pref_true: pairs.append((x, y, "a")) elif pref_false: pairs.append((y, x, "b")) processed_files.add(filepath) processed_files.add(other_path) break return pairs # Process all test types results = [] for test_type, folder in test_types.items(): dir_path = os.path.join(base_dir, folder) print(f"Processing {test_type} from {dir_path}") filepaths = glob.glob(os.path.join(dir_path, "*.mp3")) pairs = find_pairs(filepaths) print(f"Found {len(pairs)} pairs for {test_type}") results.append((pairs, test_type)) return results def validate(model, val_loader): model.eval() total_loss = 0 predictions = [] # Will store (preds, scores, labels) tuples with torch.no_grad(): for batch in tqdm(val_loader): audio, labels = batch audio, labels = audio.cuda(), labels.cuda() if run_config["model"]["objective"] == "bradley_terry": audio_a = audio[:, 0] # Shape: [4, 2, 480000] audio_b = audio[:, 1] # Shape: [4, 2, 480000] r_i = model(audio_a) r_j = model(audio_b) loss = bradley_terry_loss(r_i, r_j, labels) elif run_config["model"]["objective"] == "score": scores = model(audio) loss = torch.nn.functional.mse_loss(scores.squeeze(1), labels) else: scores = model(audio) loss = torch.nn.functional.binary_cross_entropy_with_logits( scores, labels ) total_loss += loss.item() # predictions.append( # (torch.sigmoid(scores) > 0.5, torch.sigmoid(scores), labels) # ) # Concatenate all batches # preds, scores, labels = [torch.cat(x, dim=0) for x in zip(*predictions)] # Gather from all processes # world_size = dist.get_world_size() # gathered_tensors = [] # for tensor in [preds, scores, labels]: # gathered = [torch.zeros_like(tensor) for _ in range(world_size)] # dist.all_gather(gathered, tensor) # gathered_tensors.append(torch.cat(gathered).cpu().numpy()) if dist.get_rank() == 0: # preds, scores, labels = gathered_tensors metrics = { "loss": total_loss / len(val_loader), # "accuracy": ( # (preds == labels).all(axis=1).mean() # if len(labels.shape) > 1 and labels.shape[1] > 1 # else (preds == labels).mean() # ), # "macro_f1": f1_score(labels, preds, average="macro", zero_division=0), # "micro_f1": f1_score(labels, preds, average="micro", zero_division=0), } # Add threshold-free metrics if False: try: metrics.update( { "auc": roc_auc_score(labels.ravel(), scores.ravel()), "ap": average_precision_score(labels.ravel(), scores.ravel()), } ) except ValueError: metrics.update({"auc": 0.0, "ap": 0.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, f"last_ckpt.pt") torch.save(checkpoint, checkpoint_path) def load_corruptions_config(filename): with open(filename, "r") as f: return json.load(f) from torch.nn import functional as F 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 """ idx = 0 # batch idx # Compute negative log likelihood using logsigmoid for numerical stability loss = -( (labels) * F.logsigmoid(r_i - r_j) + (1 - labels) * F.logsigmoid(r_j - r_i) ) return loss.mean() 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": 100_000, "run_name": "base-5s-reward-v1", "project_name": "ear-v2", "lr": 1e-6, "grad_clip_norm": 10.0, "preload_ckpt": None, "preload_optimizer": False, "warmup_steps": 500, }, "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", }, "dataset": { "audio_dir": "/app/suno/christian/data/reward_model/remaster_v1/audio", "train_metas": "/app/suno/christian/data/reward_model/remaster_v1/metas.jsonl", "val_metas": "/app/suno/christian/data/reward_model/remaster_v1/metas.jsonl", "batch_size": 24, "num_workers": 8, "max_chunks_per_file": None, "buffer_size": 10_000, "chunk_size_s": 5.0, "sample_rate": 48_000, "random_crop": False, }, } # setup test audio pairs # only do this on the first process if local_rank == 0: test_audio_pairs_list = process_audio_pairs("/home/christian/audio/ear-bench") print(len(test_audio_pairs_list)) # 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_dataset = PreferenceAudioDataset( metas_filepath=run_config["dataset"]["train_metas"], audio_dir=run_config["dataset"]["audio_dir"], sample_rate=run_config["dataset"]["sample_rate"], num_workers=run_config["dataset"]["num_workers"], chunk_size_s=run_config["dataset"]["chunk_size_s"], buffer_size=run_config["dataset"]["buffer_size"], max_chunks_per_file=run_config["dataset"]["max_chunks_per_file"], random_crop=run_config["dataset"]["random_crop"], ) 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 = PreferenceAudioDataset( metas_filepath=run_config["dataset"]["val_metas"], audio_dir=run_config["dataset"]["audio_dir"], sample_rate=run_config["dataset"]["sample_rate"], num_workers=run_config["dataset"]["num_workers"], chunk_size_s=run_config["dataset"]["chunk_size_s"], buffer_size=run_config["dataset"]["buffer_size"] // 10, max_chunks_per_file=run_config["dataset"]["max_chunks_per_file"], random_crop=run_config["dataset"]["random_crop"], ) 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, ) 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, label_tensor = batch audio = audio.cuda() label_tensor = label_tensor.cuda() if run_config["model"]["objective"] == "bradley_terry": audio_a = audio[:, 0] # Shape: [bs, 2, 2, 480000] audio_b = audio[:, 1] # Shape: [bs, 2, 2, 480000] # 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) elif run_config["model"]["objective"] == "score": preds = model(audio) loss = torch.nn.functional.mse_loss(preds.squeeze(1), label_tensor) else: preds = model(audio) loss = torch.nn.functional.binary_cross_entropy_with_logits( preds, 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(): if run_config["model"]["objective"] == "bradley_terry": # 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() else: preds = torch.sigmoid(preds) accuracy = ( ((preds > 0.5).float() == 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 results = {} print("Running test...") for test_audio_pairs, name in test_audio_pairs_list: scores, scores_delta = run_test(model, test_audio_pairs) results[name] = { "accuracy": sum(scores) / len(scores), "scores_delta": sum(scores_delta) / len(scores_delta), } print("Test done.") metrics_to_log = { "val/loss": val_dict["loss"], } for name, metrics in results.items(): metrics_to_log[f"test/{name}_accuracy"] = metrics["accuracy"] metrics_to_log[f"test/{name}_scores_delta"] = metrics["scores_delta"] wandb.log(metrics_to_log) save_checkpoint(model, optimizer, run_config, global_step, checkpoint_dir) print("Done!")