import os import sys os.environ['HF_HUB_CACHE'] = './checkpoints/hf_cache' import torch import torch.multiprocessing as mp import random import librosa import yaml import argparse import torchaudio import torchaudio.compliance.kaldi as kaldi import glob import time from tqdm import tqdm import shutil import accelerate from .optimizers import build_optimizer from .data.ft_dataset import build_ft_dataloader import hydra from omegaconf import DictConfig from accelerate import Accelerator from accelerate import DistributedDataParallelKwargs from accelerate.logging import get_logger EXPECTED_SEEDVC_SKIPPED_KEYS = {"estimator.f0_embedder.weight", "estimator.input_pos"} class Trainer: def __init__( self, config_path, pretrained_cfm_ckpt_path, pretrained_ar_ckpt_path, data_dir, run_name, batch_size=0, num_workers=0, steps=1000, save_interval=500, max_epochs=1000, train_cfm=True, train_ar=False, mixed_precision=None, ): self.config_path = config_path self.mixed_precision = mixed_precision # Load configuration self.config = yaml.safe_load(open(config_path)) # Setup logging directory self.log_dir = os.path.join("runs", run_name) if not os.path.exists(self.log_dir): os.makedirs(self.log_dir, exist_ok=True) shutil.copy(config_path, os.path.join(self.log_dir, os.path.basename(config_path))) # Setup accelerator ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True, broadcast_buffers=False) self.accelerator = Accelerator( project_dir=self.log_dir, split_batches=True, kwargs_handlers=[ddp_kwargs], mixed_precision=mixed_precision ) self.device = self.accelerator.device # Initialize training parameters self._init_dataloader( data_dir=data_dir, batch_size=batch_size, num_workers=num_workers, spect_params=self.config['mel_fn'], sr=self.config['sr'], ) # Initialize models and optimizers self._init_models(train_cfm=train_cfm, train_ar=train_ar) # Load checkpoint if available self._load_checkpoint(pretrained_cfm_ckpt_path, pretrained_ar_ckpt_path) # Initialize training parameters self.iters = 0 self.start_epoch = 0 self.log_interval = 10 self.max_steps = steps self.save_interval = save_interval self.max_epochs = max_epochs def _init_dataloader(self, data_dir, batch_size, num_workers, spect_params, sr): self.spect_params = spect_params self.sr = sr # Initialize dataloader self.train_dataloader = build_ft_dataloader( data_dir, spect_params, self.sr, batch_size=batch_size, num_workers=num_workers, ) def _init_models(self, train_cfm=True, train_ar=False): """Initialize models and optimizers""" assert train_cfm or train_ar, "At least one model should be trained" self.train_cfm = train_cfm self.train_ar = train_ar # Initialize main model self._init_main_model(train_cfm=train_cfm, train_ar=train_ar) # Initialize optimizers self._init_optimizers() def _init_main_model(self, train_cfm=True, train_ar=False): """Initialize the main model""" with self.accelerator.main_process_first(): cfg = DictConfig(self.config) self.model = hydra.utils.instantiate(cfg).to(self.device) for p in self.model.parameters(): p.requires_grad = False if train_cfm: for p in self.model.cfm.parameters(): p.requires_grad = True for p in self.model.cfm_length_regulator.parameters(): p.requires_grad = True if train_ar: for p in self.model.ar.parameters(): p.requires_grad = True for p in self.model.ar_length_regulator.parameters(): p.requires_grad = True def _init_optimizers(self): """Initialize optimizers and schedulers""" from .optimizers import build_single_optimizer self.optimizer, self.scheduler = build_single_optimizer( self.model, lr=2e-5, ) self.optimizer = self.accelerator.prepare(self.optimizer) self.scheduler = self.accelerator.prepare(self.scheduler) def _find_checkpoint(self, name_pattern, max_keep=1): """Find checkpoint files in the specified directory""" available_checkpoints = glob.glob(os.path.join(self.log_dir, name_pattern)) if len(available_checkpoints) > max_keep - 1: # find the checkpoint that has the highest step number latest_checkpoint = max( available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0]) ) earliest_checkpoint = min( available_checkpoints, key=lambda x: int(x.split("_")[-1].split(".")[0]) ) # delete the earliest checkpoint if ( earliest_checkpoint != latest_checkpoint and self.accelerator.is_main_process and len(available_checkpoints) > max_keep ): os.remove(earliest_checkpoint) print(f"Removed {earliest_checkpoint}") return latest_checkpoint else: return None def _load_checkpoint(self, pretrained_cfm_ckpt_path, pretrained_ar_ckpt_path): """Load checkpoint if available""" cfm_checkpoint_path = pretrained_cfm_ckpt_path or self._find_checkpoint("CFM_epoch_*_step_*.pth", max_keep=1) ar_checkpoint_path = pretrained_ar_ckpt_path or self._find_checkpoint("AR_epoch_*_step_*.pth", max_keep=1) with self.accelerator.main_process_first(): if cfm_checkpoint_path: print(f"Loading CFM checkpoint from {cfm_checkpoint_path}") if ar_checkpoint_path: print(f"Loading AR checkpoint from {ar_checkpoint_path}") self.model.load_checkpoints(cfm_checkpoint_path=cfm_checkpoint_path, ar_checkpoint_path=ar_checkpoint_path) self.model = self.accelerator.prepare(self.model) def filter_state_dict_shapes(self, params, model): model_state_dict = model.state_dict() filtered_state_dict = { k: v for k, v in params.items() if k in model_state_dict and v.shape == model_state_dict[k].shape } skipped_keys = set(params.keys()) - set(filtered_state_dict.keys()) unexpected_skipped_keys = skipped_keys - EXPECTED_SEEDVC_SKIPPED_KEYS if unexpected_skipped_keys: print( f"Warning: Skipped loading some keys due to shape mismatch: {unexpected_skipped_keys}" ) return filtered_state_dict, skipped_keys def train(self): """Main training loop""" for epoch in range(self.start_epoch, self.start_epoch + 1000): epoch_start_time = time.time() try: self.train_dataloader.sampler.set_epoch(epoch) except AttributeError: pass self.model.train() for i, batch in enumerate(tqdm(self.train_dataloader)): # Process batch self._process_batch(epoch, i, batch) if self.iters >= self.max_steps and self.accelerator.is_main_process: print("Reached max steps, stopping training") self._save_checkpoint(epoch) exit() # Log epoch completion if self.accelerator.is_main_process: print(f"Epoch {epoch} completed in {time.time() - epoch_start_time:.2f} seconds") if epoch + 1 >= self.max_epochs and self.accelerator.is_main_process: print("Reached max epochs, stopping training") self._save_checkpoint(epoch) exit() def _process_batch(self, epoch, i, batch): """Process a single batch""" # Move batch to device waves, mels, wave_lens, mel_lens = batch # Resample to 16kHz for ASR models waves_16k = torchaudio.functional.resample(waves, self.sr, 16000) wave_lengths_16k = (wave_lens.float() * 16000 / self.sr).long() # Forward pass and loss calculation with self.accelerator.autocast(): loss_ar, loss_cfm = self.model( waves_16k.to(self.device), mels.to(self.device), wave_lengths_16k.to(self.device), mel_lens.to(self.device), forward_ar=self.train_ar, forward_cfm=self.train_cfm, ) loss = loss_ar + loss_cfm self.accelerator.backward(loss) grad_norm_g = torch.nn.utils.clip_grad_norm_( self.model.parameters(), 1000.0 ) self.optimizer.step() self.scheduler.step(self.iters) self.optimizer.zero_grad() # Log training progress self._log_training_progress(epoch, i, loss, loss_ar, loss_cfm, grad_norm_g) # Save checkpoint if self.iters != 0 and self.iters % self.save_interval == 0 and self.accelerator.is_main_process: self._save_checkpoint(epoch) # Increment iteration counter self.iters += 1 def _log_training_progress(self, epoch, i, loss, loss_ar, loss_cfm, grad_norm_g): """Log training progress to tensorboard and wandb""" if self.iters % self.log_interval == 0 and self.accelerator.is_main_process: with torch.no_grad(): cur_lr = self.scheduler.get_last_lr()[0] if i != 0 else 0 # Log to console print("Epoch %d, Iteration %d, Loss: %.4f, Loss AR: %.4f, Loss CFM: %.4f, Grad Norm: %.4f, LR: %.6f" % (epoch, i, loss.item(), loss_ar.item(), loss_cfm.item(), grad_norm_g, cur_lr)) def _save_checkpoint(self, epoch): """Save model checkpoint""" print('Saving checkpoint...') if self.train_ar: state = { 'net': { 'ar': self.accelerator.unwrap_model(self.model).ar.state_dict(), 'length_regulator': self.accelerator.unwrap_model(self.model).ar_length_regulator.state_dict(), }, 'iters': self.iters, 'epoch': epoch, } save_path = os.path.join(self.log_dir, 'AR_epoch_%05d_step_%05d.pth' % (epoch, self.iters)) torch.save(state, save_path) print(f"Saved AR checkpoint to {save_path}") # Find all checkpoints and remove old ones self._remove_old_checkpoints("AR_epoch_*_step_*.pth", max_keep=1) if self.train_cfm: state = { 'net': { 'cfm': self.accelerator.unwrap_model(self.model).cfm.state_dict(), 'length_regulator': self.accelerator.unwrap_model(self.model).cfm_length_regulator.state_dict(), }, 'iters': self.iters, 'epoch': epoch, } save_path = os.path.join(self.log_dir, 'CFM_epoch_%05d_step_%05d.pth' % (epoch, self.iters)) torch.save(state, save_path) print(f"Saved CFM checkpoint to {save_path}") # Find all checkpoints and remove old ones self._remove_old_checkpoints("CFM_epoch_*_step_*.pth", max_keep=1) def _remove_old_checkpoints(self, name_pattern, max_keep=1): """Remove old checkpoints""" checkpoints = glob.glob(os.path.join(self.log_dir, name_pattern)) if len(checkpoints) > max_keep: # Sort by step checkpoints.sort(key=lambda x: int(x.split('_')[-1].split('.')[0])) # Remove all except last 1 for cp in checkpoints[:-max_keep]: os.remove(cp) def main(args): trainer = Trainer( config_path=args.config, pretrained_cfm_ckpt_path=args.pretrained_cfm_ckpt, pretrained_ar_ckpt_path=args.pretrained_ar_ckpt, data_dir=args.dataset_dir, run_name=args.run_name, batch_size=args.batch_size, steps=args.max_steps, max_epochs=args.max_epochs, save_interval=args.save_every, num_workers=args.num_workers, train_cfm=args.train_cfm, train_ar=args.train_ar, ) trainer.train() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--config', type=str, default='configs/v2/vc_wrapper.yaml') parser.add_argument('--pretrained-cfm-ckpt', type=str, default=None) parser.add_argument('--pretrained-ar-ckpt', type=str, default=None) parser.add_argument('--dataset-dir', type=str, default='/path/to/dataset') parser.add_argument('--run-name', type=str, default='my_run') parser.add_argument('--batch-size', type=int, default=2) parser.add_argument('--max-steps', type=int, default=1000) parser.add_argument('--max-epochs', type=int, default=1000) parser.add_argument('--save-every', type=int, default=500) parser.add_argument('--num-workers', type=int, default=0) parser.add_argument('--train-cfm', action='store_true', help='Train CFM model') parser.add_argument('--train-ar', action='store_true', help='Train AR model') args = parser.parse_args() main(args)