import os import json import argparse import random from pathlib import Path from datetime import datetime import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from torch.utils.tensorboard import SummaryWriter from torch.cuda.amp import autocast, GradScaler import torchaudio import librosa from tqdm import tqdm from audiotools import AudioSignal, STFTParams # Import from the provided codebase from higgs_audio_tokenizer import HiggsAudioTokenizer from quantization.distrib import broadcast_tensors, sync_buffer, is_distributed, world_size, rank from quantization.ddp_utils import set_random_seed, is_logging_process, get_timestamp # Import DAC losses and discriminator import sys sys.path.append('.') # Add current directory to path from loss import L1Loss, MultiScaleSTFTLoss, MelSpectrogramLoss, GANLoss from discriminator import Discriminator class CosineWarmupScheduler(torch.optim.lr_scheduler._LRScheduler): """Cosine scheduler with linear warmup""" def __init__(self, optimizer, warmup_steps, total_steps, eta_min=1e-6, last_epoch=-1): self.warmup_steps = warmup_steps self.total_steps = total_steps self.eta_min = eta_min super().__init__(optimizer, last_epoch) def get_lr(self): if self.last_epoch < self.warmup_steps: # Linear warmup warmup_factor = self.last_epoch / self.warmup_steps return [base_lr * warmup_factor for base_lr in self.base_lrs] else: # Cosine annealing progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps) cosine_factor = 0.5 * (1 + np.cos(np.pi * progress)) return [self.eta_min + (base_lr - self.eta_min) * cosine_factor for base_lr in self.base_lrs] class AudioDataset(Dataset): """Dataset for loading audio files from CSV""" def __init__(self, csv_path, sample_rate=44100, segment_duration=2.0, is_train=True): self.df = pd.read_csv(csv_path) self.sample_rate = sample_rate self.segment_duration = segment_duration self.segment_length = int(sample_rate * segment_duration) self.is_train = is_train # Filter out files that don't exist valid_files = [] for idx, row in self.df.iterrows(): if os.path.exists(row.iloc[0]): valid_files.append(row.iloc[0]) self.audio_paths = valid_files print(f"Found {len(self.audio_paths)} valid audio files") def __len__(self): return len(self.audio_paths) def __getitem__(self, idx): audio_path = self.audio_paths[idx] try: audio, sr = librosa.load(audio_path, sr=self.sample_rate, mono=True) = if len(audio) > self.segment_length: if self.is_train: start = random.randint(0, len(audio) - self.segment_length) else: start = 0 = audio = audio[start:start + self.segment_length] else: # Pad if too short audio = np.pad(audio, (0, self.segment_length - len(audio))) audio_tensor = torch.FloatTensor(audio).unsqueeze(0) return audio_tensor, audio_path except Exception as e: print(f"Error loading {audio_path}: {e}") # Return silence if loading fails return torch.zeros(1, self.segment_length), audio_path class BosonTrainer: def __init__(self, args): self.args = args self.distributed = False # Check if we're in a distributed environment if 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE']) > 1: self.distributed = True self.setup_ddp() self.device = torch.device(f'cuda:{args.local_rank}') else: # Single GPU mode self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.cuda.set_device(0) set_random_seed(args.seed) # Load config with open(args.config, 'r') as f: self.config = json.load(f) # Initialize models self.model = self.build_model() self.discriminator = self.build_discriminator() if args.use_discriminator else None # Setup data loaders self.train_loader, self.val_loader = self.setup_data_loaders() # Setup optimizers self.optimizer_g = torch.optim.AdamW( self.model.parameters(), lr=args.learning_rate, betas=(0.5, 0.9), weight_decay=args.weight_decay ) if self.discriminator is not None: self.optimizer_d = torch.optim.AdamW( self.discriminator.parameters(), lr=args.learning_rate * 2, # Typically discriminator learns faster betas=(0.5, 0.9), weight_decay=args.weight_decay ) # Initialize gradient scalers for mixed precision if args.use_mixed_precision: self.scaler_g = GradScaler() self.scaler_d = GradScaler() if self.discriminator is not None else None else: self.scaler_g = None self.scaler_d = None # Calculate total training steps self.total_steps = args.num_epochs * len(self.train_loader) # Setup schedulers with warmup self.scheduler_g = CosineWarmupScheduler( self.optimizer_g, warmup_steps=args.warmup_steps, total_steps=self.total_steps, eta_min=1e-6 ) if self.discriminator is not None: self.scheduler_d = CosineWarmupScheduler( self.optimizer_d, warmup_steps=args.warmup_steps, total_steps=self.total_steps, eta_min=1e-6 ) # Setup losses self.setup_losses() # Setup tensorboard if not self.distributed or rank() == 0: self.writer = SummaryWriter( log_dir=os.path.join(args.output_dir, 'logs', get_timestamp()) ) self.global_step = 0 self.start_epoch = 0 # Load checkpoint if exists if args.resume: self.load_checkpoint() def setup_ddp(self): """Initialize DDP""" if 'LOCAL_RANK' in os.environ: self.args.local_rank = int(os.environ['LOCAL_RANK']) dist.init_process_group(backend='nccl') torch.cuda.set_device(self.args.local_rank) set_random_seed(self.args.seed + rank()) def build_model(self): """Build and wrap model with DDP if needed""" print(self.config) model = HiggsAudioTokenizer( n_filters=self.config['n_filters'], D=self.config['D'], target_bandwidths=self.config['target_bandwidths'], ratios=self.config['ratios'], sample_rate=self.config['sample_rate'], bins=self.config['bins'], n_q=self.config['n_q'], codebook_dim=self.config.get('codebook_dim', None), semantic_techer=self.config['semantic_techer'], device=self.device ).to(self.device) if self.distributed: # Broadcast model parameters to ensure all ranks have same initialization broadcast_tensors(model.parameters()) # Wrap with DDP model = DDP(model, device_ids=[self.args.local_rank]) return model # def build_discriminator(self): # """Build discriminator with DDP if needed""" # # Use sample rate from config # discriminator = Discriminator( # rates=[], # No multi-rate discriminator for now # periods=[2, 3, 5, 7, 11], # fft_sizes=[2048, 1024, 512], # sample_rate=self.config['sample_rate'], # ).to(self.device) # if self.distributed: # broadcast_tensors(discriminator.parameters()) # discriminator = DDP(discriminator, device_ids=[self.args.local_rank]) # return discriminator def build_discriminator(self): discriminator = Discriminator( rates=[], # No multi-rate discriminator periods=[2, 3, 5, 7, 11], fft_sizes=[2048, 1024, 512], sample_rate=self.config['sample_rate'], # 44100 ).to(self.device) if self.distributed: broadcast_tensors(discriminator.parameters()) discriminator = DDP(discriminator, device_ids=[self.args.local_rank]) return discriminator def setup_losses(self): # Basic losses self.l1_loss = L1Loss() self.stft_loss = MultiScaleSTFTLoss( window_lengths=[2048, 1024, 512, 256, 128], loss_fn=nn.L1Loss(), clamp_eps=1e-5, mag_weight=1.0, log_weight=1.0, ) self.mel_loss = MelSpectrogramLoss( n_mels=[150, 80], window_lengths=[2048, 512], mel_fmin=[0.0, 0.0], mel_fmax=[None, None], clamp_eps=1e-5, mag_weight=1.0, log_weight=1.0, ) if self.discriminator is not None: self.gan_loss = GANLoss(self.discriminator) self.loss_weights = { 'rec': 1., # Waveform L1 loss 'stft': 1., # Multi-scale STFT loss 'mel': 45.0, # Mel-spectrogram loss 'commit': 0.25, # Commitment loss 'semantic': 1., # Semantic loss 'gen': 1., # Generator adversarial loss 'feat': 2.0, # Feature matching loss } def setup_data_loaders(self): # Split data into train/val df = pd.read_csv(self.args.data_csv) n_total = len(df) n_train = int(n_total * 0.9) # Create temporary CSV files for train/val split train_csv = '/tmp/train_audio.csv' val_csv = '/tmp/val_audio.csv' if not self.distributed or rank() == 0: df[:n_train].to_csv(train_csv, index=False) df[n_train:].to_csv(val_csv, index=False) if self.distributed: dist.barrier() # Create datasets train_dataset = AudioDataset( train_csv, sample_rate=self.config['sample_rate'], segment_duration=self.args.segment_duration, is_train=True ) val_dataset = AudioDataset( val_csv, sample_rate=self.config['sample_rate'], segment_duration=self.args.segment_duration, is_train=False ) # Create samplers and loaders if self.distributed: train_sampler = DistributedSampler(train_dataset, shuffle=True) val_sampler = DistributedSampler(val_dataset, shuffle=False) else: train_sampler = None val_sampler = None train_loader = DataLoader( train_dataset, batch_size=self.args.batch_size, sampler=train_sampler, shuffle=(train_sampler is None), num_workers=self.args.num_workers, pin_memory=True, drop_last=True ) val_loader = DataLoader( val_dataset, batch_size=self.args.batch_size, sampler=val_sampler, shuffle=False, num_workers=self.args.num_workers, pin_memory=True, drop_last=False ) return train_loader, val_loader def is_main_process(self): """Check if this is the main process""" return not self.distributed or rank() == 0 def train_epoch(self, epoch): """Train for one epoch""" self.model.train() if self.discriminator is not None: self.discriminator.train() if self.distributed: self.train_loader.sampler.set_epoch(epoch) total_losses = { 'total': 0, 'rec': 0, 'stft': 0, 'mel': 0, 'commit': 0, 'semantic': 0, 'gen': 0, 'feat': 0, 'disc': 0 } pbar = tqdm(self.train_loader, desc=f'Epoch {epoch}', disable=not self.is_main_process()) for batch_idx, (audio, paths) in enumerate(pbar): audio = audio.to(self.device) # Create AudioSignal objects for loss computation audio_signal = AudioSignal(audio, self.config['sample_rate']) # Forward pass with random bandwidth bw_idx = random.randint(0, len(self.config['target_bandwidths']) - 1) bw = self.config['target_bandwidths'][bw_idx] # Use autocast for mixed precision with autocast(dtype=torch.bfloat16, enabled=self.args.use_mixed_precision): output, commit_loss, semantic_loss, _ = self.model(audio, bw) recons_signal = AudioSignal(output, self.config['sample_rate']) use_discriminator = (self.discriminator is not None and self.global_step >= self.args.discriminator_start_step) if use_discriminator and self.global_step % self.args.disc_interval == 0: self.optimizer_d.zero_grad() with autocast(dtype=torch.bfloat16, enabled=self.args.use_mixed_precision): disc_loss = self.gan_loss.discriminator_loss(recons_signal, audio_signal) if self.scaler_d is not None: self.scaler_d.scale(disc_loss).backward() self.scaler_d.unscale_(self.optimizer_d) torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), 10.0) self.scaler_d.step(self.optimizer_d) self.scaler_d.update() else: disc_loss.backward() torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), 10.0) self.optimizer_d.step() self.scheduler_d.step() total_losses['disc'] += disc_loss.item() # Train generator losses = {} # Compute losses with autocast with autocast(dtype=torch.bfloat16, enabled=self.args.use_mixed_precision): # Reconstruction losses losses['rec'] = self.l1_loss(recons_signal, audio_signal) losses['stft'] = self.stft_loss(recons_signal, audio_signal) losses['mel'] = self.mel_loss(recons_signal, audio_signal) # losses['mel'] = torch.tensor(0.0, device=self.device) # uncomment this for the first 30k steps, it's faster if you pretrain it on semantic / commit loss first losses['commit'] = commit_loss losses['semantic'] = semantic_loss # GAN losses if discriminator is active if use_discriminator: gen_loss, feat_loss = self.gan_loss.generator_loss(recons_signal, audio_signal) losses['gen'] = gen_loss losses['feat'] = feat_loss else: # Set to zero for logging purposes losses['gen'] = torch.tensor(0.0, device=self.device) losses['feat'] = torch.tensor(0.0, device=self.device) # Total weighted loss total_loss = sum(self.loss_weights.get(k, 0) * v for k, v in losses.items() if k not in ['gen', 'feat'] or use_discriminator) # Backward pass self.optimizer_g.zero_grad() if self.scaler_g is not None: self.scaler_g.scale(total_loss).backward() self.scaler_g.unscale_(self.optimizer_g) torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.scaler_g.step(self.optimizer_g) self.scaler_g.update() else: total_loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) self.optimizer_g.step() self.scheduler_g.step() # Update metrics total_losses['total'] += total_loss.item() for k, v in losses.items(): total_losses[k] += v.item() # Update progress bar if self.is_main_process(): pbar.set_postfix({ 'loss': f'{total_loss.item():.4f}', 'rec': f'{losses["rec"].item():.4f}', 'mel': f'{losses["mel"].item():.4f}', 'commit_loss': f'{losses["commit"].item():.4f}', 'semantic_loss': f'{losses["semantic"].item():.4f}', 'lr': f'{self.scheduler_g.get_last_lr()[0]:.9f}', 'disc': 'ON' if use_discriminator else 'OFF', 'step': self.global_step }) # Log to tensorboard if self.is_main_process() and self.global_step % self.args.log_interval == 0: for k, v in losses.items(): self.writer.add_scalar(f'train/{k}_loss', v.item(), self.global_step) self.writer.add_scalar('train/total_loss', total_loss.item(), self.global_step) self.writer.add_scalar('train/lr', self.scheduler_g.get_last_lr()[0], self.global_step) self.writer.add_scalar('train/bandwidth', bw, self.global_step) self.writer.add_scalar('train/discriminator_active', float(use_discriminator), self.global_step) if use_discriminator: self.writer.add_scalar('train/disc_loss', total_losses['disc'] / max(1, batch_idx), self.global_step) if self.scaler_g is not None: self.writer.add_scalar('train/grad_scale', self.scaler_g.get_scale(), self.global_step) # Save checkpoint at step intervals if self.global_step > 0 and self.global_step % self.args.save_step_interval == 0: self.save_checkpoint_step(self.global_step) if self.is_main_process(): print(f"\nSaved checkpoint at step {self.global_step}") self.global_step += 1 # Return average losses n_batches = len(self.train_loader) return {k: v / n_batches for k, v in total_losses.items()} @torch.no_grad() def validate(self, epoch): """Validation loop""" self.model.eval() total_losses = { 'total': 0, 'rec': 0, 'stft': 0, 'mel': 0, 'commit': 0, 'semantic': 0 } audio_samples = {'train': [], 'val': []} for batch_idx, (audio, paths) in enumerate(tqdm(self.val_loader, desc='Validation', disable=not self.is_main_process())): audio = audio.to(self.device) audio_signal = AudioSignal(audio, self.config['sample_rate']) # Use medium bandwidth for validation bw = self.config['target_bandwidths'][2] # Use autocast for validation too with autocast(dtype=torch.bfloat16, enabled=self.args.use_mixed_precision): output, commit_loss, semantic_loss, _ = self.model(audio, bw) recons_signal = AudioSignal(output, self.config['sample_rate']) # Compute losses losses = { 'rec': self.l1_loss(recons_signal, audio_signal), 'stft': self.stft_loss(recons_signal, audio_signal), 'mel': self.mel_loss(recons_signal, audio_signal), 'commit': commit_loss, 'semantic': semantic_loss } total_loss = sum(self.loss_weights.get(k, 0) * v for k, v in losses.items()) total_losses['total'] += total_loss.item() for k, v in losses.items(): total_losses[k] += v.item() # Collect audio samples for tensorboard (first 3 from validation) if self.is_main_process() and len(audio_samples['val']) < 3: audio_samples['val'].append({ 'original': audio[0].cpu(), 'reconstructed': output[0].cpu(), 'path': paths[0] }) # Get train samples for comparison if self.is_main_process(): self.model.eval() for batch_idx, (audio, paths) in enumerate(self.train_loader): if len(audio_samples['train']) >= 3: break audio = audio.to(self.device) bw = self.config['target_bandwidths'][2] with autocast(dtype=torch.bfloat16, enabled=self.args.use_mixed_precision): output, _, _, _ = self.model(audio, bw) audio_samples['train'].append({ 'original': audio[0].cpu(), 'reconstructed': output[0].cpu(), 'path': paths[0] }) # Log audio samples to tensorboard if self.is_main_process(): for split in ['train', 'val']: for idx, sample in enumerate(audio_samples[split]): self.writer.add_audio( f'{split}/original_{idx}', sample['original'], epoch, sample_rate=self.config['sample_rate'] ) self.writer.add_audio( f'{split}/reconstructed_{idx}', sample['reconstructed'], epoch, sample_rate=self.config['sample_rate'] ) # Average losses n_batches = len(self.val_loader) val_metrics = {k: v / n_batches for k, v in total_losses.items()} # Log validation metrics if self.is_main_process(): for key, value in val_metrics.items(): self.writer.add_scalar(f'val/{key}_loss', value, epoch) return val_metrics def save_checkpoint(self, epoch, is_best=False): """Save model checkpoint (epoch-based)""" if not self.is_main_process(): return model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict() # Get current learning rates for verification current_lr_g = self.scheduler_g.get_last_lr()[0] checkpoint = { 'epoch': epoch, 'global_step': self.global_step, 'model_state_dict': model_state, 'optimizer_g_state_dict': self.optimizer_g.state_dict(), 'scheduler_g_state_dict': self.scheduler_g.state_dict(), 'scheduler_g_last_epoch': self.scheduler_g.last_epoch, # Explicitly save this 'current_lr_g': current_lr_g, # Save for verification 'config': self.config, 'args': self.args } # Save gradient scaler states if using mixed precision if self.scaler_g is not None: checkpoint['scaler_g_state_dict'] = self.scaler_g.state_dict() if self.discriminator is not None: disc_state = self.discriminator.module.state_dict() if self.distributed else self.discriminator.state_dict() current_lr_d = self.scheduler_d.get_last_lr()[0] checkpoint['discriminator_state_dict'] = disc_state checkpoint['optimizer_d_state_dict'] = self.optimizer_d.state_dict() checkpoint['scheduler_d_state_dict'] = self.scheduler_d.state_dict() checkpoint['scheduler_d_last_epoch'] = self.scheduler_d.last_epoch checkpoint['current_lr_d'] = current_lr_d if self.scaler_d is not None: checkpoint['scaler_d_state_dict'] = self.scaler_d.state_dict() # Save latest checkpoint checkpoint_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth') os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) torch.save(checkpoint, checkpoint_path) # Save best checkpoint if is_best: best_path = os.path.join(self.args.output_dir, 'checkpoints', 'best.pth') torch.save(checkpoint, best_path) # Save periodic checkpoint if epoch % self.args.save_interval == 0: epoch_path = os.path.join(self.args.output_dir, 'checkpoints', f'epoch_{epoch}.pth') torch.save(checkpoint, epoch_path) def save_checkpoint_step(self, step): """Save model checkpoint (step-based)""" if not self.is_main_process(): return # Get current epoch from training loop current_epoch = step // len(self.train_loader) model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict() # Get current learning rates for verification current_lr_g = self.scheduler_g.get_last_lr()[0] checkpoint = { 'epoch': current_epoch, 'global_step': step, 'model_state_dict': model_state, 'optimizer_g_state_dict': self.optimizer_g.state_dict(), 'scheduler_g_state_dict': self.scheduler_g.state_dict(), 'scheduler_g_last_epoch': self.scheduler_g.last_epoch, # Explicitly save this 'current_lr_g': current_lr_g, # Save for verification 'config': self.config, 'args': self.args } # Save gradient scaler states if using mixed precision if self.scaler_g is not None: checkpoint['scaler_g_state_dict'] = self.scaler_g.state_dict() if self.discriminator is not None: disc_state = self.discriminator.module.state_dict() if self.distributed else self.discriminator.state_dict() current_lr_d = self.scheduler_d.get_last_lr()[0] checkpoint['discriminator_state_dict'] = disc_state checkpoint['optimizer_d_state_dict'] = self.optimizer_d.state_dict() checkpoint['scheduler_d_state_dict'] = self.scheduler_d.state_dict() checkpoint['scheduler_d_last_epoch'] = self.scheduler_d.last_epoch checkpoint['current_lr_d'] = current_lr_d if self.scaler_d is not None: checkpoint['scaler_d_state_dict'] = self.scaler_d.state_dict() # Create checkpoint directory if it doesn't exist checkpoint_dir = os.path.join(self.args.output_dir, 'checkpoints') os.makedirs(checkpoint_dir, exist_ok=True) # Save step-based checkpoint step_path = os.path.join(self.args.output_dir, 'checkpoints', f'step_{step}.pth') torch.save(checkpoint, step_path) # Also update latest checkpoint latest_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth') torch.save(checkpoint, latest_path) # Keep only the last N step-based checkpoints to save disk space if self.args.keep_last_n_steps > 0: checkpoint_dir = os.path.join(self.args.output_dir, 'checkpoints') step_checkpoints = sorted([f for f in os.listdir(checkpoint_dir) if f.startswith('step_')]) if len(step_checkpoints) > self.args.keep_last_n_steps: for old_checkpoint in step_checkpoints[:-self.args.keep_last_n_steps]: os.remove(os.path.join(checkpoint_dir, old_checkpoint)) def load_checkpoint(self): checkpoint_path = os.path.join(self.args.output_dir, 'checkpoints', 'latest.pth') if os.path.exists(checkpoint_path): print(f"Loading checkpoint from {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False) # Load model state if self.distributed: self.model.module.load_state_dict(checkpoint['model_state_dict']) else: self.model.load_state_dict(checkpoint['model_state_dict']) # Load optimizer state self.optimizer_g.load_state_dict(checkpoint['optimizer_g_state_dict']) # Load scheduler state self.scheduler_g.load_state_dict(checkpoint['scheduler_g_state_dict']) # Restore scheduler's last_epoch from checkpoint if 'scheduler_g_last_epoch' in checkpoint: self.scheduler_g.last_epoch = checkpoint['scheduler_g_last_epoch'] else: self.scheduler_g.last_epoch = checkpoint['global_step'] # Force scheduler to recompute its internal state self.scheduler_g._last_lr = self.scheduler_g.get_lr() # Load gradient scaler state if using mixed precision if self.scaler_g is not None and 'scaler_g_state_dict' in checkpoint: self.scaler_g.load_state_dict(checkpoint['scaler_g_state_dict']) # Load discriminator if present if self.discriminator is not None and 'discriminator_state_dict' in checkpoint: if self.distributed: self.discriminator.module.load_state_dict(checkpoint['discriminator_state_dict']) else: self.discriminator.load_state_dict(checkpoint['discriminator_state_dict']) self.optimizer_d.load_state_dict(checkpoint['optimizer_d_state_dict']) self.scheduler_d.load_state_dict(checkpoint['scheduler_d_state_dict']) # Restore discriminator scheduler's last_epoch if 'scheduler_d_last_epoch' in checkpoint: self.scheduler_d.last_epoch = checkpoint['scheduler_d_last_epoch'] else: self.scheduler_d.last_epoch = checkpoint['global_step'] self.scheduler_d._last_lr = self.scheduler_d.get_lr() if self.scaler_d is not None and 'scaler_d_state_dict' in checkpoint: self.scaler_d.load_state_dict(checkpoint['scaler_d_state_dict']) # Restore training state self.start_epoch = checkpoint['epoch'] + 1 self.global_step = checkpoint['global_step'] # Verify learning rate restoration current_lr_g = self.scheduler_g.get_last_lr()[0] saved_lr_g = checkpoint.get('current_lr_g', None) print(f"\n{'='*60}") print(f"CHECKPOINT LOADED SUCCESSFULLY") print(f"{'='*60}") print(f"Resumed from epoch: {checkpoint['epoch']}") print(f"Global step: {self.global_step}") print(f"Scheduler last_epoch: {self.scheduler_g.last_epoch}") print(f"Current learning rate (generator): {current_lr_g:.9f}") print(f"Mixed precision: {'ENABLED' if self.args.use_mixed_precision else 'DISABLED'}") if saved_lr_g is not None: print(f"Saved learning rate (generator): {saved_lr_g:.9f}") if abs(current_lr_g - saved_lr_g) > 1e-9: print("⚠️ WARNING: Learning rate mismatch! This might indicate improper state restoration.") if self.discriminator is not None: current_lr_d = self.scheduler_d.get_last_lr()[0] saved_lr_d = checkpoint.get('current_lr_d', None) print(f"Current learning rate (discriminator): {current_lr_d:.9f}") if saved_lr_d is not None: print(f"Saved learning rate (discriminator): {saved_lr_d:.9f}") print(f"Discriminator status: {'ACTIVE' if self.global_step >= self.args.discriminator_start_step else f'INACTIVE (starts at step {self.args.discriminator_start_step})'}") print(f"Next epoch: {self.start_epoch}") print(f"Next step checkpoint at: step {((self.global_step // self.args.save_step_interval) + 1) * self.args.save_step_interval}") print(f"{'='*60}\n") g if self.global_step > 0: temp_scheduler = CosineWarmupScheduler( self.optimizer_g, self.args.warmup_steps, self.total_steps, eta_min=1e-6, last_epoch=-1 ) # Step it to the current global step for _ in range(self.global_step): temp_scheduler.step() expected_lr = temp_scheduler.get_last_lr()[0] if abs(current_lr_g - expected_lr) > 1e-9: print(f"⚠️ Learning rate verification failed!") print(f" Expected: {expected_lr:.9f}") print(f" Got: {current_lr_g:.9f}") print(" The scheduler state might not be properly restored.") else: print(f"No checkpoint found at {checkpoint_path}, starting from scratch") def train(self): """Main training loop""" best_val_loss = float('inf') # Print training configuration if self.is_main_process(): print(f"\n{'='*50}") print(f"Training Configuration:") print(f"{'='*50}") print(f"Total epochs: {self.args.num_epochs}") print(f"Steps per epoch: {len(self.train_loader)}") print(f"Total steps: {self.total_steps}") print(f"Warmup steps: {self.args.warmup_steps}") print(f"Mixed precision training: {'ENABLED (bfloat16)' if self.args.use_mixed_precision else 'DISABLED'}") print(f"Discriminator starts at step: {self.args.discriminator_start_step}") print(f"Checkpoint saving:") print(f" - Every {self.args.save_interval} epochs") print(f" - Every {self.args.save_step_interval} steps") print(f" - Keep last {self.args.keep_last_n_steps} step checkpoints") if self.start_epoch > 0: print(f"RESUMING from epoch {self.start_epoch}, step {self.global_step}") print(f"{'='*50}\n") for epoch in range(self.start_epoch, self.args.num_epochs): # IMPORTANT: Set the epoch for distributed sampler when resuming # This ensures proper data shuffling across epochs if self.distributed and hasattr(self.train_loader.sampler, 'set_epoch'): self.train_loader.sampler.set_epoch(epoch) # Train train_metrics = self.train_epoch(epoch) # Validate val_metrics = self.validate(epoch) # Log epoch metrics if self.is_main_process(): print(f"\nEpoch {epoch} Summary:") print(f"Train - Total: {train_metrics['total']:.4f}, Rec: {train_metrics['rec']:.4f}, " f"STFT: {train_metrics['stft']:.4f}, Mel: {train_metrics['mel']:.4f}, " f"Commit: {train_metrics['commit']:.4f}, Semantic: {train_metrics['semantic']:.4f}") if self.discriminator is not None: print(f" Gen: {train_metrics['gen']:.4f}, Feat: {train_metrics['feat']:.4f}, " f"Disc: {train_metrics['disc']:.4f}") print(f" Discriminator Status: {'Active' if self.global_step >= self.args.discriminator_start_step else f'Starting at step {self.args.discriminator_start_step}'}") print(f"Val - Total: {val_metrics['total']:.4f}, Rec: {val_metrics['rec']:.4f}, " f"STFT: {val_metrics['stft']:.4f}, Mel: {val_metrics['mel']:.4f}, " f"Commit: {val_metrics['commit']:.4f}, Semantic: {val_metrics['semantic']:.4f}") print(f"Current Step: {self.global_step}, Next step checkpoint at: {((self.global_step // self.args.save_step_interval) + 1) * self.args.save_step_interval}") print(f"Current LR: {self.scheduler_g.get_last_lr()[0]:.9f}") # Save checkpoint is_best = val_metrics['total'] < best_val_loss if is_best: best_val_loss = val_metrics['total'] self.save_checkpoint(epoch, is_best) # Save final model if self.is_main_process(): model_state = self.model.module.state_dict() if self.distributed else self.model.state_dict() final_path = os.path.join(self.args.output_dir, 'checkpoints', 'final.pth') torch.save({ 'model_state_dict': model_state, 'config': self.config }, final_path) # Also save just the model weights in the format expected by the original code model_only_path = os.path.join(self.args.output_dir, 'model.pth') torch.save(model_state, model_only_path) # Copy config import shutil shutil.copy(self.args.config, os.path.join(self.args.output_dir, 'config.json')) # Cleanup if self.is_main_process(): self.writer.close() if self.distributed: dist.destroy_process_group() def main(): parser = argparse.ArgumentParser(description='Train Boson Audio Codec') # Data arguments parser.add_argument('--data_csv', type=str, required=True, help='Path to CSV file containing audio file paths') parser.add_argument('--config', type=str, default='config.json', help='Path to config JSON file') # Training arguments parser.add_argument('--batch_size', type=int, default=28, help='Batch size per GPU') parser.add_argument('--num_epochs', type=int, default=100, help='Number of training epochs') parser.add_argument('--learning_rate', type=float, default=1e-4, help='Initial learning rate') parser.add_argument('--weight_decay', type=float, default=0.01, help='Weight decay') parser.add_argument('--segment_duration', type=float, default=2., help='Audio segment duration in seconds') # Mixed precision training parser.add_argument('--use_mixed_precision', action='store_true', help='Use bfloat16 mixed precision training') # Scheduler arguments parser.add_argument('--warmup_steps', type=int, default=5000, help='Number of warmup steps for cosine scheduler') # Loss arguments parser.add_argument('--use_discriminator', action='store_true', help='Use adversarial training with discriminator') parser.add_argument('--discriminator_start_step', type=int, default=25_000, help='Start training discriminator after N steps') parser.add_argument('--disc_interval', type=int, default=1, help='Train discriminator every N steps') # System arguments parser.add_argument('--output_dir', type=str, default='outputs_mp_cqt', help='Output directory for checkpoints and logs') parser.add_argument('--num_workers', type=int, default=16, help='Number of data loading workers') parser.add_argument('--seed', type=int, default=42, help='Random seed') parser.add_argument('--local_rank', type=int, default=0, help='Local rank for distributed training') # Logging arguments parser.add_argument('--log_interval', type=int, default=10, help='Log every N steps') parser.add_argument('--save_interval', type=int, default=1, help='Save checkpoint every N epochs') parser.add_argument('--save_step_interval', type=int, default=1000, help='Save checkpoint every N steps') parser.add_argument('--keep_last_n_steps', type=int, default=5, help='Keep only the last N step-based checkpoints (0 to keep all)') # Resume training parser.add_argument('--resume', action='store_true', help='Resume training from latest checkpoint') args = parser.parse_args() # Create output directory os.makedirs(args.output_dir, exist_ok=True) # Train trainer = BosonTrainer(args) trainer.train() if __name__ == '__main__': torch.set_float32_matmul_precision('high') main()