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 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 import numpy as np from scipy import signal from typing import List, Tuple def biquad( gain_db: float, cutoff_freq: float, q_factor: float, sample_rate: float, filter_type: str, ) -> Tuple[np.ndarray, np.ndarray]: """Use design parameters to generate coefficients for a specific filter type.""" A = 10 ** (gain_db / 40.0) w0 = 2.0 * np.pi * (cutoff_freq / sample_rate) alpha = np.sin(w0) / (2.0 * q_factor) cos_w0 = np.cos(w0) sqrt_A = np.sqrt(A) if filter_type == "high_shelf": b0 = A * ((A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha) b1 = -2 * A * ((A - 1) + (A + 1) * cos_w0) b2 = A * ((A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha) a0 = (A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha a1 = 2 * ((A - 1) - (A + 1) * cos_w0) a2 = (A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha elif filter_type == "low_shelf": b0 = A * ((A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha) b1 = 2 * A * ((A - 1) - (A + 1) * cos_w0) b2 = A * ((A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha) a0 = (A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha a1 = -2 * ((A - 1) + (A + 1) * cos_w0) a2 = (A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha elif filter_type == "peaking": b0 = 1 + alpha * A b1 = -2 * cos_w0 b2 = 1 - alpha * A a0 = 1 + alpha / A a1 = -2 * cos_w0 a2 = 1 - alpha / A b = np.array([b0, b1, b2]) / a0 a = np.array([1.0, a1 / a0, a2 / a0]) return b, a def apply_random_eq( x: torch.Tensor, sample_rate: float, n_bands: int = 5, prob_active: float = 0.2, dtype=np.float32, ) -> torch.Tensor: """Apply random equalization to audio signal. Args: x: Input audio signal sample_rate: Sample rate in Hz n_bands: Number of peaking bands prob_active: Probability of each band being active dtype: Output data type Returns: Equalized audio signal """ # Define parameter ranges gain_range = (-12.0, 12.0) low_shelf_freq_range = (20.0, 200.0) high_shelf_freq_range = (4000.0, 20000.0) peaking_freq_range = (200.0, 4000.0) q_range = (0.1, 10.0) # Helper function for log-uniform random sampling def random_log_uniform(min_val, max_val): log_min = np.log(min_val) log_max = np.log(max_val) return np.exp(np.random.uniform(log_min, log_max)) # convert x to numpy x = x.numpy() # Apply low-shelf filter (random chance of being active) if np.random.random() < prob_active: low_shelf_params = { "gain_db": np.random.uniform(*gain_range), "cutoff_freq": random_log_uniform(*low_shelf_freq_range), "q_factor": np.random.uniform(*q_range), } b, a = biquad( low_shelf_params["gain_db"], low_shelf_params["cutoff_freq"], low_shelf_params["q_factor"], sample_rate, "low_shelf", ) x = signal.lfilter(b, a, x) # Apply peaking filters for _ in range(n_bands): if np.random.random() < prob_active: peak_params = { "gain_db": np.random.uniform(*gain_range), "cutoff_freq": random_log_uniform(*peaking_freq_range), "q_factor": np.random.uniform(*q_range), } b, a = biquad( peak_params["gain_db"], peak_params["cutoff_freq"], peak_params["q_factor"], sample_rate, "peaking", ) x = signal.lfilter(b, a, x) # Apply high-shelf filter (random chance of being active) if np.random.random() < prob_active: high_shelf_params = { "gain_db": np.random.uniform(*gain_range), "cutoff_freq": random_log_uniform(*high_shelf_freq_range), "q_factor": np.random.uniform(*q_range), } b, a = biquad( high_shelf_params["gain_db"], high_shelf_params["cutoff_freq"], high_shelf_params["q_factor"], sample_rate, "high_shelf", ) x = signal.lfilter(b, a, x) return torch.from_numpy(x) def apply_random_tanh_distortion(audio: torch.Tensor, sample_rate: float): gain_db = random.uniform(0, 12) return apply_tanh_distortion(audio, sample_rate, gain_db) def apply_random_clipping_distortion(audio: torch.Tensor, sample_rate: float): gain_db = random.uniform(0, 12) return apply_clipping_distortion(audio, sample_rate, gain_db) def apply_random_noise(audio: torch.Tensor, sample_rate: float): noise_type = random.choice(["white", "pink"]) noise_gain = random.uniform(-48, -6) return apply_noise(audio, sample_rate, noise_gain, noise_type) def apply_random_stereo_to_mono(audio: torch.Tensor, sample_rate: float): return apply_stereo_to_mono(audio, sample_rate) def apply_random_channel_imbalance(audio: torch.Tensor, sample_rate: float): imbalance = random.uniform(-1.0, 1.0) return apply_channel_imbalance(audio, sample_rate, imbalance) def apply_random_freq_boost(audio: torch.Tensor, sample_rate: float): pass def apply_random_highpass(audio: torch.Tensor, sample_rate: float): pass # ---------------- functional corruptions with parameters ---------------- def apply_stereo_to_mono(audio: torch.Tensor, sample_rate: float): return audio.mean(dim=0, keepdims=True).repeat(2, 1) def apply_channel_imbalance( audio: torch.Tensor, sample_rate: float, imbalance: float = 0.0 ): if not -1 <= imbalance <= 1 or audio.shape[-2] != 2: raise ValueError("Invalid input") out = audio.clone() l_gain, r_gain = (1.0 - imbalance, 1.0) if imbalance > 0 else (1.0, 1.0 + imbalance) out[0, :], out[1, :] = out[0, :] * l_gain, out[1, :] * r_gain return out def apply_freq_boost( audio: torch.Tensor, sample_rate: float, freq_hz: float = 1000.0, gain_db: float = 10.0, ): # Design peaking filter with Q=0.707 w0 = 2 * math.pi * freq_hz / sample_rate alpha = math.sin(w0) / (2 * 0.707) A = 10 ** (gain_db / 40.0) b0 = 1 + alpha * A b1 = -2 * math.cos(w0) b2 = 1 - alpha * A a0 = 1 + alpha / A a1 = -2 * math.cos(w0) a2 = 1 - alpha / A # Normalize coefficients by a0 b0 = b0 / a0 b1 = b1 / a0 b2 = b2 / a0 a1 = a1 / a0 a2 = a2 / a0 a0 = 1.0 return torchaudio.functional.biquad(audio, b0, b1, b2, a0, a1, a2) def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0): return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz) def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0): return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz) def apply_bandpass( audio: torch.Tensor, sample_rate: float, central_freq: float = 1000.0, bandwidth: float = 0.707, ): return torchaudio.functional.bandpass_biquad( audio, sample_rate, central_freq, bandwidth ) def apply_tanh_distortion( audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0 ): gain_lin = 10 ** (gain_db / 20.0) return torch.tanh(audio * gain_lin) def apply_clipping_distortion( audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0 ): gain_lin = 10 ** (gain_db / 20.0) return (audio * gain_lin).clamp(-1, 1) def apply_noise( audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0, noise_type: str = "white", ): gain_lin = 10 ** (gain_db / 20.0) noise = torch.randn_like(audio) if noise_type == "white": return audio + gain_lin * noise elif noise_type == "pink": b = torch.tensor([0.049922035, -0.095993537, 0.050612699, -0.004408786]) a = torch.tensor([1, -2.494956002, 2.017265875, -0.522189400]) noise = torchaudio.functional.filtfilt(noise, a, b) noise /= noise.abs().max() return audio + gain_lin * noise else: raise ValueError(f"Invalid noise type: {noise_type}") def apply_dc_offset( audio: torch.Tensor, sample_rate: float, offset: float = 0.0, mode: str = "constant" ): if mode == "constant": return audio + offset elif mode == "ramp": ramp = torch.linspace(0, offset, audio.shape[-1]) return audio + ramp else: raise ValueError(f"Invalid mode: {mode}") def apply_flip_polarity(audio: torch.Tensor, sample_rate: float): audio = audio.clone() channel = torch.randint(0, audio.shape[0], (1,)).item() audio[channel] *= -1 return audio def apply_audio_codec_advanced( audio: torch.Tensor, sample_rate: float, bit_rate: int = 16000, n_passes: int = 1, codec_type: str = "mp3", ): for _ in range(n_passes): if codec_type == "mp3": effector = torchaudio.io.AudioEffector( format="mp3", codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate), ) elif codec_type == "ogg-vorbis": effector = torchaudio.io.AudioEffector( format="ogg", encoder="vorbis", codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate), ) elif codec_type == "opus": effector = torchaudio.io.AudioEffector( format="ogg", encoder="opus", codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate), ) else: raise ValueError(f"Invalid codec type: {codec_type}") audio = effector.apply(audio.T, sample_rate).T return audio def apply_audio_codec( audio: torch.Tensor, sample_rate: float, bit_rate: int = 16000, n_passes: int = 1, format_str: str = "mp3", ): for _ in range(n_passes): effector = torchaudio.io.AudioEffector( format=format_str, codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate), ) audio = effector.apply(audio.T, sample_rate).T return audio def apply_hum( audio: torch.Tensor, sample_rate: float, amplitude: float = 0.0, freq: float = 0.0 ): t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate hum = amplitude * torch.sin(2 * np.pi * freq * t) # Add harmonics at 2x hum += (amplitude * 0.5) * torch.sin(2 * np.pi * 2 * freq * t) return audio + hum.expand_as(audio) def apply_comb_filter( audio: torch.Tensor, sample_rate: float, delay_ms: float = 0.0, gain_db: float = 0.0 ): delay_samples = int(delay_ms * sample_rate / 1000) gain_lin = 10 ** (gain_db / 20.0) delayed = torch.roll(audio, shifts=delay_samples, dims=-1) return audio + gain_lin * delayed # this function may be bad def apply_reduce_bit_depth(audio: torch.Tensor, sample_rate: float, bits: int = 8): steps = 2**bits return (audio.clamp(-1, 1) * 0.5 + 0.5) * (steps - 1) // 1 / (steps - 1) * 2 - 1 def apply_add_clicks(audio: torch.Tensor, sample_rate: float, density: float = 0.001): mask = torch.rand_like(audio) < density clicks = (torch.rand_like(audio) * 2 - 1) * mask return audio + clicks def apply_stereo_width(audio: torch.Tensor, sample_rate: float, width: float = 1.0): left, right = audio[0], audio[1] mid = (left + right) * 0.5 side = (left - right) * 0.5 side = ( side * width ) # when width < 1, side is narrower, when width > 1, side is wider return torch.stack([mid + side, mid - side]) def apply_spectral_mask( audio: torch.Tensor, sample_rate: float, threshold: float = -60, ratio: float = 0.5, n_fft: int = 2048, ): window = torch.hann_window(n_fft).to(audio.device) spec = torch.stft(audio, n_fft, n_fft // 4, window=window, return_complex=True) mask = torch.where(20 * torch.log10(torch.abs(spec) + 1e-8) < threshold, ratio, 1.0) return torch.istft( spec * mask, n_fft, n_fft // 4, window=window, length=audio.shape[-1] ) def apply_time_stretch(audio: torch.Tensor, sample_rate: float, rate: float = 1.0): effects = [ ["tempo", str(rate)], ] return torchaudio.sox_effects.apply_effects_tensor(audio, sample_rate, effects)[0] def apply_wow_flutter( audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1 ): t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate mod = depth * torch.sin(2 * torch.pi * rate * t) # Convert modulation to sample offsets offsets = (mod * sample_rate).long() # Apply time-varying delay output = torch.zeros_like(audio) for i in range(audio.shape[-1]): idx = max(0, min(i + offsets[i].item(), audio.shape[-1] - 1)) output[..., i] = audio[..., idx] return output def apply_wow_flutter_fast( audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1 ): t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate mod = depth * torch.sin(2 * torch.pi * rate * t) indices = torch.arange(audio.shape[-1], device=audio.device) indices = indices + (mod * sample_rate).long() indices = torch.clamp(indices, 0, audio.shape[-1] - 1) while indices.dim() < audio.dim(): indices = indices.unsqueeze(0) indices = indices.expand_as(audio) output = torch.gather(audio, -1, indices) return output def apply_stereo_fold(audio, sample_rate): mono = audio.mean(dim=0, keepdim=True) # Add phase issues return torch.cat([mono, -mono], dim=0) def apply_ring_modulation(audio, sample_rate, freq=440, mix=0.2): samples = audio.shape[-1] t = torch.linspace(0, samples / sample_rate, samples, device=audio.device) freq = freq + 10 * torch.sin(2 * torch.pi * 0.5 * t) phase = 2 * torch.pi * freq * t carrier = torch.sin(phase).view(1, -1) # Changed from (1,1,-1) to (1,-1) modulated = audio * carrier return (1 - mix) * audio + mix * modulated def apply_white_noise_burst( audio, sample_rate, noise_level=0.1, min_burst_length=500, max_burst_length=8000, p_burst=0.01, ): # Use shortest burst length to determine number of segments num_segments = audio.shape[-1] // min_burst_length # Generate random burst lengths and levels burst_lengths = torch.randint( min_burst_length, max_burst_length, (num_segments,), device=audio.device ) burst_levels = noise_level * (0.5 + torch.rand(num_segments, device=audio.device)) burst_mask = ( torch.rand(num_segments, device=audio.device) < p_burst ).bool() # Changed to bool # Create index tensor for the full audio length indices = torch.arange(audio.shape[-1], device=audio.device) # Create cumulative positions positions = torch.cumsum(burst_lengths, dim=0) starts = torch.cat([torch.tensor([0], device=audio.device), positions[:-1]]) # Create mask using broadcasting mask = torch.zeros(audio.shape[-1], device=audio.device) valid_mask = (indices.unsqueeze(0) >= starts.unsqueeze(1)) & ( indices.unsqueeze(0) < positions.unsqueeze(1) ) valid_mask = valid_mask & burst_mask.unsqueeze(1) # Now both are boolean # Convert boolean mask to burst levels mask = (valid_mask.float() * burst_levels.unsqueeze(1)).max(dim=0)[0] # Expand mask to match audio dimensions mask = mask.view(1, -1).expand_as(audio) # Apply noise noise = torch.randn_like(audio) return audio + noise * mask def apply_quantize_zero(audio, sample_rate, threshold=0.001): return torch.where(torch.abs(audio) < threshold, 0, audio) def apply_phase_randomize( audio: torch.Tensor, sample_rate: float, block_size: int = 2048, mix: float = 0.75 ): window = torch.hann_window(block_size, device=audio.device) # Process each channel output = [] for channel in audio: stft = torch.stft(channel, block_size, window=window, return_complex=True) mag = stft.abs() random_phase = torch.exp(2j * torch.pi * torch.rand_like(stft)) channel_out = torch.istft( mag * random_phase, block_size, window=window, length=channel.shape[-1] ) output.append(channel_out) return (1 - mix) * audio + mix * torch.stack(output) def apply_reverb( audio: torch.Tensor, sample_rate: float, reverberance: int = 50, # 0-100 hf_damping: int = 50, # 0-100 room_scale: int = 100, # 0-100 stereo_depth: int = 100, # 0-100 pre_delay: float = 0, # 0-200ms wet_gain: float = 0, ): # -10-10 dB effects = [ [ "reverb", str(reverberance), str(hf_damping), str(room_scale), str(stereo_depth), str(pre_delay), str(wet_gain), ] ] out, _ = torchaudio.sox_effects.apply_effects_tensor(audio, sample_rate, effects) return out # Separate function mapping corruption_functions = { "stereo_to_mono": apply_stereo_to_mono, "channel_imbalance": apply_channel_imbalance, "lowpass": apply_lowpass, "bandpass": apply_bandpass, "highpass": apply_highpass, "tanh_distortion": apply_tanh_distortion, "clipping_distortion": apply_clipping_distortion, "noise": apply_noise, "hum": apply_hum, "comb_filter": apply_comb_filter, "reduce_bit_depth": apply_reduce_bit_depth, "add_clicks": apply_add_clicks, "reverb": apply_reverb, "audio_codec": apply_audio_codec, "dc_offset": apply_dc_offset, "flip_polarity": apply_flip_polarity, "stereo_width": apply_stereo_width, "spectral_mask": apply_spectral_mask, "time_stretch": apply_time_stretch, "wow_flutter": apply_wow_flutter_fast, "stereo_fold": apply_stereo_fold, "ring_modulation": apply_ring_modulation, "white_noise_burst": apply_white_noise_burst, "quantize_zero": apply_quantize_zero, "phase_randomize": apply_phase_randomize, "freq_boost": apply_freq_boost, "audio_codec_advanced": apply_audio_codec_advanced, } # now we have a function that takes a preset and applies the relevant corruptions to the audio # Modified apply_preset function to use both config and functions def apply_preset(audio: torch.Tensor, sr: float, preset: dict, functions: dict): chs, seq_len = audio.shape for corruption_name, corruption_info in preset.items(): audio = functions[corruption_name]( audio.clone(), sr, **corruption_info["params"] ) # audio = torch.clamp(audio, -1, 1) # clip to -1, 1 # repeat pad to original length if audio.shape[-1] < seq_len: audio = audio.repeat(1, seq_len) # crop to original length audio = audio[..., :seq_len] return audio def create_label_encoder(corruptions_dict: dict) -> Dict[str, int]: """ Create a mapping from all possible corruption parameter combinations to indices. For corruptions with multiple parameters, creates labels for all combinations. """ label_to_idx = {} idx = 0 for corruption_name, corruption_info in corruptions_dict.items(): params = corruption_info["params"] # If corruption has no parameters if not params: label = f"{corruption_name}" label_to_idx[label] = idx idx += 1 continue # Get all parameter names and their possible values param_names = list(params.keys()) param_values = [params[name] for name in param_names] # Generate all possible combinations of parameter values for values in itertools.product(*param_values): # Create parameter string param_str = ",".join( f"{name}={value}" for name, value in zip(param_names, values) ) label = f"{corruption_name}:{param_str}" label_to_idx[label] = idx idx += 1 return label_to_idx def sample_n_corruptions(max_corruptions: int = 5, p: float = 0.5): probs = np.array([(1 - p) ** i * p for i in range(max_corruptions)]) probs = probs / probs.sum() return np.random.choice(np.arange(max_corruptions), p=probs) + 1 # Similarly, we need to update how we generate labels in the preset function def generate_random_preset( corruptions_dict: dict, max_corruptions: int = 10, no_corruption_probability: float = 0.01, ): """ Generate a random preset and its corresponding labels. Handles corruptions with multiple parameters. """ if random.random() < no_corruption_probability: return {}, set() max_corruptions = len(corruptions_dict) # sample n_corruptions from exponential distribution n_corruptions = sample_n_corruptions(max_corruptions, p=0.5) if n_corruptions == 0: return {}, set() if len(corruptions_dict) == 1: selected_corruptions = [list(corruptions_dict.keys())[0]] else: selected_corruptions = random.sample( list(corruptions_dict.keys()), n_corruptions ) preset = {} labels = set() for corruption_name in selected_corruptions: corruption_info = corruptions_dict[corruption_name] params = {} # If no parameters, just add the corruption name if not corruption_info["params"]: preset[corruption_name] = {"params": {}} labels.add(corruption_name) continue # Generate parameters and create combined label param_strs = [] for param_name, param_values in corruption_info["params"].items(): param_value = random.choice(param_values) params[param_name] = param_value param_strs.append(f"{param_name}={param_value}") preset[corruption_name] = {"params": params} # Create single label with all parameters label = f"{corruption_name}:{','.join(param_strs)}" labels.add(label) return preset, labels def labels_to_tensor( labels: Set[str], label_encoder: Dict[str, int], device: str = "cpu" ) -> torch.Tensor: """ Convert a set of labels to a binary tensor. """ output = torch.zeros(len(label_encoder), dtype=torch.float32, device=device) for label in labels: if label in label_encoder: output[label_encoder[label]] = 1.0 return output 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 class CorruptAudioDataset(torch.utils.data.Dataset): def __init__( self, manifest_path: str, label_encoder: Dict[str, int], corruptions: Dict[str, Dict[str, Any]], sample_rate: int, max_corruptions: int = 10, no_corruption_probability: float = 0.5, num_workers: int = 1, chunk_size_s: float = 5.0, buffer_size: int = 50_000, max_chunks_per_file: int = 100, num_versions: int = 1, random_crop: bool = False, ): self.manifest_path = manifest_path 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_corruptions = max_corruptions self.no_corruption_probability = no_corruption_probability self.max_chunks_per_file = max_chunks_per_file self.num_versions = num_versions 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 *= 1.25 self.preprocess_chunk_size_samples = int(self.preprocess_chunk_size_samples) # assert self.num_versions > 1, "num_versions must be greater than 1" self.label_encoder = label_encoder self.num_labels = len(self.label_encoder) self.corruptions = corruptions self.items_since_last_reload = buffer_size # force a reload self.buffer = [] self.loss_fn = auraloss.freq.MelSTFTLoss( sample_rate=self.sample_rate, fft_size=2048, win_length=2048, hop_size=1024, n_mels=64, ) # self.loss_fn = auraloss.time.SISDRLoss() # load manifest with open(manifest_path, "r") as f: self.filepaths = [line.strip() for line in f.readlines()] print(f"Loaded {len(self.filepaths)} filepaths from {manifest_path}") def __len__(self): return self.buffer_size * self.num_workers def _reload_buffer(self): self.buffer = [] rand_idxs = torch.randperm(len(self.filepaths)) # max rand_idxs repeat endlessly rand_idxs = itertools.cycle(rand_idxs) # pbar = tqdm(rand_idxs, total=len(self.filepaths), desc="Loading audio buffer") for idx in rand_idxs: if len(self.buffer) >= self.buffer_size: break try: filepath = self.filepaths[idx] audio, sr = torchaudio.load(filepath) if sr != self.sample_rate: audio = torchaudio.functional.resample(audio, sr, self.sample_rate) # Pad if needed to ensure consistent chunk size if audio.shape[-1] < self.preprocess_chunk_size_samples: continue # Split into chunks chunks = audio.unfold( -1, self.preprocess_chunk_size_samples, self.preprocess_chunk_size_samples, ) chunks = chunks.chunk(chunks.shape[1], dim=1) # Filter chunks by minimum length valid_chunks = [ chunk.squeeze(1) for chunk in chunks if chunk.shape[-1] >= self.preprocess_chunk_size_samples ] # filter out chunks of silence valid_chunks = [ chunk for chunk in valid_chunks if (chunk.abs() ** 2).mean() > 0.001 ] # limit to max_chunks_per_file valid_chunks = valid_chunks[: self.max_chunks_per_file] self.buffer.extend(valid_chunks) # pbar.set_postfix({"buffer_size": len(self.buffer)}) except Exception as e: print(f"Error loading {filepath}: {e}") continue self.items_since_last_reload = 0 def __getitem__(self, _): # reload buffer if needed if self.items_since_last_reload >= len(self.buffer): self._reload_buffer() # sample a random audio from the buffer buffer_idx = np.random.randint(0, len(self.buffer)) audio = self.buffer[buffer_idx].clone() corrupted_audios = [] label_tensors = [] mse_losses = [] for version_idx in range(self.num_versions): preset, labels = generate_random_preset( self.corruptions, self.max_corruptions, self.no_corruption_probability, ) # only corrupt the audio if version_idx >= 1 # and both corrupt 50% of the time if version_idx >= 1: corrupted_audio = ( apply_preset( audio.clone(), self.sample_rate, preset, corruption_functions, ) if preset else audio.clone() ) else: corrupted_audio = audio.clone() # random crop to chunk size if self.random_crop: start_idx = np.random.randint( 0, corrupted_audio.shape[-1] - self.chunk_size_samples ) end_idx = start_idx + self.chunk_size_samples corrupted_audio_crop = corrupted_audio[..., start_idx:end_idx] audio_crop = audio[..., start_idx:end_idx] else: corrupted_audio_crop = corrupted_audio audio_crop = audio # lets clip anything out of range here # corrupted_audio_crop = torch.clamp(corrupted_audio_crop, -1, 1) # label_tensor = labels_to_tensor(labels, self.label_encoder) self.items_since_last_reload += 1 if np.random.uniform() < 0.5: # peak normalize corrupted_audio_crop = ( corrupted_audio_crop / corrupted_audio_crop.abs().max().clamp(1e-8) ) audio_crop = audio_crop / audio_crop.abs().max().clamp(1e-8) if np.random.uniform() < 0.5: gain_reduction_db = np.random.uniform(-12, 0) corrupted_audio_crop *= 10 ** (gain_reduction_db / 20.0) corrupted_audios.append(corrupted_audio_crop) if self.num_versions == 1: return corrupted_audios[0], mse_losses[0] # create a new label for which version has lower mse # mse_losses = torch.tensor(mse_losses) # label_tensor = mse_losses.argmin().float() label_tensor = torch.tensor([0.0]) # in this case always return 0 return torch.stack(corrupted_audios), label_tensor import torchaudio import numpy as np from concurrent.futures import ThreadPoolExecutor from tqdm import tqdm 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) class UpsampleAudioDataset(torch.utils.data.Dataset): def __init__( self, manifest_path: str, audio_dir: str, sample_rate: int, source_audio_manifest_path: str = None, chunk_size_s: float = 5.0, buffer_size: int = 1000, peak_normalize: bool = True, random_gain: bool = True, ): self.manifest_path = manifest_path self.audio_dir = audio_dir self.sample_rate = sample_rate self.chunk_size_s = chunk_size_s self.chunk_size_samples = int(chunk_size_s * sample_rate) self.buffer_size = buffer_size self.peak_normalize = peak_normalize self.random_gain = random_gain self.source_audio_manifest_path = source_audio_manifest_path with open(manifest_path, "r") as f: self.example_ids = [line.strip() for line in f.readlines()] print(f"Loaded {len(self.example_ids)} example ids from {manifest_path}") if self.source_audio_manifest_path is not None: with open(self.source_audio_manifest_path, "r") as f: self.source_audio_filepaths = [line.strip() for line in f.readlines()] print( f"Loaded {len(self.source_audio_filepaths)} source audio filepaths from {self.source_audio_manifest_path}" ) self.items_since_last_reload = self.buffer_size self.buffer = [] def __len__(self): return self.buffer_size def _reload_buffer(self): self.buffer = [] pbar = tqdm( np.random.permutation(len(self.example_ids)), total=len(self.example_ids), desc="Reloading buffer", ) for idx in pbar: if self.source_audio_manifest_path is not None and False: source_audio_filepath = np.random.choice(self.source_audio_filepaths) source_audio, sr = torchaudio.load(source_audio_filepath) # split into chunks source_audio_chunks = source_audio.unfold( -1, self.chunk_size_samples, self.chunk_size_samples ) for i in range(source_audio_chunks.shape[1]): source_audio_chunk = source_audio_chunks[:, i, :] if np.random.uniform() < 0.5: # get a random chunk from the source audio rand_idx = np.random.randint(0, source_audio_chunks.shape[1]) random_source_audio_chunk = source_audio_chunks[ :, rand_idx, : ].clone() else: random_source_audio_chunk = source_audio_chunk.clone() corrupted_source_audio_chunk = corrupt_audio( random_source_audio_chunk, self.sample_rate ) self.buffer.append( (corrupted_source_audio_chunk, source_audio_chunk, 1) ) example_id = self.example_ids[idx] # load the audio original_audio_path = os.path.join( self.audio_dir, f"{example_id}_original.mp3" ) codec_audio_path = os.path.join(self.audio_dir, f"{example_id}_codec.mp3") diff_ctx_audio_path = os.path.join( self.audio_dir, f"{example_id}_diff_ctx.mp3" ) diff_no_ctx_audio_path = os.path.join( self.audio_dir, f"{example_id}_diff_no_ctx.mp3" ) # load the audio codec_audio, sr = torchaudio.load(codec_audio_path) diff_ctx_audio, sr = torchaudio.load(diff_ctx_audio_path) diff_no_ctx_audio, sr = torchaudio.load(diff_no_ctx_audio_path) original_audio, sr = torchaudio.load(original_audio_path) # first crop out the first 30 seconds of the ctx_audio diff_ctx_audio = diff_ctx_audio[..., 30 * self.sample_rate :] # now we need to split the ctx_audio into num_chunks ctx_audio_chunks = diff_ctx_audio.unfold( -1, self.chunk_size_samples, self.chunk_size_samples ) no_ctx_audio_chunks = diff_no_ctx_audio.unfold( -1, self.chunk_size_samples, self.chunk_size_samples ) codec_audio_chunks = codec_audio.unfold( -1, self.chunk_size_samples, self.chunk_size_samples ) original_audio_chunks = original_audio.unfold( -1, self.chunk_size_samples, self.chunk_size_samples ) # Get the minimum number of chunks across all audio types to ensure alignment num_chunks = min( ctx_audio_chunks.shape[1], no_ctx_audio_chunks.shape[1], codec_audio_chunks.shape[1], original_audio_chunks.shape[1], ) # Construct all valid comparison pairs and add to buffer # Each pair is (worse_audio, better_audio, 1) where 1 indicates the second audio is better for i in range(num_chunks): j = i # use the same index for both # ctx is worse than no_ctx self.buffer.append( (ctx_audio_chunks[:, i, :], no_ctx_audio_chunks[:, j, :], 1) ) # ctx is worse than codec self.buffer.append( (ctx_audio_chunks[:, i, :], codec_audio_chunks[:, j, :], 1) ) # ctx is worse than original self.buffer.append( (ctx_audio_chunks[:, i, :], original_audio_chunks[:, j, :], 1) ) # no_ctx is worse than codec self.buffer.append( (no_ctx_audio_chunks[:, i, :], codec_audio_chunks[:, j, :], 1) ) # no_ctx is worse than original self.buffer.append( ( no_ctx_audio_chunks[:, i, :], original_audio_chunks[:, j, :], 1, ) ) # codec is worse than original self.buffer.append( (codec_audio_chunks[:, i, :], original_audio_chunks[:, j, :], 1) ) # Add corrupted original vs original comparison # Get the current original audio chunk # original_audio_chunk = original_audio_chunks[:, j, :] # Apply corruption to create a corrupted version # corrupted_audio_chunk = corrupt_audio( # original_audio_chunk.clone(), self.sample_rate # ) # Add the comparison: corrupted original is worse than original # self.buffer.append((corrupted_audio_chunk, original_audio_chunk, 1)) pbar.set_postfix({"buffer_size": len(self.buffer)}) if len(self.buffer) >= self.buffer_size: print(f"Buffer size reached {len(self.buffer)}/{self.buffer_size}") self.items_since_last_reload = 0 break def __getitem__(self, idx): if self.items_since_last_reload >= len(self.buffer): self._reload_buffer() self.items_since_last_reload += 1 audio_a, audio_b, label = self.buffer[idx] if self.peak_normalize: if np.random.uniform() < 0.2: audio_a = audio_a / audio_a.abs().max().clamp(1e-8) if np.random.uniform() < 0.2: audio_b = audio_b / audio_b.abs().max().clamp(1e-8) if self.random_gain: if np.random.uniform() < 0.2: gain_reduction_db = np.random.uniform(-12, 0) audio_a *= 10 ** (gain_reduction_db / 20.0) if np.random.uniform() < 0.2: gain_reduction_db = np.random.uniform(-12, 0) audio_b *= 10 ** (gain_reduction_db / 20.0) return audio_a, audio_b, label 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, 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) # 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 run_test(model, audio_pairs): model.eval() scores = [] scores_delta = [] assert len(audio_pairs) > 0, "No audio pairs found" for audio_a, audio_b 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) _, score_b = model.module.get_score(audio_b) score_a = score_a.item() score_b = score_b.item() # a is always the positive # b is always the negative scores.append(score_a > score_b) scores_delta.append(score_a - score_b) return scores, scores_delta def process_audio_pairs(csv_path: str): """ Process audio pairs from different test directories and return paired data. Args: csv_path: Path to the csv file containing the test data Returns: Dict of test types with audio pairs as values positive is first, negative is second """ tests = {} # load csv file df = pd.read_csv(csv_path) # find unique test types test_types = df["test_name"].unique() print(test_types) for test_type in test_types: # get all rows for this test type test_df = df[df["test_name"] == test_type] # get all unique audio pairs positive_files = test_df["positive"].tolist() negative_files = test_df["negative"].tolist() # combinbe into one list of tuples audio_filepath_pairs = list(zip(positive_files, negative_files)) # now we need to read the audio files audio_pairs = [] for positive_file, negative_file in tqdm(audio_filepath_pairs): audio_a = torchaudio.load(positive_file)[0] audio_b = torchaudio.load(negative_file)[0] audio_pairs.append((audio_a, audio_b)) tests[test_type] = audio_pairs return tests 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_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) 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_j - r_i) + (1 - labels) * F.logsigmoid(r_i - r_j) ) 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-compare-v6", "project_name": "ear-v2", "lr": 1e-5, "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": { # "corruptions_config": "/home/christian/code/christian/metadata/corruptions_config_v6.json", # "train_manifest": "/app/suno/data/audio_2ch_48khz_lg/ear_train_filtered_v2_with_gens.csv", # "val_manifest": "/app/suno/data/audio_2ch_48khz_lg/ear_val.csv", "train_manifest": "/home/christian/code/christian/metadata/ear/upsample_manifest_tr.txt", "val_manifest": "/home/christian/code/christian/metadata/ear/upsample_manifest_val.txt", "audio_dir": "/app/suno/christian/data/outputs/v45_2b_step_2_600_000/discogs_subset_sampled_metas", "train_source_audio_manifest": "/home/christian/code/christian/metadata/ear/ear_train_filtered_v1.csv", "val_source_audio_manifest": "/home/christian/code/christian/metadata/ear/ear_val.csv", "max_corruptions": 3, "no_corruption_probability": 0.0, "batch_size": 24, "num_workers": 1, "chunk_size_s": 5.0, "buffer_size": 100000, "sample_rate": 48_000, "random_crop": True, }, } # create label encoder # corruptions = load_corruptions_config(run_config["dataset"]["corruptions_config"]) # label_encoder = create_label_encoder(corruptions) # num_labels = len(label_encoder) # print(f"Number of labels: {num_labels}") # run_config["model"]["num_labels"] = num_labels # setup test audio pairs # only do this on the first process if local_rank == 0: test_audio_pairs_list = process_audio_pairs( "/home/christian/code/christian/metadata/ear/ear_bench_test_manifest.csv" ) 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_filepaths = run_config["dataset"]["train_manifest"] train_dataset = UpsampleAudioDataset( run_config["dataset"]["train_manifest"], run_config["dataset"]["audio_dir"], sample_rate=run_config["dataset"]["sample_rate"], source_audio_manifest_path=run_config["dataset"]["train_source_audio_manifest"], # num_workers=run_config["dataset"]["num_workers"], chunk_size_s=run_config["dataset"]["chunk_size_s"], buffer_size=run_config["dataset"]["buffer_size"], # max_corruptions=run_config["dataset"]["max_corruptions"], # no_corruption_probability=run_config["dataset"]["no_corruption_probability"], # num_versions=2 if run_config["model"]["objective"] == "bradley_terry" else 1, # 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_filepaths = run_config["dataset"]["val_manifest"] val_dataset = UpsampleAudioDataset( run_config["dataset"]["val_manifest"], run_config["dataset"]["audio_dir"], sample_rate=run_config["dataset"]["sample_rate"], source_audio_manifest_path=run_config["dataset"]["val_source_audio_manifest"], # 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_corruptions=run_config["dataset"]["max_corruptions"], # no_corruption_probability=run_config["dataset"]["no_corruption_probability"], # num_versions=2 if run_config["model"]["objective"] == "bradley_terry" else 1, # 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] 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) 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 name, test_audio_pairs in test_audio_pairs_list.items(): 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!")