def get_mel(config, model): mel_fn = MelSpectrogramFixed(sample_rate=config.data_config.sample_rate, n_fft=config.data_config.n_fft, win_length=config.data_config.win_length, hop_length=config.data_config.hop_length, f_min=config.data_config.f_min, f_max=config.data_config.f_max, n_mels=config.data_config.n_mels, window_fn=torch.hann_window).cuda() dataset = AudioDataset(config, training=False) test_batch = dataset.sample_test_batch(1) # n_iter = 25 # path_to_store_schedule = f'schedules/default/{n_iter}iters.pt' # iters_best_schedule, stats = iters_schedule_grid_search( # model, config, # n_iter=n_iter, # betas_range=(1e-06, 0.01), # test_batch_size=1, step=1, # path_to_store_schedule=path_to_store_schedule, # save_stats_for_grid=True, # verbose=True, n_jobs=4 # ) i = 0 for test in tqdm(test_batch): mel = mel_fn(test[None].cuda()) start = datetime.now() t = time() outputs = model.forward(mel, store_intermediate_states=False) end = datetime.now() print("Time infer: ", str(time() - t)) outputs = outputs.cpu().squeeze() save_path = str(i) + '.wav' i += 1 torchaudio.save(save_path, outputs, sample_rate=config.data_config.sample_rate) inference_time = (end - start).total_seconds() rtf = compute_rtf(outputs, inference_time, sample_rate=config.data_config.sample_rate) show_message(f'Done. RTF estimate:{np.std(rtf)}')
def run(config, args): show_message('Initializing logger...', verbose=args.verbose) logger = Logger(config) show_message('Initializing model...', verbose=args.verbose) model = WaveGrad(config).cuda() show_message(f'Number of parameters: {model.nparams}', verbose=args.verbose) mel_fn = MelSpectrogramFixed(sample_rate=config.data_config.sample_rate, n_fft=config.data_config.n_fft, win_length=config.data_config.win_length, hop_length=config.data_config.hop_length, f_min=config.data_config.f_min, f_max=config.data_config.f_max, n_mels=config.data_config.n_mels, window_fn=torch.hann_window).cuda() show_message('Initializing optimizer, scheduler and losses...', verbose=args.verbose) optimizer = torch.optim.Adam(params=model.parameters(), lr=config.training_config.lr) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=config.training_config.scheduler_step_size, gamma=config.training_config.scheduler_gamma) show_message('Initializing data loaders...', verbose=args.verbose) train_dataset = AudioDataset(config, training=True) train_dataloader = DataLoader(train_dataset, batch_size=config.training_config.batch_size, drop_last=True) test_dataset = AudioDataset(config, training=False) test_dataloader = DataLoader(test_dataset, batch_size=1) test_batch = test_dataset.sample_test_batch( config.training_config.n_samples_to_test) if config.training_config.continue_training: show_message('Loading latest checkpoint to continue training...', verbose=args.verbose) model, optimizer, iteration = logger.load_latest_checkpoint( model, optimizer) epoch_size = len(train_dataset) // config.training_config.batch_size epoch_start = iteration // epoch_size else: iteration = 0 epoch_start = 0 # Log ground truth test batch audios = { f'audio_{index}/gt': audio for index, audio in enumerate(test_batch) } logger.log_audios(0, audios) specs = { f'mel_{index}/gt': mel_fn(audio.cuda()).cpu().squeeze() for index, audio in enumerate(test_batch) } logger.log_specs(0, specs) show_message('Start training...', verbose=args.verbose) try: for epoch in range(epoch_start, config.training_config.n_epoch): # Training step model.set_new_noise_schedule( init=torch.linspace, init_kwargs={ 'steps': config.training_config.training_noise_schedule.n_iter, 'start': config.training_config.training_noise_schedule. betas_range[0], 'end': config.training_config.training_noise_schedule. betas_range[1] }) for i, batch in enumerate(train_dataloader): tic_iter = time.time() batch = batch.cuda() mels = mel_fn(batch) # Training step model.zero_grad() loss = model.compute_loss(mels, batch) loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_( parameters=model.parameters(), max_norm=config.training_config.grad_clip_threshold) optimizer.step() toc_iter = time.time() dur_iter = toc_iter - tic_iter dur_iter = np.round(dur_iter, 4) loss_stats = { 'total_loss': loss.item(), 'grad_norm': grad_norm.item() } iter_size = len(train_dataloader) logger.log_training(epoch, iteration, i, iter_size, loss_stats, dur_iter, int(dur_iter * iter_size), verbose=args.verbose) iteration += 1 # Test step if epoch % config.training_config.test_interval == 0: model.set_new_noise_schedule( init=torch.linspace, init_kwargs={ 'steps': config.training_config.test_noise_schedule.n_iter, 'start': config.training_config.test_noise_schedule. betas_range[0], 'end': config.training_config.test_noise_schedule. betas_range[1] }) with torch.no_grad(): # Calculate test set loss test_loss = 0 for i, batch in enumerate(test_dataloader): batch = batch.cuda() mels = mel_fn(batch) test_loss_ = model.compute_loss(mels, batch) test_loss += test_loss_ test_loss /= (i + 1) loss_stats = {'total_loss': test_loss.item()} # Restore random batch from test dataset audios = {} specs = {} test_l1_loss = 0 test_l1_spec_loss = 0 average_rtf = 0 for index, test_sample in enumerate(test_batch): test_sample = test_sample[None].cuda() test_mel = mel_fn(test_sample.cuda()) start = datetime.now() y_0_hat = model.forward( test_mel, store_intermediate_states=False) y_0_hat_mel = mel_fn(y_0_hat) end = datetime.now() generation_time = (end - start).total_seconds() average_rtf += compute_rtf( y_0_hat, generation_time, config.data_config.sample_rate) test_l1_loss += torch.nn.L1Loss()(y_0_hat, test_sample).item() test_l1_spec_loss += torch.nn.L1Loss()( y_0_hat_mel, test_mel).item() audios[f'audio_{index}/predicted'] = y_0_hat.cpu( ).squeeze() specs[f'mel_{index}/predicted'] = y_0_hat_mel.cpu( ).squeeze() average_rtf /= len(test_batch) show_message(f'Device: GPU. average_rtf={average_rtf}', verbose=args.verbose) test_l1_loss /= len(test_batch) loss_stats['l1_test_batch_loss'] = test_l1_loss test_l1_spec_loss /= len(test_batch) loss_stats['l1_spec_test_batch_loss'] = test_l1_spec_loss logger.log_test(epoch, loss_stats, verbose=args.verbose) #logger.log_audios(epoch, audios) #logger.log_specs(epoch, specs) logger.save_checkpoint(iteration, model, optimizer) if epoch % (epoch // 10 + 1) == 0: scheduler.step() except KeyboardInterrupt: print('KeyboardInterrupt: training has been stopped.') return
def run(config, args): if 'LOCAL_RANK' in os.environ and 'WORLD_SIZE' in os.environ: local_rank = int(os.environ['LOCAL_RANK']) world_size = int(os.environ['WORLD_SIZE']) else: local_rank = args.rank world_size = args.world_size distributed_run = world_size > 1 torch.manual_seed(args.seed + local_rank) np.random.seed(args.seed + local_rank) # if local_rank == 0: # if not os.path.exists(args.output): # os.makedirs(args.output) torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = False if distributed_run: init_distributed(args, world_size, local_rank) device = torch.device('cuda' if args.cuda else 'cpu') if local_rank == 0: print("start training") print("args", args) print("config", config) ############################################# # model if local_rank == 0: print("load model") model = WaveGrad(config).cuda() # optimizer amp config if local_rank == 0: print("configure optimizer and amp") kw = dict(lr=args.learning_rate, betas=(0.9, 0.98), eps=1e-9, weight_decay=args.weight_decay) if args.optimizer == 'adam': optimizer = FusedAdam(model.parameters(), **kw) elif args.optimizer == 'lamb': optimizer = FusedLAMB(model.parameters(), **kw) elif args.optimizer == 'pytorch': optimizer = torch.optim.Adam(model.parameters(), **kw) else: raise ValueError if args.amp: model, optimizer = amp.initialize(model, optimizer, opt_level="O1") if distributed_run: model = DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True) start_epoch = [1] start_iter = [0] ################ #load checkpoint if args.checkpoint_path is not None: ch_fpath = args.checkpoint_path load_checkpoint(local_rank, model, optimizer, start_epoch, start_iter, config, args.amp, ch_fpath, world_size) start_epoch = start_epoch[0] total_iter = start_iter[0] # dataloader ########################################################## if local_rank == 0: print("load dataset") if local_rank == 0: print("prepare train dataset") train_dataset = AudioDataset(config, training=True) # distributed sampler if distributed_run: train_sampler, shuffle = DistributedSampler(train_dataset), False else: train_sampler, shuffle = None, True train_loader = DataLoader(train_dataset, num_workers=1, shuffle=shuffle, sampler=train_sampler, batch_size=args.batch_size, pin_memory=False, drop_last=True) # ground truth samples if local_rank == 0: print("prepare test_dataset") test_dataset = AudioDataset(config, training=False) test_loader = DataLoader(test_dataset, batch_size=1) test_batch = test_dataset.sample_test_batch( config.training_config.n_samples_to_test) # Log ground truth test batch if local_rank == 0: print("save truth wave and mel") mel_fn = MelSpectrogramFixed(sample_rate=config.data_config.sample_rate, n_fft=config.data_config.n_fft, win_length=config.data_config.win_length, hop_length=config.data_config.hop_length, f_min=config.data_config.f_min, f_max=config.data_config.f_max, n_mels=config.data_config.n_mels, window_fn=torch.hann_window).cuda() audios = { f'audio_{index}/gt': audio for index, audio in enumerate(test_batch) } specs = { f'mel_{index}/gt': mel_fn(audio.cuda()).cpu().squeeze() for index, audio in enumerate(test_batch) } ####### loop start #epoch iteration = 0 model.train() val_loss = 0.0 torch.cuda.synchronize() if local_rank == 0: print("epoch start") for epoch in range(start_epoch, args.epochs + 1): tic_epoch = time.time() epoch_loss = 0.0 if distributed_run: train_loader.sampler.set_epoch(epoch) accumulated_steps = 0 iter_loss = 0 epoch_iter = 0 #iteration = 0 num_iters = len(train_loader) // args.gradient_accumulation_steps model.module.set_new_noise_schedule( # 1000 default init=torch.linspace, init_kwargs={ 'steps': config.training_config.training_noise_schedule.n_iter, 'start': config.training_config.training_noise_schedule.betas_range[0], 'end': config.training_config.training_noise_schedule.betas_range[1] } ) for i, batch in enumerate(train_loader): tic_iter = time.time() old_lr = optimizer.param_groups[0]['lr'] adjust_learning_rate(iteration, optimizer, args.learning_rate, args.warmup_steps) new_lr = optimizer.param_groups[0]['lr'] model.zero_grad() batch = batch.cuda() mels = mel_fn(batch) # Training step model.zero_grad() loss = model.module.compute_loss(mels, batch) if args.amp: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() if distributed_run: reduced_loss = reduce_tensor(loss.data, world_size).item() else: reduced_loss = loss.item() # if np.isnan(reduced_loss): # raise Exception("loss is NaN") iter_loss += reduced_loss if args.amp: grad_norm = torch.nn.utils.clip_grad_norm_( amp.master_params(optimizer), args.grad_clip_thresh) else: grad_norm = torch.nn.utils.clip_grad_norm_( model.parameters(), args.grad_clip_thresh) optimizer.step() toc_iter = time.time() dur_iter = toc_iter - tic_iter epoch_loss += iter_loss iter_size = len(train_loader) dur_epoch_est = iter_size * dur_iter if local_rank == 0: print( "\nepoch {:4d} | iter {:>12d} {:>3d}/{:3d} | {:3.2f}s/iter est {:4.2f}s/epoch | losses {:>12.6f} {:>12.6f} LR {:e}--> {:e}" .format(epoch, iteration, i, iter_size, dur_iter, dur_epoch_est, iter_loss, grad_norm, old_lr, new_lr), end='') iter_loss = 0 iteration += 1 # Finished epoch toc_epoch = time.time() dur_epoch = toc_epoch - tic_epoch if local_rank == 0: print("for {}item, {:4.2f}s/epoch ".format( iter_size, dur_epoch)) # Test step if epoch % config.training_config.test_interval == 0: model.module.set_new_noise_schedule( # 50 for default init=torch.linspace, init_kwargs={ 'steps': config.training_config.test_noise_schedule.n_iter, 'start': config.training_config.test_noise_schedule.betas_range[0], 'end': config.training_config.test_noise_schedule.betas_range[1] } ) if (epoch % args.epochs_per_checkpoint == 0): ch_path = os.path.join(args.output, "WaveGrad_ch_{:d}.pt".format(epoch)) save_checkpoint(local_rank, model, optimizer, epoch, iteration, config, args.amp, ch_path)
def run_training(rank, config, args): if args.n_gpus > 1: init_distributed(rank, args.n_gpus, config.dist_config) torch.cuda.set_device(f'cuda:{rank}') show_message('Initializing logger...', verbose=args.verbose, rank=rank) logger = Logger(config, rank=rank) show_message('Initializing model...', verbose=args.verbose, rank=rank) model = WaveGrad(config).cuda() show_message(f'Number of WaveGrad parameters: {model.nparams}', verbose=args.verbose, rank=rank) mel_fn = MelSpectrogramFixed(**config.data_config).cuda() show_message('Initializing optimizer, scheduler and losses...', verbose=args.verbose, rank=rank) optimizer = torch.optim.Adam(params=model.parameters(), lr=config.training_config.lr) scheduler = torch.optim.lr_scheduler.StepLR( optimizer, step_size=config.training_config.scheduler_step_size, gamma=config.training_config.scheduler_gamma) if config.training_config.use_fp16: scaler = torch.cuda.amp.GradScaler() show_message('Initializing data loaders...', verbose=args.verbose, rank=rank) train_dataset = AudioDataset(config, training=True) train_sampler = DistributedSampler( train_dataset) if args.n_gpus > 1 else None train_dataloader = DataLoader(train_dataset, batch_size=config.training_config.batch_size, sampler=train_sampler, drop_last=True) if rank == 0: test_dataset = AudioDataset(config, training=False) test_dataloader = DataLoader(test_dataset, batch_size=1) test_batch = test_dataset.sample_test_batch( config.training_config.n_samples_to_test) if config.training_config.continue_training: show_message('Loading latest checkpoint to continue training...', verbose=args.verbose, rank=rank) model, optimizer, iteration = logger.load_latest_checkpoint( model, optimizer) epoch_size = len(train_dataset) // config.training_config.batch_size epoch_start = iteration // epoch_size else: iteration = 0 epoch_start = 0 # Log ground truth test batch if rank == 0: audios = { f'audio_{index}/gt': audio for index, audio in enumerate(test_batch) } logger.log_audios(0, audios) specs = { f'mel_{index}/gt': mel_fn(audio.cuda()).cpu().squeeze() for index, audio in enumerate(test_batch) } logger.log_specs(0, specs) if args.n_gpus > 1: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) show_message(f'INITIALIZATION IS DONE ON RANK {rank}.') show_message('Start training...', verbose=args.verbose, rank=rank) try: for epoch in range(epoch_start, config.training_config.n_epoch): # Training step model.train() (model if args.n_gpus == 1 else model.module).set_new_noise_schedule( init=torch.linspace, init_kwargs={ 'steps': config.training_config.training_noise_schedule.n_iter, 'start': config.training_config.training_noise_schedule. betas_range[0], 'end': config.training_config.training_noise_schedule. betas_range[1] }) for batch in ( tqdm(train_dataloader, leave=False) \ if args.verbose and rank == 0 else train_dataloader ): model.zero_grad() batch = batch.cuda() mels = mel_fn(batch) if config.training_config.use_fp16: with torch.cuda.amp.autocast(): loss = (model if args.n_gpus == 1 else model.module).compute_loss(mels, batch) scaler.scale(loss).backward() scaler.unscale_(optimizer) else: loss = (model if args.n_gpus == 1 else model.module).compute_loss(mels, batch) loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_( parameters=model.parameters(), max_norm=config.training_config.grad_clip_threshold) if config.training_config.use_fp16: scaler.step(optimizer) scaler.update() else: optimizer.step() loss_stats = { 'total_loss': loss.item(), 'grad_norm': grad_norm.item() } logger.log_training(iteration, loss_stats, verbose=False) iteration += 1 # Test step after epoch on rank==0 GPU if epoch % config.training_config.test_interval == 0 and rank == 0: model.eval() (model if args.n_gpus == 1 else model.module).set_new_noise_schedule( init=torch.linspace, init_kwargs={ 'steps': config.training_config.test_noise_schedule.n_iter, 'start': config.training_config.test_noise_schedule. betas_range[0], 'end': config.training_config.test_noise_schedule. betas_range[1] }) with torch.no_grad(): # Calculate test set loss test_loss = 0 for i, batch in enumerate( tqdm(test_dataloader) \ if args.verbose and rank == 0 else test_dataloader ): batch = batch.cuda() mels = mel_fn(batch) test_loss_ = (model if args.n_gpus == 1 else model.module).compute_loss(mels, batch) test_loss += test_loss_ test_loss /= (i + 1) loss_stats = {'total_loss': test_loss.item()} # Restore random batch from test dataset audios = {} specs = {} test_l1_loss = 0 test_l1_spec_loss = 0 average_rtf = 0 for index, test_sample in enumerate(test_batch): test_sample = test_sample[None].cuda() test_mel = mel_fn(test_sample.cuda()) start = datetime.now() y_0_hat = (model if args.n_gpus == 1 else model.module).forward( test_mel, store_intermediate_states=False) y_0_hat_mel = mel_fn(y_0_hat) end = datetime.now() generation_time = (end - start).total_seconds() average_rtf += compute_rtf( y_0_hat, generation_time, config.data_config.sample_rate) test_l1_loss += torch.nn.L1Loss()(y_0_hat, test_sample).item() test_l1_spec_loss += torch.nn.L1Loss()( y_0_hat_mel, test_mel).item() audios[f'audio_{index}/predicted'] = y_0_hat.cpu( ).squeeze() specs[f'mel_{index}/predicted'] = y_0_hat_mel.cpu( ).squeeze() average_rtf /= len(test_batch) show_message(f'Device: GPU. average_rtf={average_rtf}', verbose=args.verbose) test_l1_loss /= len(test_batch) loss_stats['l1_test_batch_loss'] = test_l1_loss test_l1_spec_loss /= len(test_batch) loss_stats['l1_spec_test_batch_loss'] = test_l1_spec_loss logger.log_test(iteration, loss_stats, verbose=args.verbose) logger.log_audios(iteration, audios) logger.log_specs(iteration, specs) logger.save_checkpoint( iteration, model if args.n_gpus == 1 else model.module, optimizer) if epoch % (epoch // 10 + 1) == 0: scheduler.step() except KeyboardInterrupt: print('KeyboardInterrupt: training has been stopped.') cleanup() return