# train_codec_hybrid_temporal.py # ============================================================================== # Training script for Hybrid TTS Codec with WINDOWED Temporal Style Encoder # # KEY DIFFERENCE: # - Style encoder uses large overlapping windows (~3 sec) for timbre # - Output: [B, num_windows, 32] - coarse temporal, captures speaker/timbre # - Prevents overfitting to phoneme-level content # # Architecture: # WindowedTemporalStyleEncoder: mel [B, 40, T] -> style [B, num_windows, 32] # ProsodyEncoder: (pitch, energy, windowed_style) -> latent # FSQ: latent -> tokens (prosody ONLY quantized) # TextDownsampler: text -> text_down (continuous) # Decoder: (quantized_prosody, text_down, windowed_style) -> waveform # # Usage: # python train_codec_hybrid_temporal.py -p Configs/config_codec_hybrid_temporal.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 import build_model, load_checkpoint, gaussian_upsample 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 # Hybrid Codec with Temporal Style from Modules.codec_decoder_hybrid_temporal import HybridTTSCodecVocoderTemporal 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 def spec_augment(mel, T_mask_param=20, F_mask_param=10, num_T_masks=2, num_F_masks=1): """ Apply SpecAugment (Time and Frequency Masking) to mel spectrograms. Used for style encoder regularization to prevent content leakage. Args: mel: [B, n_mels, T] """ B, n_mels, T = mel.shape cloned_mel = mel.clone() # Frequency masking for i in range(B): for _ in range(num_F_masks): f_width = random.randint(0, F_mask_param) f_start = random.randint(0, max(1, n_mels - f_width)) cloned_mel[i, f_start:f_start+f_width, :] = 0 # Time masking for i in range(B): for _ in range(num_T_masks): t_width = random.randint(0, T_mask_param) t_start = random.randint(0, max(1, T - t_width)) cloned_mel[i, :, t_start:t_start+t_width] = 0 return cloned_mel @click.command() @click.option('-p', '--config_path', default='Configs/config_codec_hybrid_temporal.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_hybrid_temporal_codec.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 40 n_mels = config.get('n_mels', 40) 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 = 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, # persistent_workers=True in build_dataloader prevents FD leak 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 - LR decay during pre-TMA, constant during TMA 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)), # Decay fraction "epochs": epochs, "steps_per_epoch": len(train_dataloader), } # Log LR schedule 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} (linear decay)") print(f"[LR Schedule] TMA (epochs {decay_epochs}-{epochs}): constant {scheduler_params['min_lr']:.2e}") # Build base models (text encoder, style encoder for comparison, discriminators) model = build_model(model_params) 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']}") # Force LR update after loading checkpoint (checkpoint LR overrides config) 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} (overriding checkpoint LR)") # 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_style = loss_params.get('lambda_style', 1.0) # Style prediction loss # Noise injection for style encoder regularization style_noise_std = config.get('style_noise_std', 0.2) # Gaussian noise std for mel input to style encoder style_dropout_prob = config.get('style_dropout_prob', 0.5) # Probability of using predicted style during training print(f"[Regularization] Style encoder mel noise std: {style_noise_std}") print(f"[Regularization] Style dropout (scheduled sampling) prob: {style_dropout_prob}") best_loss = float('inf') # Compression level codec_strides = model_params.codec_strides 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 required_upsample = hop_length // istft_hop current_upsample = int(np.prod(upsample_rates)) print(f"[Hybrid Temporal Codec] changed ptich13") print(f"[Hybrid Temporal Codec] Compression: {codec_compression}x") print(f"[Hybrid Temporal Codec] Upsample Rates: {upsample_rates} (Total: {np.prod(upsample_rates)})") print(f"[Hybrid Temporal Codec] Token Rate: {44100 / (hop_length * codec_compression):.2f} Hz") print(f"[Hybrid Temporal Codec] Style: Windowed [B, num_windows, 32] (~3sec windows, ~1sec stride)") print(f"[Hybrid Temporal Codec] n_mels: {model_params.n_mels}") print(f"[Hybrid Temporal Codec] max_len: {max_len}") # FSQ levels for prosody quantization fsq_levels = model_params.get('fsq_levels', [4] * 6) codebook_size = np.prod(fsq_levels) print(f"[Hybrid Temporal Codec] FSQ Levels: {fsq_levels} ({codebook_size} codes)") print(f"[Hybrid Temporal Codec] Text stays continuous (not quantized)!") 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, languages_ids, mels, mel_input_length, _target_ph, _target_ph_l, context_ph, context_lens, durations_tg, pitches, energies, ) = tensors with torch.no_grad(): text_pad_mask = length_to_mask(_target_ph_l) # Extract language embedding safely (handle DDP wrapper) lang_emb = accelerator.unwrap_model(model.text_encoder).language_emb(languages_ids).detach() # Text encoding - stays continuous (NOT quantized) t_en = model.text_encoder( _target_ph, _target_ph_l, # language_id=languages_ids, language_emb=lang_emb ) asr_features, _ = gaussian_upsample( t_en, durations_tg, mel_input_length, sigma_sq=10, token_mask=~text_pad_mask ) # Gather lengths for distributed training 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 en, gt, wav, f0_list, n0_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) en.append(asr_features[bib, :, (random_start * 2):((random_start+mel_len) * 2)]) gt.append(mels[bib, :, (random_start * 2):((random_start+mel_len) * 2)]) y = waves[bib][(random_start * 2) * 441:((random_start+mel_len) * 2) * 441] wav.append(y.to(device)) f0_list.append(pitches[bib, (random_start * 2):((random_start+mel_len) * 2)]) n0_list.append(energies[bib, (random_start * 2):((random_start+mel_len) * 2)]) en = torch.stack(en) # [B, 512, T] - text embeddings (continuous) gt = torch.stack(gt).detach() # [B, n_mels, T] F0 = torch.stack(f0_list).float().detach() # [B, T] - pitch N0 = torch.stack(n0_list).float().detach() # [B, T] - energy wav = torch.stack(wav).float().detach() # [B, T_audio] if gt.shape[-1] < 80: continue # Add Gaussian noise to mel for style encoder regularization (prevent overfitting) # Apply SpecAugment to prevent style encoder from learning text/content # REDUCED from T=30,F=8 to T=15,F=4 to preserve prosody depth if model.codec.training: # 1. Gaussian Noise if style_noise_std > 0: mel_input = gt + torch.randn_like(gt) * style_noise_std else: mel_input = gt.clone() # Milder SpecAugment to preserve prosody/emotion information mel_input = spec_augment(mel_input, T_mask_param=20, F_mask_param=4) else: mel_input = gt current_prob = min(0.8, 0.5+(epoch / 20)) current_prob_pitch = min(0.8, 0.3+(epoch / 40)) if random.random() < current_prob: use_predicted_style = True else: use_predicted_style = False if random.random() < current_prob_pitch: use_predicted_f0 = True else: use_predicted_f0 = False codec_out = model.codec( pitch=F0, energy=N0, text_emb=en, mel=mel_input, # Pass augmented/noisy mel for regularization use_predicted_style=use_predicted_style, use_predicted_f0=use_predicted_f0, language_emb=lang_emb ) 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'] style_temporal = codec_out['style_temporal'] # [B, num_windows, 32] - GT from windowed style encoder style_pred = codec_out['style_pred'] # [B, T_pred, 32] - predicted from quantized prosody style_up = codec_out['style_up'] # [B, T, 32] - upsampled GT style # Match lengths 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) # Style Prediction Loss - compare predicted style to GT style from encoder # style_temporal: [B, num_windows, 32] (windowed, coarse) # style_pred: [B, T_pred, 32] (from prosody latent, different temporal resolution) # Interpolate both to common temporal resolution before comparing style_pred_t = style_pred.transpose(1, 2) # [B, 32, T_pred] style_temporal_t = style_temporal.transpose(1, 2) # [B, 32, num_windows] common_len = max(style_pred_t.shape[-1], style_temporal_t.shape[-1]) style_pred_interp = F.interpolate(style_pred_t, size=common_len, mode='linear', align_corners=False) style_temporal_interp = F.interpolate(style_temporal_t, size=common_len, mode='linear', align_corners=False) loss_style = F.mse_loss(style_pred_interp, style_temporal_interp.detach()) # === 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 (multi-resolution) loss_mel = stft_loss(y_rec.squeeze(), wav.detach()) # Mag/Phase loss for iSTFT loss_mag_phase = mag_phase_loss(mag, phase, wav.detach()) # Codebook utilization (for logging) 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_style * loss_style ) 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_style * loss_style ) running_loss += accelerator.gather(loss_mel).mean().item() accelerator.backward(g_loss) optimizer.step('text_encoder') optimizer.step('codec') iters += 1 entropy_loss = commitment_loss.item() if torch.is_tensor(commitment_loss) else commitment_loss if (i + 1) % log_interval == 0 and accelerator.is_main_process: # Get current learning rate current_lr = optimizer.optimizers['codec'].param_groups[0]['lr'] # Log temporal style shape style_shape = style_temporal.shape 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'SLM: {loss_slm.item():.5f}, ' f'Style: {loss_style.item():.5f}, ' f'Entropy: {entropy_loss:.5f}, ' f'FSQ Util: {perplexity:.1f}% ({unique_codes}/{codebook_size}), ' f'StyleShape: {list(style_shape)}', 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/slm_loss', loss_slm.item(), iters) writer.add_scalar('train/style_loss', loss_style.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 # =============================================================== # CATCH OOM ERROR # =============================================================== except torch.cuda.OutOfMemoryError: if accelerator.is_main_process: print(f"[Warning] OOM at epoch {epoch}, step {i}. Flushing cache and skipping batch.") # Clear gradients to prevent accumulation from failed step for key in model: model[key].zero_grad() # Force garbage collection and empty cache torch.cuda.empty_cache() import gc gc.collect() continue # === Validation === loss_test = 0 _ = [model[key].eval() for key in model] eval_samples = [] # Store tuple: (mel_len, mel, asr, f0, n, lang, wav) 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 text_pad_mask = length_to_mask(_target_ph_l) lang_emb = model.text_encoder.language_emb(languages_ids) # Text encoding (continuous) t_en = model.text_encoder( _target_ph, _target_ph_l, language_emb=lang_emb ) asr_features, _ = gaussian_upsample( t_en, durations_tg, mel_input_length, sigma_sq=10, token_mask=~text_pad_mask ) # 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], asr_features[bib], pitches[bib], energies[bib], lang_emb[bib] if lang_emb is not None else None, waves[bib] )) mel_len = min([int(mel_input_length.min().item() / 2 - 1), max_len // 2]) en, gt, wav_list, f0_list, n_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) en.append(asr_features[bib, :, (random_start * 2):((random_start + mel_len) * 2)]) gt.append(mels[bib, :, (random_start * 2):((random_start + mel_len) * 2)]) y = waves[bib][(random_start * 2) * 441:((random_start + mel_len) * 2) * 441] wav_list.append(y.to(device)) f0_list.append(pitches[bib, (random_start * 2):((random_start + mel_len) * 2)]) n_list.append(energies[bib, (random_start * 2):((random_start + mel_len) * 2)]) wav = torch.stack(wav_list).float().detach() F0_curve = torch.stack(f0_list).detach() N_curve = torch.stack(n_list).detach() en = torch.stack(en) gt = torch.stack(gt).detach() # Hybrid temporal codec forward codec_out = model.codec( pitch=F0_curve, energy=N_curve, text_emb=en, mel=gt, language_emb=lang_emb ) y_rec = codec_out['wav'] # Match lengths 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, en_sample, f0_sample, n_sample, lang_emb_sample, gt_wav = sample mel_length = int(mel_len_val) gt_mel = gt_mel[:, :mel_length].unsqueeze(0).to(device) en_sample = en_sample[:, :mel_length].unsqueeze(0).to(device) f0_sample = f0_sample[:mel_length].unsqueeze(0).float().to(device) n_sample = n_sample[:mel_length].unsqueeze(0).float().to(device) lang_emb_sample = lang_emb_sample.unsqueeze(0).to(device) if lang_emb_sample is not None else None # Full forward (with GT style from mel) codec_out = model.codec(f0_sample, n_sample, en_sample, gt_mel, language_emb=lang_emb_sample) y_rec = codec_out['wav'] style_temporal = codec_out['style_temporal'] style_pred = codec_out['style_pred'] writer.add_audio(f'eval/hybrid_temporal_codec_{bib}', y_rec.cpu().numpy().squeeze(), epoch, sample_rate=sr) # Log style shapes print(f" Sample {bib}: style_temporal shape = {style_temporal.shape}, style_pred shape = {style_pred.shape}") # Test tokenize + decode path WITH GT style tokens, text_down, style_temp = model.codec.tokenize(f0_sample, n_sample, en_sample, gt_mel) y_from_tokens_pred_style = model.codec.decode_tokens(tokens, text_down, style_temporal=None, language_emb=lang_emb_sample) writer.add_audio(f'eval/from_tokens_pred_style_{bib}', y_from_tokens_pred_style.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, 'epoch': epoch, } save_path = osp.join(log_dir, f'epoch_hybrid_temporal_codec_{epoch:05d}.pth') torch.save(state, save_path) # Cleanup to prevent file descriptor leaks 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, 'epoch': epoch, } torch.save(state, osp.join(log_dir, 'hybrid_temporal_codec_final.pth')) if __name__ == "__main__": main()