# train_codec_mel_speaker.py # ============================================================================== # Training script for Mel-based Codec with LEARNABLE SPEAKER EMBEDDINGS # # KEY FEATURES: # - Input: mel spectrogram + pitch (NO text embeddings, NO energy) # - Uses learnable speaker embeddings: nn.Embedding(num_speakers, speaker_dim) # - Speaker conditioning via AdaIN1d throughout decoder # - F0 prediction from latent representation # Architecture: # MelPitchEncoder: (mel, pitch) -> latent # FSQ: latent -> tokens (spectral + prosody quantized) # Decoder: (quantized_latent, speaker_emb) -> waveform # F0 Predictor: latent -> pitch prediction # Usage: # python train_codec_mel_speaker.py -p Configs/config_codec_mel_speaker.yml # ============================================================================== import os import os.path as osp import sys import yaml import shutil import numpy as np import torch import click import warnings warnings.simplefilter('ignore') import random from munch import Munch from torch import nn import torch.nn.functional as F import torchaudio from torch.nn.utils.rnn import pad_sequence from models_mel_speaker import build_model_mel_speaker, load_checkpoint from meldataset import build_dataloader from utils import get_data_path_list, log_print, length_to_mask, recursive_munch from losses import MultiResolutionSTFTLoss, MagPhaseLoss, GeneratorLoss, DiscriminatorLoss, WavLMLoss from optimizers import build_optimizer import time from accelerate import Accelerator from accelerate.utils import LoggerType from accelerate import DistributedDataParallelKwargs from torch.utils.tensorboard import SummaryWriter import logging from accelerate.logging import get_logger logger = get_logger(__name__, log_level="DEBUG") # Fix "Too many open files" error with DataLoader workers import torch.multiprocessing torch.multiprocessing.set_sharing_strategy('file_system') # Increase file descriptor limit import resource soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (min(65536, hard), hard)) import gc # ============================================================================== # Helper function to expand phoneme durations to frame-level targets # ============================================================================== def expand_durations_to_frames(phonemes, durations, max_frames): """ Expand phoneme IDs by their durations to get frame-level phoneme targets. Args: phonemes: [B, T_phonemes] - phoneme IDs (with padding) durations: [B, T_phonemes] - duration for each phoneme (in frames at hop_length=441) max_frames: int - maximum number of frames to output Returns: expanded: [B, max_frames] - phoneme ID at each frame mask: [B, max_frames] - valid frames mask (1 where valid, 0 where padded) """ B = phonemes.shape[0] device = phonemes.device expanded = torch.zeros(B, max_frames, dtype=torch.long, device=device) mask = torch.zeros(B, max_frames, dtype=torch.float, device=device) for b in range(B): frame_idx = 0 for p_idx in range(phonemes.shape[1]): dur = int(durations[b, p_idx].item()) if dur <= 0: continue phoneme_id = phonemes[b, p_idx].item() end_idx = min(frame_idx + dur, max_frames) if frame_idx >= max_frames: break expanded[b, frame_idx:end_idx] = phoneme_id mask[b, frame_idx:end_idx] = 1.0 frame_idx = end_idx return expanded, mask def crop_phonemes_and_durations(phonemes, durations, start_frame, end_frame): """ Crop phonemes and durations to match a cropped mel region. Args: phonemes: [T_phonemes] - phoneme IDs (with padding at start and end) durations: [T_phonemes] - duration for each phoneme in frames start_frame: int - start frame of the crop end_frame: int - end frame of the crop (exclusive) Returns: cropped_phonemes: [T_cropped] - phoneme IDs in the cropped region cropped_durations: [T_cropped] - durations in the cropped region """ # Find which phonemes fall in the cropped region frame_pos = 0 cropped_phonemes = [] cropped_durations = [] for i in range(len(phonemes)): dur = int(durations[i].item()) if dur <= 0: continue phone_start = frame_pos phone_end = frame_pos + dur # Check if this phoneme overlaps with our crop region if phone_end > start_frame and phone_start < end_frame: # Calculate how much of this phoneme is in the crop region overlap_start = max(phone_start, start_frame) overlap_end = min(phone_end, end_frame) cropped_dur = overlap_end - overlap_start if cropped_dur > 0: cropped_phonemes.append(phonemes[i].item()) cropped_durations.append(cropped_dur) frame_pos = phone_end # Early exit if we've passed the crop region if frame_pos >= end_frame: break return cropped_phonemes, cropped_durations @click.command() @click.option('-p', '--config_path', default='Configs/config_codec_mel_speaker.yml', type=str) def main(config_path): config = yaml.safe_load(open(config_path)) log_dir = config['log_dir'] if not osp.exists(log_dir): os.makedirs(log_dir, exist_ok=True) shutil.copy(config_path, osp.join(log_dir, osp.basename(config_path))) ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( project_dir=log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs], mixed_precision='bf16' ) if accelerator.is_main_process: writer = SummaryWriter(log_dir + "/tensorboard") file_handler = logging.FileHandler(osp.join(log_dir, 'train_codec_mel_speaker.log')) file_handler.setLevel(logging.DEBUG) file_handler.setFormatter(logging.Formatter('%(levelname)s:%(asctime)s: %(message)s')) logger.logger.addHandler(file_handler) # Config batch_size = config.get('batch_size', 10) device = accelerator.device epochs = config.get('epochs_codec', 100) save_freq = config.get('save_freq', 2) log_interval = config.get('log_interval', 10) # n_mels - default to 80 n_mels = config.get('n_mels', 80) data_params = config.get('data_params', None) sr = config['preprocess_params'].get('sr', 44100) train_path = data_params['train_data'] val_path = data_params['val_data'] root_path = data_params['root_path'] min_length = data_params['min_length'] max_len = config.get('max_len', 7500) train_list, val_list = get_data_path_list(train_path, val_path) # Determine dump directory based on n_mels dump_dir = config.get('dump_dir', f"./dump_{n_mels}") # Dataloaders train_dataloader = build_dataloader( train_list, root_path, alignment_path="../data_preparation_v2/alignments_train.safetensors", feature_dir=f"{dump_dir}/train", stats_path=f"{dump_dir}/speaker_stats.json", min_length=min_length, batch_size=batch_size, num_workers=16, dataset_config={}, device=device ) val_dataloader = build_dataloader( val_list, root_path, alignment_path="../data_preparation_v2/alignments_val.safetensors", feature_dir=f"{dump_dir}/val", stats_path=f"{dump_dir}/speaker_stats.json", min_length=min_length, batch_size=20, validation=True, num_workers=0, device=device, dataset_config={} ) # Model params model_params = recursive_munch(config['model_params']) # Get hop_length from config hop_length = config['preprocess_params']['spect_params'].get('hop_length', 441) # Scheduler params scheduler_params = { "max_lr": float(config['optimizer_params'].get('lr', 1e-5)), "min_lr": float(config['optimizer_params'].get('min_lr', 1e-5)), "pct_start": float(config['optimizer_params'].get('pct_start', 0.33)), "epochs": epochs, "steps_per_epoch": len(train_dataloader), } decay_epochs = int(epochs * scheduler_params["pct_start"]) print(f"[LR Schedule] Pre-TMA (epochs 0-{decay_epochs}): {scheduler_params['max_lr']:.2e} -> {scheduler_params['min_lr']:.2e}") print(f"[LR Schedule] TMA (epochs {decay_epochs}-{epochs}): constant {scheduler_params['min_lr']:.2e}") # Build model with speaker embeddings model = build_model_mel_speaker(model_params) # Log speaker embedding info num_speakers = model_params.get('num_speakers', 11) speaker_dim = model_params.get('speaker_dim', 128) print(f"[Mel Speaker Codec] Number of speakers: {num_speakers} (IDs 0-{num_speakers-1})") print(f"[Mel Speaker Codec] Speaker embedding dim: {speaker_dim}") print(f"[Mel Speaker Codec] Speaker embedding shape: {model.codec.speaker_embedding.weight.shape}") print(f"[Mel Speaker Codec] Input n_mels: {n_mels}") scheduler_params_dict = {key: scheduler_params.copy() for key in model} # Prepare models for k in model: model[k] = accelerator.prepare(model[k]) model[k] = model[k].to(device) train_dataloader, val_dataloader = accelerator.prepare(train_dataloader, val_dataloader) optimizer = build_optimizer( {key: model[key].parameters() for key in model}, scheduler_params_dict={key: scheduler_params.copy() for key in model}, lr=float(config['optimizer_params'].get('lr', 1e-5)) ) for k, v in optimizer.optimizers.items(): optimizer.optimizers[k] = accelerator.prepare(optimizer.optimizers[k]) optimizer.schedulers[k] = accelerator.prepare(optimizer.schedulers[k]) # Load pretrained start_epoch = 0 iters = 0 with accelerator.main_process_first(): if config.get('pretrained_model', '') != '': model, optimizer, start_epoch, iters = load_checkpoint( model, optimizer, config['pretrained_model'], load_only_params=config.get('load_only_params', False) ) print(f"Loaded pretrained model from {config['pretrained_model']}") desired_lr = float(config['optimizer_params'].get('lr', 1e-5)) for key in optimizer.optimizers: for param_group in optimizer.optimizers[key].param_groups: param_group['lr'] = desired_lr print(f"Reset learning rate to {desired_lr:.2e}") # Losses loss_params = Munch(config.get('loss_params', {})) TMA_epoch = loss_params.get('TMA_epoch', 5) stft_loss = MultiResolutionSTFTLoss().to(device) mag_phase_loss = MagPhaseLoss( n_fft=model_params.decoder.gen_istft_n_fft, hop_length=model_params.decoder.gen_istft_hop_size ).to(device) gl = GeneratorLoss(model.mpd, model.msd).to(device) dl = DiscriminatorLoss(model.mpd, model.msd).to(device) wl = WavLMLoss(model_sr=sr, slm_sr=16000).to(device) # Loss weights lambda_mel = loss_params.get('lambda_mel', 45.0) lambda_gen = loss_params.get('lambda_gen', 1.0) lambda_slm = loss_params.get('lambda_slm', 1.0) lambda_mag = loss_params.get('lambda_mag', 0.1) lambda_entropy = loss_params.get('lambda_entropy', 10.0) lambda_f0 = loss_params.get('lambda_f0',1.0) lambda_phoneme = loss_params.get('lambda_phoneme', 0.5) # Phoneme prediction loss weight best_loss = float('inf') # Compression level codec_strides = model_params.get('codec_strides', [2]) upsample_rates = model_params.decoder.upsample_rates codec_compression = 1 for s in codec_strides: codec_compression *= s istft_hop = model_params.decoder.gen_istft_hop_size print(f"[Mel Speaker Codec] Compression: {codec_compression}x") print(f"[Mel Speaker Codec] Upsample Rates: {upsample_rates} (Total: {np.prod(upsample_rates)})") print(f"[Mel Speaker Codec] Token Rate: {44100 / (hop_length * codec_compression):.2f} Hz") print(f"[Mel Speaker Codec] max_len: {max_len}") # FSQ levels fsq_levels = model_params.get('fsq_levels', [4] * 6) codebook_size = np.prod(fsq_levels) print(f"[Mel Speaker Codec] FSQ Levels: {fsq_levels} ({codebook_size} codes)") for epoch in range(start_epoch, epochs): running_loss = 0 start_time = time.time() _ = [model[key].train() for key in model] for i, batch in enumerate(train_dataloader): try: # =============================================================== # START OF TRAINING STEP # =============================================================== waves = batch[1] tensors = [b.to(device) for b in batch[2:13] if isinstance(b, torch.Tensor)] ( speakers_ids, # [B] - speaker IDs languages_ids, mels, # [B, n_mels, T] - mel spectrogram mel_input_length, _target_ph, _target_ph_l, context_ph, context_lens, durations_tg, pitches, # [B, T] - pitch (log F0) energies, # [B, T] - energy (not used in this model) ) = tensors # Gather lengths mel_input_length_all = accelerator.gather(mel_input_length) mel_len = min([int(mel_input_length_all.min().item() / 2 - 1), max_len // 2]) mel_len_st = int(mel_input_length.min().item() / 2 - 1) # Prepare batches (random crop for training) mel_list, wav_list, f0_list, spk_list = [], [], [], [] phoneme_list, phoneme_len_list = [], [] lang_list = [] dur_list = [] cropped_phoneme_list = [] cropped_dur_list = [] for bib in range(len(mel_input_length)): mel_length = int(mel_input_length[bib].item() / 2) random_start = np.random.randint(0, mel_length - mel_len) # Crop mel spectrogram mel_start_frame = random_start * 2 mel_end_frame = (random_start + mel_len) * 2 mel_crop = mels[bib, :, mel_start_frame:mel_end_frame] mel_list.append(mel_crop) # Crop waveform y = waves[bib][mel_start_frame * hop_length:mel_end_frame * hop_length] wav_list.append(y.to(device)) # Crop pitch f0_list.append(pitches[bib, mel_start_frame:mel_end_frame]) # Speaker ID spk_list.append(speakers_ids[bib]) # Language ID lang_list.append(languages_ids[bib]) # Phonemes and durations for this sample (full sequence for reference) phoneme_list.append(_target_ph[bib]) dur_list.append(durations_tg[bib]) # Crop phonemes and durations to match the cropped mel region cropped_ph, cropped_dur = crop_phonemes_and_durations( _target_ph[bib], durations_tg[bib], mel_start_frame, mel_end_frame ) cropped_phoneme_list.append(cropped_ph) cropped_dur_list.append(cropped_dur) mel_batch = torch.stack(mel_list).detach() # [B, n_mels, T] F0 = torch.stack(f0_list).float().detach() # [B, T] wav = torch.stack(wav_list).float().detach() # [B, T_audio] speaker_ids_batch = torch.stack(spk_list) # [B] - speaker IDs language_ids_batch = torch.stack(lang_list) # [B] - language IDs if mel_batch.shape[-1] < 80: continue current_prob_pitch = min(0.8, 0.3 + (epoch / 20)) use_predicted_f0 = random.random() < current_prob_pitch # =============================================================== # CODEC FORWARD WITH SPEAKER IDS AND LANGUAGE IDS # =============================================================== # Convert to padded tensors (handling potential length mismatches if simply stacking) if len(cropped_phoneme_list) > 0: # Find max length in this batch max_len_ph = max(len(x) for x in cropped_phoneme_list) padded_phonemes = torch.zeros((len(cropped_phoneme_list), max_len_ph), dtype=torch.long, device=device) padded_durations = torch.zeros((len(cropped_dur_list), max_len_ph), dtype=torch.float, device=device) for b in range(len(cropped_phoneme_list)): L = len(cropped_phoneme_list[b]) if L > 0: padded_phonemes[b, :L] = torch.tensor(cropped_phoneme_list[b], device=device) padded_durations[b, :L] = torch.tensor(cropped_dur_list[b], dtype=torch.float, device=device) else: padded_phonemes = None padded_durations = None codec_out = model.codec( mel=mel_batch, pitch=F0, speaker_ids=speaker_ids_batch, language_ids=language_ids_batch, use_predicted_f0=use_predicted_f0, target_phonemes=padded_phonemes, target_durations=padded_durations, ) y_rec = codec_out['wav'] mag = codec_out['mag'] phase = codec_out['phase'] commitment_loss = codec_out['commitment_loss'] tokens = codec_out['tokens'] f0_pred = codec_out['f0_pred'] speaker_emb = codec_out['speaker_emb'] # [B, speaker_dim] learned speaker embedding # Get phoneme loss directly from model (calculated at token rate) loss_phoneme = codec_out.get('phoneme_loss') if loss_phoneme is None: loss_phoneme = torch.tensor(0.0, device=device) # Match lengths for reconstruction min_len = min(y_rec.shape[-1], wav.shape[-1]) y_rec = y_rec[..., :min_len] wav = wav[..., :min_len] # Match F0 lengths min_len_f0 = min(f0_pred.shape[-1], F0.shape[-1]) f0_pred = f0_pred[..., :min_len_f0] F0_target = F0[..., :min_len_f0] # F0 Loss loss_f0 = F.mse_loss(f0_pred, F0_target) # =============================================================== # DISCRIMINATOR STEP # =============================================================== if epoch >= TMA_epoch: optimizer.zero_grad() d_loss = dl(wav.detach().unsqueeze(1).float(), y_rec.detach()).mean() accelerator.backward(d_loss) optimizer.step('msd') optimizer.step('mpd') else: d_loss = torch.tensor(0.0, device=device) # === Generator Step === optimizer.zero_grad() # STFT loss loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) # Mag/Phase loss loss_mag_phase = mag_phase_loss(mag, phase, wav.detach()) # Codebook utilization with torch.no_grad(): B, n_q, T = tokens.shape flat_tokens = tokens.reshape(-1) unique_codes = len(torch.unique(flat_tokens)) codebook_size = np.prod(model.codec.fsq_levels) perplexity = unique_codes / codebook_size * 100 if epoch >= TMA_epoch: loss_gen = gl(wav.detach().unsqueeze(1).float(), y_rec).mean() loss_slm = wl(wav.detach(), y_rec).mean() g_loss = ( lambda_mel * loss_mel + lambda_gen * loss_gen + lambda_slm * loss_slm + lambda_mag * loss_mag_phase + lambda_entropy * commitment_loss + lambda_f0 * loss_f0 + lambda_phoneme * loss_phoneme ) else: loss_gen = torch.tensor(0.0, device=device) loss_slm = torch.tensor(0.0, device=device) g_loss = ( lambda_mel * loss_mel + lambda_mag * loss_mag_phase + lambda_entropy * commitment_loss + lambda_f0 * loss_f0 + lambda_phoneme * loss_phoneme ) running_loss += accelerator.gather(loss_mel).mean().item() accelerator.backward(g_loss) optimizer.step('codec') iters += 1 entropy_loss = commitment_loss.item() if torch.is_tensor(commitment_loss) else commitment_loss phoneme_loss_val = loss_phoneme.item() if torch.is_tensor(loss_phoneme) else loss_phoneme if (i + 1) % log_interval == 0 and accelerator.is_main_process: current_lr = optimizer.optimizers['codec'].param_groups[0]['lr'] # Log unique speakers in batch unique_speakers = torch.unique(speaker_ids_batch).tolist() log_print( f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_dataloader)}], ' f'LR: {current_lr:.2e}, ' f'Mel: {running_loss / log_interval:.5f}, ' f'Gen: {loss_gen.item():.5f}, ' f'Disc: {d_loss.item():.5f}, ' f'F0: {loss_f0.item():.5f}, ' f'Ph: {phoneme_loss_val:.5f}, ' f'SLM: {loss_slm.item():.5f}, ' f'Entropy: {entropy_loss:.5f}, ' f'FSQ Util: {perplexity:.1f}% ({unique_codes}/{codebook_size}), ' f'Speakers: {unique_speakers}', logger ) writer.add_scalar('train/learning_rate', current_lr, iters) writer.add_scalar('train/mel_loss', running_loss / log_interval, iters) writer.add_scalar('train/gen_loss', loss_gen.item(), iters) writer.add_scalar('train/d_loss', d_loss.item(), iters) writer.add_scalar('train/f0_loss', loss_f0.item(), iters) writer.add_scalar('train/phoneme_loss', phoneme_loss_val, iters) writer.add_scalar('train/slm_loss', loss_slm.item(), iters) writer.add_scalar('train/entropy_loss', entropy_loss, iters) writer.add_scalar('train/fsq_utilization', perplexity, iters) writer.add_scalar('train/unique_codes', unique_codes, iters) writer.add_scalar('train/mag_phase_loss', loss_mag_phase.item(), iters) running_loss = 0 except torch.cuda.OutOfMemoryError: if accelerator.is_main_process: print(f"[Warning] OOM at epoch {epoch}, step {i}. Flushing cache.") for key in model: model[key].zero_grad() torch.cuda.empty_cache() gc.collect() continue # === Validation === loss_test = 0 _ = [model[key].eval() for key in model] eval_samples = [] # Store (mel_len, mel, f0, wav, speaker_id) with torch.no_grad(): iters_test = 0 for batch_idx, batch in enumerate(val_dataloader): try: waves = batch[1] tensors = [b.to(device) for b in batch[2:13] if isinstance(b, torch.Tensor)] ( speakers_ids, languages_ids, mels, mel_input_length, _target_ph, _target_ph_l, context_ph, context_lens, durations_tg, pitches, energies, ) = tensors # Collect samples for evaluation if len(eval_samples) < 20: for bib in range(len(mel_input_length)): if len(eval_samples) >= 20: break eval_samples.append(( mel_input_length[bib].item(), mels[bib], pitches[bib], waves[bib], speakers_ids[bib], languages_ids[bib], )) mel_len = min([int(mel_input_length.min().item() / 2 - 1), max_len // 2]) mel_list, wav_list, f0_list, spk_list, lang_list = [], [], [], [], [] for bib in range(len(mel_input_length)): mel_length = int(mel_input_length[bib].item() / 2) random_start = np.random.randint(0, mel_length - mel_len) mel_list.append(mels[bib, :, (random_start * 2):((random_start + mel_len) * 2)]) y = waves[bib][(random_start * 2) * hop_length:((random_start + mel_len) * 2) * hop_length] wav_list.append(y.to(device)) f0_list.append(pitches[bib, (random_start * 2):((random_start + mel_len) * 2)]) spk_list.append(speakers_ids[bib]) lang_list.append(languages_ids[bib]) wav = torch.stack(wav_list).float().detach() F0_curve = torch.stack(f0_list).detach() mel_batch = torch.stack(mel_list).detach() speaker_ids_batch = torch.stack(spk_list) language_ids_batch = torch.stack(lang_list) # Forward with speaker IDs codec_out = model.codec( mel=mel_batch, pitch=F0_curve, speaker_ids=speaker_ids_batch, language_ids=language_ids_batch, ) y_rec = codec_out['wav'] min_len = min(y_rec.shape[-1], wav.shape[-1]) y_rec = y_rec[..., :min_len] wav = wav[..., :min_len] loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) loss_test += accelerator.gather(loss_mel).mean().item() iters_test += 1 except Exception as e: print(f"Error in validation loop: {e}") continue if accelerator.is_main_process: print(f'Epoch: {epoch + 1}') current_val_loss = loss_test / iters_test if iters_test > 0 else 0 log_print(f'Validation loss: {current_val_loss:.3f}\n', logger) writer.add_scalar('eval/mel_loss', current_val_loss, epoch + 1) # Generate audio samples with torch.no_grad(): for bib, sample in enumerate(eval_samples): # mel_len_val, gt_mel, f0_sample, gt_wav, speaker_id, language_id = sample mel_length = int(mel_len_val) gt_mel = gt_mel[:, :mel_length].unsqueeze(0).to(device) f0_sample = f0_sample[:mel_length].unsqueeze(0).float().to(device) speaker_id_sample = speaker_id.unsqueeze(0).to(device) language_id_sample = language_id.unsqueeze(0).to(device) # Forward with speaker ID (use GT f0) codec_out = model.codec( gt_mel, f0_sample, speaker_id_sample, language_ids=language_id_sample, use_predicted_f0=False ) y_rec = codec_out['wav'] speaker_emb = codec_out['speaker_emb'] f0_pred = codec_out['f0_pred'] writer.add_audio( f'eval/mel_codec_{bib}_spk{speaker_id.item()}', y_rec.cpu().numpy().squeeze(), epoch, sample_rate=sr ) # Test with predicted F0 codec_out_pred_f0 = model.codec( gt_mel, f0_sample, speaker_id_sample, language_ids=language_id_sample, use_predicted_f0=True ) y_rec_pred_f0 = codec_out_pred_f0['wav'] writer.add_audio( f'eval/mel_codec_pred_f0_{bib}_spk{speaker_id.item()}', y_rec_pred_f0.cpu().numpy().squeeze(), epoch, sample_rate=sr ) print(f" Sample {bib}: speaker_id={speaker_id.item()}, f0_pred shape={f0_pred.shape}") # Test tokenize + decode path tokens, spk_emb = model.codec.tokenize( gt_mel, f0_sample, speaker_id_sample, language_ids=language_id_sample ) y_from_tokens, _ = model.codec.decode_tokens( tokens, speaker_id_sample, language_id_sample, f0=f0_sample ) writer.add_audio( f'eval/from_tokens_{bib}_spk{speaker_id.item()}', y_from_tokens.cpu().numpy().squeeze(), epoch, sample_rate=sr ) # Save GT audio gt_wav_len = int(mel_length * hop_length) if len(gt_wav) > gt_wav_len: gt_wav = gt_wav[:gt_wav_len] writer.add_audio(f'eval/gt_{bib}', gt_wav.cpu().numpy(), epoch, sample_rate=sr) # Save checkpoint if epoch % save_freq == 0: if accelerator.is_main_process: print('Saving checkpoint...') state = { 'net': {key: model[key].state_dict() for key in model}, 'optimizer': optimizer.state_dict(), 'iters': iters, 'val_loss': loss_test / iters_test if iters_test > 0 else 0, 'epoch': epoch, } save_path = osp.join(log_dir, f'epoch_mel_speaker_codec_{epoch:05d}.pth') torch.save(state, save_path) gc.collect() torch.cuda.empty_cache() # Final save if accelerator.is_main_process: print('Saving final checkpoint...') net_state = { key: accelerator.unwrap_model(model[key]).state_dict() for key in model } state = { 'net': net_state, 'optimizer': optimizer.state_dict(), 'iters': iters, 'val_loss': loss_test / iters_test if iters_test > 0 else 0, 'epoch': epoch, } torch.save(state, osp.join(log_dir, 'mel_speaker_codec_final.pth')) if __name__ == "__main__": main()