def validate(model, epoch, total_iter, criterion, valset, batch_size, collate_fn, distributed_run, batch_to_gpu, ema=False): """Handles all the validation scoring and printing""" was_training = model.training model.eval() tik = time.perf_counter() with torch.no_grad(): val_sampler = DistributedSampler(valset) if distributed_run else None val_loader = DataLoader(valset, num_workers=4, shuffle=False, sampler=val_sampler, batch_size=batch_size, pin_memory=False, collate_fn=collate_fn) val_meta = defaultdict(float) val_num_frames = 0 for i, batch in enumerate(val_loader): x, y, num_frames = batch_to_gpu(batch) y_pred = model(x) loss, meta = criterion(y_pred, y, is_training=False, meta_agg='sum') if distributed_run: for k, v in meta.items(): val_meta[k] += reduce_tensor(v, 1) val_num_frames += reduce_tensor(num_frames.data, 1).item() else: for k, v in meta.items(): val_meta[k] += v val_num_frames = num_frames.item() val_meta = {k: v / len(valset) for k, v in val_meta.items()} val_meta['took'] = time.perf_counter() - tik log( (epoch, ) if epoch is not None else (), tb_total_steps=total_iter, subset='val_ema' if ema else 'val', data=OrderedDict([('loss', val_meta['loss'].item()), ('mel_loss', val_meta['mel_loss'].item()), ('frames/s', num_frames.item() / val_meta['took']), ('took', val_meta['took'])]), ) if was_training: model.train() return val_meta
def evaluate(epoch, step, val_loader, val_feat_proc, detokenize, ema_model, loss_fn, greedy_decoder, use_amp): ema_model.eval() start_time = time.time() agg = {'losses': [], 'preds': [], 'txts': [], 'idx': []} logging.log_start(logging.constants.EVAL_START, metadata=dict(epoch_num=epoch)) for i, batch in enumerate(val_loader): print(f'{val_loader.pipeline_type} evaluation: {i:>10}/{len(val_loader):<10}', end='\r') audio, audio_lens, txt, txt_lens = batch feats, feat_lens = val_feat_proc([audio, audio_lens]) log_probs, log_prob_lens = ema_model(feats, feat_lens, txt, txt_lens) loss = loss_fn(log_probs[:, :log_prob_lens.max().item()], log_prob_lens, txt, txt_lens) pred = greedy_decoder.decode(ema_model, feats, feat_lens) agg['losses'] += helpers.gather_losses([loss.cpu()]) agg['preds'] += helpers.gather_predictions([pred], detokenize) agg['txts'] += helpers.gather_transcripts([txt.cpu()], [txt_lens.cpu()], detokenize) wer, loss = process_evaluation_epoch(agg) logging.log_event(logging.constants.EVAL_ACCURACY, value=wer, metadata=dict(epoch_num=epoch)) logging.log_end(logging.constants.EVAL_STOP, metadata=dict(epoch_num=epoch)) log((epoch,), step, 'dev_ema', {'loss': loss, 'wer': 100.0 * wer, 'took': time.time() - start_time}) ema_model.train() return wer
def evaluate(epoch, step, val_loader, val_feat_proc, labels, model, ema_model, ctc_loss, greedy_decoder, use_amp, use_dali=False): for model, subset in [(model, 'dev'), (ema_model, 'dev_ema')]: if model is None: continue model.eval() start_time = time.time() agg = {'losses': [], 'preds': [], 'txts': []} for batch in val_loader: if use_dali: # with DALI, the data is already on GPU feat, feat_lens, txt, txt_lens = batch if val_feat_proc is not None: feat, feat_lens = val_feat_proc(feat, feat_lens, use_amp) else: batch = [t.cuda(non_blocking=True) for t in batch] audio, audio_lens, txt, txt_lens = batch feat, feat_lens = val_feat_proc(audio, audio_lens, use_amp) log_probs, enc_lens = model.forward(feat, feat_lens) loss = ctc_loss(log_probs, txt, enc_lens, txt_lens) pred = greedy_decoder(log_probs) agg['losses'] += helpers.gather_losses([loss]) agg['preds'] += helpers.gather_predictions([pred], labels) agg['txts'] += helpers.gather_transcripts([txt], [txt_lens], labels) wer, loss = process_evaluation_epoch(agg) log((epoch,), step, subset, {'loss': loss, 'wer': 100.0 * wer, 'took': time.time() - start_time}) model.train() return wer
def main(): parser = argparse.ArgumentParser(description='PyTorch FastPitch Training', allow_abbrev=False) parser = parse_args(parser) args, _ = parser.parse_known_args() distributed_run = args.world_size > 1 torch.manual_seed(args.seed + args.local_rank) np.random.seed(args.seed + args.local_rank) if args.local_rank == 0: if not os.path.exists(args.output): os.makedirs(args.output) log_fpath = args.log_file or os.path.join(args.output, 'nvlog.json') tb_subsets = ['train', 'val'] if args.ema_decay > 0.0: tb_subsets.append('val_ema') logger.init(log_fpath, args.output, enabled=(args.local_rank == 0), tb_subsets=tb_subsets) logger.parameters(vars(args), tb_subset='train') parser = models.parse_model_args('FastPitch', parser) args, unk_args = parser.parse_known_args() if len(unk_args) > 0: raise ValueError(f'Invalid options {unk_args}') torch.backends.cudnn.benchmark = args.cudnn_benchmark if distributed_run: init_distributed(args, args.world_size, args.local_rank) device = torch.device('cuda' if args.cuda else 'cpu') model_config = models.get_model_config('FastPitch', args) model = models.get_model('FastPitch', model_config, device) # Store pitch mean/std as params to translate from Hz during inference with open(args.pitch_mean_std_file, 'r') as f: stats = json.load(f) model.pitch_mean[0] = stats['mean'] model.pitch_std[0] = stats['std'] 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) else: raise ValueError scaler = torch.cuda.amp.GradScaler(enabled=args.amp) #if args.amp: #model, optimizer = amp.initialize(model, optimizer, opt_level="O1") if args.ema_decay > 0: ema_model = copy.deepcopy(model) else: ema_model = None 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] assert args.checkpoint_path is None or args.resume is False, ( "Specify a single checkpoint source") if args.checkpoint_path is not None: ch_fpath = args.checkpoint_path elif args.resume: ch_fpath = last_checkpoint(args.output) else: ch_fpath = None if ch_fpath is not None: load_checkpoint(args.local_rank, model, ema_model, optimizer, start_epoch, start_iter, model_config, args.amp, ch_fpath, args.world_size) start_epoch = start_epoch[0] total_iter = start_iter[0] criterion = loss_functions.get_loss_function( 'FastPitch', dur_predictor_loss_scale=args.dur_predictor_loss_scale, pitch_predictor_loss_scale=args.pitch_predictor_loss_scale) collate_fn = data_functions.get_collate_function('FastPitch') trainset = data_functions.get_data_loader('FastPitch', args.dataset_path, args.training_files, args) valset = data_functions.get_data_loader('FastPitch', args.dataset_path, args.validation_files, args) if distributed_run: train_sampler, shuffle = DistributedSampler(trainset), False else: train_sampler, shuffle = None, True train_loader = DataLoader(trainset, num_workers=16, shuffle=shuffle, sampler=train_sampler, batch_size=args.batch_size, pin_memory=False, drop_last=True, collate_fn=collate_fn) batch_to_gpu = data_functions.get_batch_to_gpu('FastPitch') model.train() torch.cuda.synchronize() for epoch in range(start_epoch, args.epochs + 1): epoch_start_time = time.perf_counter() epoch_loss = 0.0 epoch_mel_loss = 0.0 epoch_num_frames = 0 epoch_frames_per_sec = 0.0 if distributed_run: train_loader.sampler.set_epoch(epoch) accumulated_steps = 0 iter_loss = 0 iter_num_frames = 0 iter_meta = {} epoch_iter = 0 num_iters = len(train_loader) // args.gradient_accumulation_steps for batch in train_loader: if accumulated_steps == 0: if epoch_iter == num_iters: break total_iter += 1 epoch_iter += 1 iter_start_time = time.perf_counter() adjust_learning_rate(total_iter, optimizer, args.learning_rate, args.warmup_steps) model.zero_grad() x, y, num_frames = batch_to_gpu(batch) #AMP upstream autocast with torch.cuda.amp.autocast(enabled=args.amp): y_pred = model(x, use_gt_durations=True) loss, meta = criterion(y_pred, y) loss /= args.gradient_accumulation_steps meta = { k: v / args.gradient_accumulation_steps for k, v in meta.items() } if args.amp: #with amp.scale_loss(loss, optimizer) as scaled_loss: #scaled_loss.backward() scaler.scale(loss).backward() else: loss.backward() if distributed_run: reduced_loss = reduce_tensor(loss.data, args.world_size).item() reduced_num_frames = reduce_tensor(num_frames.data, 1).item() meta = { k: reduce_tensor(v, args.world_size) for k, v in meta.items() } else: reduced_loss = loss.item() reduced_num_frames = num_frames.item() if np.isnan(reduced_loss): raise Exception("loss is NaN") accumulated_steps += 1 iter_loss += reduced_loss iter_num_frames += reduced_num_frames iter_meta = {k: iter_meta.get(k, 0) + meta.get(k, 0) for k in meta} if accumulated_steps % args.gradient_accumulation_steps == 0: logger.log_grads_tb(total_iter, model) if args.amp: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_thresh) scaler.step(optimizer) scaler.update() #optimizer.zero_grad(set_to_none=True) optimizer.zero_grad() else: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_thresh) optimizer.step() apply_ema_decay(model, ema_model, args.ema_decay) iter_time = time.perf_counter() - iter_start_time iter_mel_loss = iter_meta['mel_loss'].item() epoch_frames_per_sec += iter_num_frames / iter_time epoch_loss += iter_loss epoch_num_frames += iter_num_frames epoch_mel_loss += iter_mel_loss logger.log( (epoch, epoch_iter, num_iters), tb_total_steps=total_iter, subset='train', data=OrderedDict([ ('loss', iter_loss), ('mel_loss', iter_mel_loss), ('frames/s', iter_num_frames / iter_time), ('took', iter_time), ('lrate', optimizer.param_groups[0]['lr']) ]), ) accumulated_steps = 0 iter_loss = 0 iter_num_frames = 0 iter_meta = {} # Finished epoch epoch_time = time.perf_counter() - epoch_start_time logger.log( (epoch, ), tb_total_steps=None, subset='train_avg', data=OrderedDict([('loss', epoch_loss / epoch_iter), ('mel_loss', epoch_mel_loss / epoch_iter), ('frames/s', epoch_num_frames / epoch_time), ('took', epoch_time)]), ) validate(model, epoch, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu, use_gt_durations=True) if args.ema_decay > 0: validate(ema_model, epoch, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu, use_gt_durations=True, ema=True) if (epoch > 0 and args.epochs_per_checkpoint > 0 and (epoch % args.epochs_per_checkpoint == 0) and args.local_rank == 0): checkpoint_path = os.path.join(args.output, f"FastPitch_checkpoint_{epoch}.pt") save_checkpoint(args.local_rank, model, ema_model, optimizer, scaler, epoch, total_iter, model_config, args.amp, checkpoint_path) logger.flush() # Finished training logger.log( (), tb_total_steps=None, subset='train_avg', data=OrderedDict([('loss', epoch_loss / epoch_iter), ('mel_loss', epoch_mel_loss / epoch_iter), ('frames/s', epoch_num_frames / epoch_time), ('took', epoch_time)]), ) validate(model, None, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu, use_gt_durations=True) if (epoch > 0 and args.epochs_per_checkpoint > 0 and (epoch % args.epochs_per_checkpoint != 0) and args.local_rank == 0): checkpoint_path = os.path.join(args.output, f"FastPitch_checkpoint_{epoch}.pt") save_checkpoint(args.local_rank, model, ema_model, optimizer, scaler, epoch, total_iter, model_config, args.amp, checkpoint_path)
def main(): parser = argparse.ArgumentParser(description='PyTorch FastPitch Training', allow_abbrev=False) parser = parse_args(parser) args, _ = parser.parse_known_args() if args.p_arpabet > 0.0: cmudict.initialize(args.cmudict_path, keep_ambiguous=True) distributed_run = args.world_size > 1 torch.manual_seed(args.seed + args.local_rank) np.random.seed(args.seed + args.local_rank) if args.local_rank == 0: if not os.path.exists(args.output): os.makedirs(args.output) log_fpath = args.log_file or os.path.join(args.output, 'nvlog.json') tb_subsets = ['train', 'val'] if args.ema_decay > 0.0: tb_subsets.append('val_ema') logger.init(log_fpath, args.output, enabled=(args.local_rank == 0), tb_subsets=tb_subsets) logger.parameters(vars(args), tb_subset='train') parser = models.parse_model_args('FastPitch', parser) args, unk_args = parser.parse_known_args() if len(unk_args) > 0: raise ValueError(f'Invalid options {unk_args}') torch.backends.cudnn.benchmark = args.cudnn_benchmark if distributed_run: init_distributed(args, args.world_size, args.local_rank) device = torch.device('cuda' if args.cuda else 'cpu') model_config = models.get_model_config('FastPitch', args) model = models.get_model('FastPitch', model_config, device) attention_kl_loss = AttentionBinarizationLoss() # Store pitch mean/std as params to translate from Hz during inference model.pitch_mean[0] = args.pitch_mean model.pitch_std[0] = args.pitch_std 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) else: raise ValueError scaler = torch.cuda.amp.GradScaler(enabled=args.amp) if args.ema_decay > 0: ema_model = copy.deepcopy(model) else: ema_model = None 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] assert args.checkpoint_path is None or args.resume is False, ( "Specify a single checkpoint source") if args.checkpoint_path is not None: ch_fpath = args.checkpoint_path elif args.resume: ch_fpath = last_checkpoint(args.output) else: ch_fpath = None if ch_fpath is not None: load_checkpoint(args, model, ema_model, optimizer, scaler, start_epoch, start_iter, model_config, ch_fpath) start_epoch = start_epoch[0] total_iter = start_iter[0] criterion = FastPitchLoss( dur_predictor_loss_scale=args.dur_predictor_loss_scale, pitch_predictor_loss_scale=args.pitch_predictor_loss_scale, attn_loss_scale=args.attn_loss_scale) collate_fn = TTSCollate() if args.local_rank == 0: prepare_tmp(args.pitch_online_dir) trainset = TTSDataset(audiopaths_and_text=args.training_files, **vars(args)) valset = TTSDataset(audiopaths_and_text=args.validation_files, **vars(args)) if distributed_run: train_sampler, shuffle = DistributedSampler(trainset), False else: train_sampler, shuffle = None, True # 4 workers are optimal on DGX-1 (from epoch 2 onwards) train_loader = DataLoader(trainset, num_workers=4, shuffle=shuffle, sampler=train_sampler, batch_size=args.batch_size, pin_memory=True, persistent_workers=True, drop_last=True, collate_fn=collate_fn) if args.ema_decay: mt_ema_params = init_multi_tensor_ema(model, ema_model) model.train() bmark_stats = BenchmarkStats() torch.cuda.synchronize() for epoch in range(start_epoch, args.epochs + 1): epoch_start_time = time.perf_counter() epoch_loss = 0.0 epoch_mel_loss = 0.0 epoch_num_frames = 0 epoch_frames_per_sec = 0.0 if distributed_run: train_loader.sampler.set_epoch(epoch) accumulated_steps = 0 iter_loss = 0 iter_num_frames = 0 iter_meta = {} iter_start_time = time.perf_counter() epoch_iter = 0 num_iters = len(train_loader) // args.grad_accumulation for batch in train_loader: if accumulated_steps == 0: if epoch_iter == num_iters: break total_iter += 1 epoch_iter += 1 adjust_learning_rate(total_iter, optimizer, args.learning_rate, args.warmup_steps) model.zero_grad(set_to_none=True) x, y, num_frames = batch_to_gpu(batch) with torch.cuda.amp.autocast(enabled=args.amp): y_pred = model(x) loss, meta = criterion(y_pred, y) if (args.kl_loss_start_epoch is not None and epoch >= args.kl_loss_start_epoch): if args.kl_loss_start_epoch == epoch and epoch_iter == 1: print('Begin hard_attn loss') _, _, _, _, _, _, _, _, attn_soft, attn_hard, _, _ = y_pred binarization_loss = attention_kl_loss(attn_hard, attn_soft) kl_weight = min( (epoch - args.kl_loss_start_epoch) / args.kl_loss_warmup_epochs, 1.0) * args.kl_loss_weight meta['kl_loss'] = binarization_loss.clone().detach( ) * kl_weight loss += kl_weight * binarization_loss else: meta['kl_loss'] = torch.zeros_like(loss) kl_weight = 0 binarization_loss = 0 loss /= args.grad_accumulation meta = {k: v / args.grad_accumulation for k, v in meta.items()} if args.amp: scaler.scale(loss).backward() else: loss.backward() if distributed_run: reduced_loss = reduce_tensor(loss.data, args.world_size).item() reduced_num_frames = reduce_tensor(num_frames.data, 1).item() meta = { k: reduce_tensor(v, args.world_size) for k, v in meta.items() } else: reduced_loss = loss.item() reduced_num_frames = num_frames.item() if np.isnan(reduced_loss): raise Exception("loss is NaN") accumulated_steps += 1 iter_loss += reduced_loss iter_num_frames += reduced_num_frames iter_meta = {k: iter_meta.get(k, 0) + meta.get(k, 0) for k in meta} if accumulated_steps % args.grad_accumulation == 0: logger.log_grads_tb(total_iter, model) if args.amp: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_thresh) scaler.step(optimizer) scaler.update() else: torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_thresh) optimizer.step() if args.ema_decay > 0.0: apply_multi_tensor_ema(args.ema_decay, *mt_ema_params) iter_mel_loss = iter_meta['mel_loss'].item() iter_kl_loss = iter_meta['kl_loss'].item() iter_time = time.perf_counter() - iter_start_time epoch_frames_per_sec += iter_num_frames / iter_time epoch_loss += iter_loss epoch_num_frames += iter_num_frames epoch_mel_loss += iter_mel_loss log( (epoch, epoch_iter, num_iters), tb_total_steps=total_iter, subset='train', data=OrderedDict([ ('loss', iter_loss), ('mel_loss', iter_mel_loss), ('kl_loss', iter_kl_loss), ('kl_weight', kl_weight), ('frames/s', iter_num_frames / iter_time), ('took', iter_time), ('lrate', optimizer.param_groups[0]['lr']) ]), ) accumulated_steps = 0 iter_loss = 0 iter_num_frames = 0 iter_meta = {} iter_start_time = time.perf_counter() # Finished epoch epoch_loss /= epoch_iter epoch_mel_loss /= epoch_iter epoch_time = time.perf_counter() - epoch_start_time log( (epoch, ), tb_total_steps=None, subset='train_avg', data=OrderedDict([('loss', epoch_loss), ('mel_loss', epoch_mel_loss), ('frames/s', epoch_num_frames / epoch_time), ('took', epoch_time)]), ) bmark_stats.update(epoch_num_frames, epoch_loss, epoch_mel_loss, epoch_time) validate(model, epoch, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu) if args.ema_decay > 0: validate(ema_model, epoch, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu, ema=True) maybe_save_checkpoint(args, model, ema_model, optimizer, scaler, epoch, total_iter, model_config) logger.flush() # Finished training if len(bmark_stats) > 0: log((), tb_total_steps=None, subset='train_avg', data=bmark_stats.get(args.benchmark_epochs_num)) validate(model, None, total_iter, criterion, valset, args.batch_size, collate_fn, distributed_run, batch_to_gpu)
def main(): args = parse_args() assert (torch.cuda.is_available()) assert args.prediction_frequency % args.log_frequency == 0 torch.backends.cudnn.benchmark = args.cudnn_benchmark # set up distributed training multi_gpu = int(os.environ.get('WORLD_SIZE', 1)) > 1 if multi_gpu: torch.cuda.set_device(args.local_rank) dist.init_process_group(backend='nccl', init_method='env://') world_size = dist.get_world_size() print_once(f'Distributed training with {world_size} GPUs\n') else: world_size = 1 torch.manual_seed(args.seed + args.local_rank) np.random.seed(args.seed + args.local_rank) random.seed(args.seed + args.local_rank) init_log(args) cfg = config.load(args.model_config) config.apply_config_overrides(cfg, args) symbols = helpers.add_ctc_blank(cfg['labels']) assert args.grad_accumulation_steps >= 1 assert args.batch_size % args.grad_accumulation_steps == 0 batch_size = args.batch_size // args.grad_accumulation_steps print_once('Setting up datasets...') train_dataset_kw, train_features_kw = config.input(cfg, 'train') val_dataset_kw, val_features_kw = config.input(cfg, 'val') use_dali = args.dali_device in ('cpu', 'gpu') if use_dali: assert train_dataset_kw['ignore_offline_speed_perturbation'], \ "DALI doesn't support offline speed perturbation" # pad_to_max_duration is not supported by DALI - have simple padders if train_features_kw['pad_to_max_duration']: train_feat_proc = BaseFeatures( pad_align=train_features_kw['pad_align'], pad_to_max_duration=True, max_duration=train_features_kw['max_duration'], sample_rate=train_features_kw['sample_rate'], window_size=train_features_kw['window_size'], window_stride=train_features_kw['window_stride']) train_features_kw['pad_to_max_duration'] = False else: train_feat_proc = None if val_features_kw['pad_to_max_duration']: val_feat_proc = BaseFeatures( pad_align=val_features_kw['pad_align'], pad_to_max_duration=True, max_duration=val_features_kw['max_duration'], sample_rate=val_features_kw['sample_rate'], window_size=val_features_kw['window_size'], window_stride=val_features_kw['window_stride']) val_features_kw['pad_to_max_duration'] = False else: val_feat_proc = None train_loader = DaliDataLoader( gpu_id=args.local_rank, dataset_path=args.dataset_dir, config_data=train_dataset_kw, config_features=train_features_kw, json_names=args.train_manifests, batch_size=batch_size, grad_accumulation_steps=args.grad_accumulation_steps, pipeline_type="train", device_type=args.dali_device, symbols=symbols) val_loader = DaliDataLoader(gpu_id=args.local_rank, dataset_path=args.dataset_dir, config_data=val_dataset_kw, config_features=val_features_kw, json_names=args.val_manifests, batch_size=batch_size, pipeline_type="val", device_type=args.dali_device, symbols=symbols) else: train_dataset_kw, train_features_kw = config.input(cfg, 'train') train_dataset = AudioDataset(args.dataset_dir, args.train_manifests, symbols, **train_dataset_kw) train_loader = get_data_loader(train_dataset, batch_size, multi_gpu=multi_gpu, shuffle=True, num_workers=4) train_feat_proc = FilterbankFeatures(**train_features_kw) val_dataset_kw, val_features_kw = config.input(cfg, 'val') val_dataset = AudioDataset(args.dataset_dir, args.val_manifests, symbols, **val_dataset_kw) val_loader = get_data_loader(val_dataset, batch_size, multi_gpu=multi_gpu, shuffle=False, num_workers=4, drop_last=False) val_feat_proc = FilterbankFeatures(**val_features_kw) dur = train_dataset.duration / 3600 dur_f = train_dataset.duration_filtered / 3600 nsampl = len(train_dataset) print_once(f'Training samples: {nsampl} ({dur:.1f}h, ' f'filtered {dur_f:.1f}h)') if train_feat_proc is not None: train_feat_proc.cuda() if val_feat_proc is not None: val_feat_proc.cuda() steps_per_epoch = len(train_loader) // args.grad_accumulation_steps # set up the model model = Jasper(encoder_kw=config.encoder(cfg), decoder_kw=config.decoder(cfg, n_classes=len(symbols))) model.cuda() ctc_loss = CTCLossNM(n_classes=len(symbols)) greedy_decoder = GreedyCTCDecoder() print_once(f'Model size: {num_weights(model) / 10**6:.1f}M params\n') # optimization kw = {'lr': args.lr, 'weight_decay': args.weight_decay} if args.optimizer == "novograd": optimizer = Novograd(model.parameters(), **kw) elif args.optimizer == "adamw": optimizer = AdamW(model.parameters(), **kw) else: raise ValueError(f'Invalid optimizer "{args.optimizer}"') scaler = torch.cuda.amp.GradScaler(enabled=args.amp) adjust_lr = lambda step, epoch, optimizer: lr_policy( step, epoch, args.lr, optimizer, steps_per_epoch=steps_per_epoch, warmup_epochs=args.warmup_epochs, hold_epochs=args.hold_epochs, num_epochs=args.epochs, policy=args.lr_policy, min_lr=args.min_lr, exp_gamma=args.lr_exp_gamma) if args.ema > 0: ema_model = copy.deepcopy(model) else: ema_model = None if multi_gpu: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.local_rank], output_device=args.local_rank) if args.pyprof: pyprof.init(enable_function_stack=True) # load checkpoint meta = {'best_wer': 10**6, 'start_epoch': 0} checkpointer = Checkpointer(args.output_dir, 'Jasper', args.keep_milestones) if args.resume: args.ckpt = checkpointer.last_checkpoint() or args.ckpt if args.ckpt is not None: checkpointer.load(args.ckpt, model, ema_model, optimizer, scaler, meta) start_epoch = meta['start_epoch'] best_wer = meta['best_wer'] epoch = 1 step = start_epoch * steps_per_epoch + 1 if args.pyprof: torch.autograd.profiler.emit_nvtx().__enter__() profiler.start() # training loop model.train() # pre-allocate if args.pre_allocate_range is not None: n_feats = train_features_kw['n_filt'] pad_align = train_features_kw['pad_align'] a, b = args.pre_allocate_range for n_frames in range(a, b + pad_align, pad_align): print_once( f'Pre-allocation ({batch_size}x{n_feats}x{n_frames})...') feat = torch.randn(batch_size, n_feats, n_frames, device='cuda') feat_lens = torch.ones(batch_size, device='cuda').fill_(n_frames) txt = torch.randint(high=len(symbols) - 1, size=(batch_size, 100), device='cuda') txt_lens = torch.ones(batch_size, device='cuda').fill_(100) with torch.cuda.amp.autocast(enabled=args.amp): log_probs, enc_lens = model(feat, feat_lens) del feat loss = ctc_loss(log_probs, txt, enc_lens, txt_lens) loss.backward() model.zero_grad() torch.cuda.empty_cache() bmark_stats = BenchmarkStats() for epoch in range(start_epoch + 1, args.epochs + 1): if multi_gpu and not use_dali: train_loader.sampler.set_epoch(epoch) epoch_utts = 0 epoch_loss = 0 accumulated_batches = 0 epoch_start_time = time.time() epoch_eval_time = 0 for batch in train_loader: if accumulated_batches == 0: step_loss = 0 step_utts = 0 step_start_time = time.time() if use_dali: # with DALI, the data is already on GPU feat, feat_lens, txt, txt_lens = batch if train_feat_proc is not None: feat, feat_lens = train_feat_proc(feat, feat_lens) else: batch = [t.cuda(non_blocking=True) for t in batch] audio, audio_lens, txt, txt_lens = batch feat, feat_lens = train_feat_proc(audio, audio_lens) # Use context manager to prevent redundant accumulation of gradients if (multi_gpu and accumulated_batches + 1 < args.grad_accumulation_steps): ctx = model.no_sync() else: ctx = empty_context() with ctx: with torch.cuda.amp.autocast(enabled=args.amp): log_probs, enc_lens = model(feat, feat_lens) loss = ctc_loss(log_probs, txt, enc_lens, txt_lens) loss /= args.grad_accumulation_steps if multi_gpu: reduced_loss = reduce_tensor(loss.data, world_size) else: reduced_loss = loss if torch.isnan(reduced_loss).any(): print_once(f'WARNING: loss is NaN; skipping update') continue else: step_loss += reduced_loss.item() step_utts += batch[0].size(0) * world_size epoch_utts += batch[0].size(0) * world_size accumulated_batches += 1 scaler.scale(loss).backward() if accumulated_batches % args.grad_accumulation_steps == 0: epoch_loss += step_loss scaler.step(optimizer) scaler.update() adjust_lr(step, epoch, optimizer) optimizer.zero_grad() apply_ema(model, ema_model, args.ema) if step % args.log_frequency == 0: preds = greedy_decoder(log_probs) wer, pred_utt, ref = greedy_wer(preds, txt, txt_lens, symbols) if step % args.prediction_frequency == 0: print_once(f' Decoded: {pred_utt[:90]}') print_once(f' Reference: {ref[:90]}') step_time = time.time() - step_start_time log( (epoch, step % steps_per_epoch or steps_per_epoch, steps_per_epoch), step, 'train', { 'loss': step_loss, 'wer': 100.0 * wer, 'throughput': step_utts / step_time, 'took': step_time, 'lrate': optimizer.param_groups[0]['lr'] }) step_start_time = time.time() if step % args.eval_frequency == 0: tik = time.time() wer = evaluate(epoch, step, val_loader, val_feat_proc, symbols, model, ema_model, ctc_loss, greedy_decoder, args.amp, use_dali) if wer < best_wer and epoch >= args.save_best_from: checkpointer.save(model, ema_model, optimizer, scaler, epoch, step, best_wer, is_best=True) best_wer = wer epoch_eval_time += time.time() - tik step += 1 accumulated_batches = 0 # end of step # DALI iterator need to be exhausted; # if not using DALI, simulate drop_last=True with grad accumulation if not use_dali and step > steps_per_epoch * epoch: break epoch_time = time.time() - epoch_start_time epoch_loss /= steps_per_epoch log( (epoch, ), None, 'train_avg', { 'throughput': epoch_utts / epoch_time, 'took': epoch_time, 'loss': epoch_loss }) bmark_stats.update(epoch_utts, epoch_time, epoch_loss) if epoch % args.save_frequency == 0 or epoch in args.keep_milestones: checkpointer.save(model, ema_model, optimizer, scaler, epoch, step, best_wer) if 0 < args.epochs_this_job <= epoch - start_epoch: print_once(f'Finished after {args.epochs_this_job} epochs.') break # end of epoch if args.pyprof: profiler.stop() torch.autograd.profiler.emit_nvtx().__exit__(None, None, None) log((), None, 'train_avg', bmark_stats.get(args.benchmark_epochs_num)) if epoch == args.epochs: evaluate(epoch, step, val_loader, val_feat_proc, symbols, model, ema_model, ctc_loss, greedy_decoder, args.amp, use_dali) checkpointer.save(model, ema_model, optimizer, scaler, epoch, step, best_wer) flush_log()
def main(): logging.configure_logger('RNNT') logging.log_start(logging.constants.INIT_START) args = parse_args() assert(torch.cuda.is_available()) assert args.prediction_frequency is None or args.prediction_frequency % args.log_frequency == 0 torch.backends.cudnn.benchmark = args.cudnn_benchmark # set up distributed training multi_gpu = int(os.environ.get('WORLD_SIZE', 1)) > 1 if multi_gpu: torch.cuda.set_device(args.local_rank) dist.init_process_group(backend='nccl', init_method='env://') world_size = dist.get_world_size() print_once(f'Distributed training with {world_size} GPUs\n') else: world_size = 1 if args.seed is not None: logging.log_event(logging.constants.SEED, value=args.seed) torch.manual_seed(args.seed + args.local_rank) np.random.seed(args.seed + args.local_rank) random.seed(args.seed + args.local_rank) # np_rng is used for buckets generation, and needs the same seed on every worker np_rng = np.random.default_rng(seed=args.seed) init_log(args) cfg = config.load(args.model_config) config.apply_duration_flags(cfg, args.max_duration) assert args.grad_accumulation_steps >= 1 assert args.batch_size % args.grad_accumulation_steps == 0, f'{args.batch_size} % {args.grad_accumulation_steps} != 0' logging.log_event(logging.constants.GRADIENT_ACCUMULATION_STEPS, value=args.grad_accumulation_steps) batch_size = args.batch_size // args.grad_accumulation_steps logging.log_event(logging.constants.SUBMISSION_BENCHMARK, value=logging.constants.RNNT) logging.log_event(logging.constants.SUBMISSION_ORG, value='my-organization') logging.log_event(logging.constants.SUBMISSION_DIVISION, value=logging.constants.CLOSED) # closed or open logging.log_event(logging.constants.SUBMISSION_STATUS, value=logging.constants.ONPREM) # on-prem/cloud/research logging.log_event(logging.constants.SUBMISSION_PLATFORM, value='my platform') logging.log_end(logging.constants.INIT_STOP) if multi_gpu: torch.distributed.barrier() logging.log_start(logging.constants.RUN_START) if multi_gpu: torch.distributed.barrier() print_once('Setting up datasets...') ( train_dataset_kw, train_features_kw, train_splicing_kw, train_specaugm_kw, ) = config.input(cfg, 'train') ( val_dataset_kw, val_features_kw, val_splicing_kw, val_specaugm_kw, ) = config.input(cfg, 'val') logging.log_event(logging.constants.DATA_TRAIN_MAX_DURATION, value=train_dataset_kw['max_duration']) logging.log_event(logging.constants.DATA_SPEED_PERTURBATON_MAX, value=train_dataset_kw['speed_perturbation']['max_rate']) logging.log_event(logging.constants.DATA_SPEED_PERTURBATON_MIN, value=train_dataset_kw['speed_perturbation']['min_rate']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_FREQ_N, value=train_specaugm_kw['freq_masks']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_FREQ_MIN, value=train_specaugm_kw['min_freq']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_FREQ_MAX, value=train_specaugm_kw['max_freq']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_TIME_N, value=train_specaugm_kw['time_masks']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_TIME_MIN, value=train_specaugm_kw['min_time']) logging.log_event(logging.constants.DATA_SPEC_AUGMENT_TIME_MAX, value=train_specaugm_kw['max_time']) logging.log_event(logging.constants.GLOBAL_BATCH_SIZE, value=batch_size * world_size * args.grad_accumulation_steps) tokenizer_kw = config.tokenizer(cfg) tokenizer = Tokenizer(**tokenizer_kw) class PermuteAudio(torch.nn.Module): def forward(self, x): return (x[0].permute(2, 0, 1), *x[1:]) train_augmentations = torch.nn.Sequential( train_specaugm_kw and features.SpecAugment(optim_level=args.amp, **train_specaugm_kw) or torch.nn.Identity(), features.FrameSplicing(optim_level=args.amp, **train_splicing_kw), PermuteAudio(), ) val_augmentations = torch.nn.Sequential( val_specaugm_kw and features.SpecAugment(optim_level=args.amp, **val_specaugm_kw) or torch.nn.Identity(), features.FrameSplicing(optim_level=args.amp, **val_splicing_kw), PermuteAudio(), ) logging.log_event(logging.constants.DATA_TRAIN_NUM_BUCKETS, value=args.num_buckets) if args.num_buckets is not None: sampler = dali_sampler.BucketingSampler( args.num_buckets, batch_size, world_size, args.epochs, np_rng ) else: sampler = dali_sampler.SimpleSampler() train_loader = DaliDataLoader(gpu_id=args.local_rank, dataset_path=args.dataset_dir, config_data=train_dataset_kw, config_features=train_features_kw, json_names=args.train_manifests, batch_size=batch_size, sampler=sampler, grad_accumulation_steps=args.grad_accumulation_steps, pipeline_type="train", device_type=args.dali_device, tokenizer=tokenizer) val_loader = DaliDataLoader(gpu_id=args.local_rank, dataset_path=args.dataset_dir, config_data=val_dataset_kw, config_features=val_features_kw, json_names=args.val_manifests, batch_size=args.val_batch_size, sampler=dali_sampler.SimpleSampler(), pipeline_type="val", device_type=args.dali_device, tokenizer=tokenizer) train_feat_proc = train_augmentations val_feat_proc = val_augmentations train_feat_proc.cuda() val_feat_proc.cuda() steps_per_epoch = len(train_loader) // args.grad_accumulation_steps logging.log_event(logging.constants.TRAIN_SAMPLES, value=train_loader.dataset_size) logging.log_event(logging.constants.EVAL_SAMPLES, value=val_loader.dataset_size) # set up the model rnnt_config = config.rnnt(cfg) logging.log_event(logging.constants.MODEL_WEIGHTS_INITIALIZATION_SCALE, value=args.weights_init_scale) if args.weights_init_scale is not None: rnnt_config['weights_init_scale'] = args.weights_init_scale if args.hidden_hidden_bias_scale is not None: rnnt_config['hidden_hidden_bias_scale'] = args.hidden_hidden_bias_scale model = RNNT(n_classes=tokenizer.num_labels + 1, **rnnt_config) model.cuda() blank_idx = tokenizer.num_labels loss_fn = RNNTLoss(blank_idx=blank_idx) logging.log_event(logging.constants.EVAL_MAX_PREDICTION_SYMBOLS, value=args.max_symbol_per_sample) greedy_decoder = RNNTGreedyDecoder( blank_idx=blank_idx, max_symbol_per_sample=args.max_symbol_per_sample) print_once(f'Model size: {num_weights(model) / 10**6:.1f}M params\n') opt_eps=1e-9 logging.log_event(logging.constants.OPT_NAME, value='lamb') logging.log_event(logging.constants.OPT_BASE_LR, value=args.lr) logging.log_event(logging.constants.OPT_LAMB_EPSILON, value=opt_eps) logging.log_event(logging.constants.OPT_LAMB_LR_DECAY_POLY_POWER, value=args.lr_exp_gamma) logging.log_event(logging.constants.OPT_LR_WARMUP_EPOCHS, value=args.warmup_epochs) logging.log_event(logging.constants.OPT_LAMB_LR_HOLD_EPOCHS, value=args.hold_epochs) logging.log_event(logging.constants.OPT_LAMB_BETA_1, value=args.beta1) logging.log_event(logging.constants.OPT_LAMB_BETA_2, value=args.beta2) logging.log_event(logging.constants.OPT_GRADIENT_CLIP_NORM, value=args.clip_norm) logging.log_event(logging.constants.OPT_LR_ALT_DECAY_FUNC, value=True) logging.log_event(logging.constants.OPT_LR_ALT_WARMUP_FUNC, value=True) logging.log_event(logging.constants.OPT_LAMB_LR_MIN, value=args.min_lr) logging.log_event(logging.constants.OPT_WEIGHT_DECAY, value=args.weight_decay) # optimization kw = {'params': model.param_groups(args.lr), 'lr': args.lr, 'weight_decay': args.weight_decay} initial_lrs = [group['lr'] for group in kw['params']] print_once(f'Starting with LRs: {initial_lrs}') optimizer = FusedLAMB(betas=(args.beta1, args.beta2), eps=opt_eps, **kw) adjust_lr = lambda step, epoch: lr_policy( step, epoch, initial_lrs, optimizer, steps_per_epoch=steps_per_epoch, warmup_epochs=args.warmup_epochs, hold_epochs=args.hold_epochs, min_lr=args.min_lr, exp_gamma=args.lr_exp_gamma) if args.amp: model, optimizer = amp.initialize( models=model, optimizers=optimizer, opt_level='O1', max_loss_scale=512.0) if args.ema > 0: ema_model = copy.deepcopy(model).cuda() else: ema_model = None logging.log_event(logging.constants.MODEL_EVAL_EMA_FACTOR, value=args.ema) if multi_gpu: model = DistributedDataParallel(model) # load checkpoint meta = {'best_wer': 10**6, 'start_epoch': 0} checkpointer = Checkpointer(args.output_dir, 'RNN-T', args.keep_milestones, args.amp) if args.resume: args.ckpt = checkpointer.last_checkpoint() or args.ckpt if args.ckpt is not None: checkpointer.load(args.ckpt, model, ema_model, optimizer, meta) start_epoch = meta['start_epoch'] best_wer = meta['best_wer'] last_wer = meta['best_wer'] epoch = 1 step = start_epoch * steps_per_epoch + 1 # training loop model.train() for epoch in range(start_epoch + 1, args.epochs + 1): logging.log_start(logging.constants.BLOCK_START, metadata=dict(first_epoch_num=epoch, epoch_count=1)) logging.log_start(logging.constants.EPOCH_START, metadata=dict(epoch_num=epoch)) epoch_utts = 0 accumulated_batches = 0 epoch_start_time = time.time() for batch in train_loader: if accumulated_batches == 0: adjust_lr(step, epoch) optimizer.zero_grad() step_utts = 0 step_start_time = time.time() all_feat_lens = [] audio, audio_lens, txt, txt_lens = batch feats, feat_lens = train_feat_proc([audio, audio_lens]) all_feat_lens += feat_lens log_probs, log_prob_lens = model(feats, feat_lens, txt, txt_lens) loss = loss_fn(log_probs[:, :log_prob_lens.max().item()], log_prob_lens, txt, txt_lens) loss /= args.grad_accumulation_steps del log_probs, log_prob_lens if torch.isnan(loss).any(): print_once(f'WARNING: loss is NaN; skipping update') else: if args.amp: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() loss_item = loss.item() del loss step_utts += batch[0].size(0) * world_size epoch_utts += batch[0].size(0) * world_size accumulated_batches += 1 if accumulated_batches % args.grad_accumulation_steps == 0: if args.clip_norm is not None: torch.nn.utils.clip_grad_norm_( getattr(model, 'module', model).parameters(), max_norm=args.clip_norm, norm_type=2) total_norm = 0.0 try: if args.log_norm: for p in getattr(model, 'module', model).parameters(): param_norm = p.grad.data.norm(2) total_norm += param_norm.item() ** 2 total_norm = total_norm ** (1. / 2) except AttributeError as e: print_once(f'Exception happened: {e}') total_norm = 0.0 optimizer.step() apply_ema(model, ema_model, args.ema) if step % args.log_frequency == 0: if args.prediction_frequency is None or step % args.prediction_frequency == 0: preds = greedy_decoder.decode(model, feats, feat_lens) wer, pred_utt, ref = greedy_wer( preds, txt, txt_lens, tokenizer.detokenize) print_once(f' Decoded: {pred_utt[:90]}') print_once(f' Reference: {ref[:90]}') wer = {'wer': 100 * wer} else: wer = {} step_time = time.time() - step_start_time log((epoch, step % steps_per_epoch or steps_per_epoch, steps_per_epoch), step, 'train', {'loss': loss_item, **wer, # optional entry 'throughput': step_utts / step_time, 'took': step_time, 'grad-norm': total_norm, 'seq-len-min': min(all_feat_lens).item(), 'seq-len-max': max(all_feat_lens).item(), 'lrate': optimizer.param_groups[0]['lr']}) step_start_time = time.time() step += 1 accumulated_batches = 0 # end of step logging.log_end(logging.constants.EPOCH_STOP, metadata=dict(epoch_num=epoch)) epoch_time = time.time() - epoch_start_time log((epoch,), None, 'train_avg', {'throughput': epoch_utts / epoch_time, 'took': epoch_time}) if epoch % args.val_frequency == 0: wer = evaluate(epoch, step, val_loader, val_feat_proc, tokenizer.detokenize, ema_model, loss_fn, greedy_decoder, args.amp) last_wer = wer if wer < best_wer and epoch >= args.save_best_from: checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer, is_best=True) best_wer = wer save_this_epoch = (args.save_frequency is not None and epoch % args.save_frequency == 0) \ or (epoch in args.keep_milestones) if save_this_epoch: checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer) logging.log_end(logging.constants.BLOCK_STOP, metadata=dict(first_epoch_num=epoch)) if last_wer <= args.target: logging.log_end(logging.constants.RUN_STOP, metadata={'status': 'success'}) print_once(f'Finished after {args.epochs_this_job} epochs.') break if 0 < args.epochs_this_job <= epoch - start_epoch: print_once(f'Finished after {args.epochs_this_job} epochs.') break # end of epoch log((), None, 'train_avg', {'throughput': epoch_utts / epoch_time}) if last_wer > args.target: logging.log_end(logging.constants.RUN_STOP, metadata={'status': 'aborted'}) if epoch == args.epochs: evaluate(epoch, step, val_loader, val_feat_proc, tokenizer.detokenize, ema_model, loss_fn, greedy_decoder, args.amp) flush_log() if args.save_at_the_end: checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)