def main():
parser = get_parser()
args = parser.parse_args()
log_fpath = args.log_file or str(Path(args.output_dir, 'nvlog_infer.json'))
log_fpath = unique_log_fpath(log_fpath)
dllogger.init(backends=[
JSONStreamBackend(Verbosity.DEFAULT, log_fpath),
StdOutBackend(Verbosity.VERBOSE, metric_format=stdout_metric_format)
])
[dllogger.log("PARAMETER", {k: v}) for k, v in vars(args).items()]
for step in ['DNN', 'data+DNN', 'data']:
for c in [0.99, 0.95, 0.9, 0.5]:
cs = 'avg' if c == 0.5 else f'{int(100*c)}%'
dllogger.metadata(f'{step.lower()}_latency_{c}', {
'name': f'{step} latency {cs}',
'format': ':>7.2f',
'unit': 'ms'
})
dllogger.metadata('eval_wer', {
'name': 'WER',
'format': ':>3.3f',
'unit': '%'
})
if args.cpu:
device = torch.device('cpu')
else:
assert torch.cuda.is_available()
device = torch.device('cuda')
torch.backends.cudnn.benchmark = args.cudnn_benchmark
if args.seed is not None:
torch.manual_seed(args.seed + args.local_rank)
np.random.seed(args.seed + args.local_rank)
random.seed(args.seed + args.local_rank)
# set up distributed training
multi_gpu = not args.cpu and int(os.environ.get('WORLD_SIZE', 1)) > 1
if multi_gpu:
torch.cuda.set_device(args.local_rank)
distrib.init_process_group(backend='nccl', init_method='env://')
print_once(f'Inference with {distrib.get_world_size()} GPUs')
cfg = config.load(args.model_config)
if args.max_duration is not None:
cfg['input_val']['audio_dataset']['max_duration'] = args.max_duration
cfg['input_val']['filterbank_features'][
'max_duration'] = args.max_duration
if args.pad_to_max_duration:
assert cfg['input_val']['audio_dataset']['max_duration'] > 0
cfg['input_val']['audio_dataset']['pad_to_max_duration'] = True
cfg['input_val']['filterbank_features']['pad_to_max_duration'] = True
use_dali = args.dali_device in ('cpu', 'gpu')
(dataset_kw, features_kw, splicing_kw, _, _) = config.input(cfg, 'val')
tokenizer_kw = config.tokenizer(cfg)
tokenizer = Tokenizer(**tokenizer_kw)
optim_level = 3 if args.amp else 0
feature_proc = torch.nn.Sequential(
torch.nn.Identity(),
torch.nn.Identity(),
features.FrameSplicing(optim_level=optim_level, **splicing_kw),
features.FillPadding(optim_level=optim_level, ),
)
# dataset
data_loader = DaliDataLoader(gpu_id=args.local_rank or 0,
dataset_path=args.dataset_dir,
config_data=dataset_kw,
config_features=features_kw,
json_names=[args.val_manifest],
batch_size=args.batch_size,
sampler=dali_sampler.SimpleSampler(),
pipeline_type="val",
device_type=args.dali_device,
tokenizer=tokenizer)
model = RNNT(n_classes=tokenizer.num_labels + 1, **config.rnnt(cfg))
if args.ckpt is not None:
print(f'Loading the model from {args.ckpt} ...')
checkpoint = torch.load(args.ckpt, map_location="cpu")
key = 'ema_state_dict' if args.ema else 'state_dict'
state_dict = checkpoint[key]
model.load_state_dict(state_dict, strict=True)
model.to(device)
model.eval()
if feature_proc is not None:
feature_proc.to(device)
feature_proc.eval()
if args.amp:
model = amp.initialize(model, opt_level='O3')
if multi_gpu:
model = DistributedDataParallel(model)
agg = {'txts': [], 'preds': [], 'logits': []}
dur = {'data': [], 'dnn': [], 'data+dnn': []}
rep_loader = chain(*repeat(data_loader, args.repeats))
rep_len = args.repeats * len(data_loader)
blank_idx = tokenizer.num_labels
greedy_decoder = RNNTGreedyDecoder(blank_idx=blank_idx)
def sync_time():
torch.cuda.synchronize() if device.type == 'cuda' else None
return time.perf_counter()
sz = []
with torch.no_grad():
for it, batch in enumerate(tqdm.tqdm(rep_loader, total=rep_len)):
if use_dali:
feats, feat_lens, txt, txt_lens = batch
if feature_proc is not None:
feats, feat_lens = feature_proc([feats, feat_lens])
else:
batch = [t.cuda(non_blocking=True) for t in batch]
audio, audio_lens, txt, txt_lens = batch
feats, feat_lens = feature_proc([audio, audio_lens])
feats = feats.permute(2, 0, 1)
if args.amp:
feats = feats.half()
sz.append(feats.size(0))
t1 = sync_time()
log_probs, log_prob_lens = model(feats, feat_lens, txt, txt_lens)
t2 = sync_time()
# burn-in period; wait for a new loader due to num_workers
if it >= 1 and (args.repeats == 1 or it >= len(data_loader)):
dur['data'].append(t1 - t0)
dur['dnn'].append(t2 - t1)
dur['data+dnn'].append(t2 - t0)
if txt is not None:
agg['txts'] += helpers.gather_transcripts([txt], [txt_lens],
tokenizer.detokenize)
preds = greedy_decoder.decode(model, feats, feat_lens)
agg['preds'] += helpers.gather_predictions([preds],
tokenizer.detokenize)
if 0 < args.steps < it:
break
t0 = sync_time()
# communicate the results
if args.transcribe_wav:
for idx, p in enumerate(agg['preds']):
print_once(f'Prediction {idx+1: >3}: {p}')
elif args.transcribe_filelist:
pass
else:
wer, loss = process_evaluation_epoch(agg)
if not multi_gpu or distrib.get_rank() == 0:
dllogger.log(step=(), data={'eval_wer': 100 * wer})
if args.save_predictions:
with open(args.save_predictions, 'w') as f:
f.write('\n'.join(agg['preds']))
# report timings
if len(dur['data']) >= 20:
ratios = [0.9, 0.95, 0.99]
for stage in dur:
lat = durs_to_percentiles(dur[stage], ratios)
for k in [0.99, 0.95, 0.9, 0.5]:
kk = str(k).replace('.', '_')
dllogger.log(step=(),
data={f'{stage.lower()}_latency_{kk}': lat[k]})
else:
# TODO measure at least avg latency
print_once('Not enough samples to measure latencies.')
def main():
"""
Launches text to speech (inference).
Inference is executed on a single GPU.
"""
parser = argparse.ArgumentParser(description='PyTorch FastPitch Inference',
allow_abbrev=False)
parser = parse_args(parser)
args, unk_args = parser.parse_known_args()
if args.p_arpabet > 0.0:
cmudict.initialize(args.cmudict_path, keep_ambiguous=True)
torch.backends.cudnn.benchmark = args.cudnn_benchmark
if args.output is not None:
Path(args.output).mkdir(parents=False, exist_ok=True)
log_fpath = args.log_file or str(Path(args.output, 'nvlog_infer.json'))
log_fpath = unique_log_fpath(log_fpath)
DLLogger.init(backends=[JSONStreamBackend(Verbosity.DEFAULT, log_fpath),
StdOutBackend(Verbosity.VERBOSE,
metric_format=stdout_metric_format)])
init_inference_metadata()
[DLLogger.log("PARAMETER", {k: v}) for k, v in vars(args).items()]
device = torch.device('cuda' if args.cuda else 'cpu')
if args.fastpitch != 'SKIP':
generator = load_and_setup_model(
'FastPitch', parser, args.fastpitch, args.amp, device,
unk_args=unk_args, forward_is_infer=True, ema=args.ema,
jitable=args.torchscript)
if args.torchscript:
generator = torch.jit.script(generator)
else:
generator = None
if args.waveglow != 'SKIP':
with warnings.catch_warnings():
warnings.simplefilter("ignore")
waveglow = load_and_setup_model(
'WaveGlow', parser, args.waveglow, args.amp, device,
unk_args=unk_args, forward_is_infer=True, ema=args.ema)
denoiser = Denoiser(waveglow).to(device)
waveglow = getattr(waveglow, 'infer', waveglow)
else:
waveglow = None
if len(unk_args) > 0:
raise ValueError(f'Invalid options {unk_args}')
fields = load_fields(args.input)
batches = prepare_input_sequence(
fields, device, args.symbol_set, args.text_cleaners, args.batch_size,
args.dataset_path, load_mels=(generator is None), p_arpabet=args.p_arpabet)
# Use real data rather than synthetic - FastPitch predicts len
for _ in tqdm(range(args.warmup_steps), 'Warmup'):
with torch.no_grad():
if generator is not None:
b = batches[0]
mel, *_ = generator(b['text'])
if waveglow is not None:
audios = waveglow(mel, sigma=args.sigma_infer).float()
_ = denoiser(audios, strength=args.denoising_strength)
gen_measures = MeasureTime(cuda=args.cuda)
waveglow_measures = MeasureTime(cuda=args.cuda)
gen_kw = {'pace': args.pace,
'speaker': args.speaker,
'pitch_tgt': None,
'pitch_transform': build_pitch_transformation(args)}
if args.torchscript:
gen_kw.pop('pitch_transform')
print('NOTE: Pitch transforms are disabled with TorchScript')
all_utterances = 0
all_samples = 0
all_letters = 0
all_frames = 0
reps = args.repeats
log_enabled = reps == 1
log = lambda s, d: DLLogger.log(step=s, data=d) if log_enabled else None
for rep in (tqdm(range(reps), 'Inference') if reps > 1 else range(reps)):
for b in batches:
if generator is None:
log(rep, {'Synthesizing from ground truth mels'})
mel, mel_lens = b['mel'], b['mel_lens']
else:
with torch.no_grad(), gen_measures:
mel, mel_lens, *_ = generator(b['text'], **gen_kw)
gen_infer_perf = mel.size(0) * mel.size(2) / gen_measures[-1]
all_letters += b['text_lens'].sum().item()
all_frames += mel.size(0) * mel.size(2)
log(rep, {"fastpitch_frames/s": gen_infer_perf})
log(rep, {"fastpitch_latency": gen_measures[-1]})
if args.save_mels:
for i, mel_ in enumerate(mel):
m = mel_[:, :mel_lens[i].item()].permute(1, 0)
fname = b['output'][i] if 'output' in b else f'mel_{i}.npy'
mel_path = Path(args.output, Path(fname).stem + '.npy')
np.save(mel_path, m.cpu().numpy())
if waveglow is not None:
with torch.no_grad(), waveglow_measures:
audios = waveglow(mel, sigma=args.sigma_infer)
audios = denoiser(audios.float(),
strength=args.denoising_strength
).squeeze(1)
all_utterances += len(audios)
all_samples += sum(audio.size(0) for audio in audios)
waveglow_infer_perf = (
audios.size(0) * audios.size(1) / waveglow_measures[-1])
log(rep, {"waveglow_samples/s": waveglow_infer_perf})
log(rep, {"waveglow_latency": waveglow_measures[-1]})
if args.output is not None and reps == 1:
for i, audio in enumerate(audios):
audio = audio[:mel_lens[i].item() * args.stft_hop_length]
if args.fade_out:
fade_len = args.fade_out * args.stft_hop_length
fade_w = torch.linspace(1.0, 0.0, fade_len)
audio[-fade_len:] *= fade_w.to(audio.device)
audio = audio / torch.max(torch.abs(audio))
fname = b['output'][i] if 'output' in b else f'audio_{i}.wav'
audio_path = Path(args.output, fname)
write(audio_path, args.sampling_rate, audio.cpu().numpy())
if generator is not None and waveglow is not None:
log(rep, {"latency": (gen_measures[-1] + waveglow_measures[-1])})
log_enabled = True
if generator is not None:
gm = np.sort(np.asarray(gen_measures))
rtf = all_samples / (all_utterances * gm.mean() * args.sampling_rate)
log((), {"avg_fastpitch_letters/s": all_letters / gm.sum()})
log((), {"avg_fastpitch_frames/s": all_frames / gm.sum()})
log((), {"avg_fastpitch_latency": gm.mean()})
log((), {"avg_fastpitch_RTF": rtf})
log((), {"90%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.90) / 2) * gm.std()})
log((), {"95%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.95) / 2) * gm.std()})
log((), {"99%_fastpitch_latency": gm.mean() + norm.ppf((1.0 + 0.99) / 2) * gm.std()})
if waveglow is not None:
wm = np.sort(np.asarray(waveglow_measures))
rtf = all_samples / (all_utterances * wm.mean() * args.sampling_rate)
log((), {"avg_waveglow_samples/s": all_samples / wm.sum()})
log((), {"avg_waveglow_latency": wm.mean()})
log((), {"avg_waveglow_RTF": rtf})
log((), {"90%_waveglow_latency": wm.mean() + norm.ppf((1.0 + 0.90) / 2) * wm.std()})
log((), {"95%_waveglow_latency": wm.mean() + norm.ppf((1.0 + 0.95) / 2) * wm.std()})
log((), {"99%_waveglow_latency": wm.mean() + norm.ppf((1.0 + 0.99) / 2) * wm.std()})
if generator is not None and waveglow is not None:
m = gm + wm
rtf = all_samples / (all_utterances * m.mean() * args.sampling_rate)
log((), {"avg_samples/s": all_samples / m.sum()})
log((), {"avg_letters/s": all_letters / m.sum()})
log((), {"avg_latency": m.mean()})
log((), {"avg_RTF": rtf})
log((), {"90%_latency": m.mean() + norm.ppf((1.0 + 0.90) / 2) * m.std()})
log((), {"95%_latency": m.mean() + norm.ppf((1.0 + 0.95) / 2) * m.std()})
log((), {"99%_latency": m.mean() + norm.ppf((1.0 + 0.99) / 2) * m.std()})
DLLogger.flush()
def main():
parser = get_parser()
args = parser.parse_args()
log_fpath = args.log_file or str(Path(args.output_dir, 'nvlog_infer.json'))
log_fpath = unique_log_fpath(log_fpath)
dllogger.init(backends=[JSONStreamBackend(Verbosity.DEFAULT, log_fpath),
StdOutBackend(Verbosity.VERBOSE,
metric_format=stdout_metric_format)])
[dllogger.log("PARAMETER", {k: v}) for k, v in vars(args).items()]
for step in ['DNN', 'data+DNN', 'data']:
for c in [0.99, 0.95, 0.9, 0.5]:
cs = 'avg' if c == 0.5 else f'{int(100*c)}%'
dllogger.metadata(f'{step.lower()}_latency_{c}',
{'name': f'{step} latency {cs}',
'format': ':>7.2f', 'unit': 'ms'})
dllogger.metadata(
'eval_wer', {'name': 'WER', 'format': ':>3.2f', 'unit': '%'})
if args.cpu:
device = torch.device('cpu')
else:
assert torch.cuda.is_available()
device = torch.device('cuda')
torch.backends.cudnn.benchmark = args.cudnn_benchmark
if args.seed is not None:
torch.manual_seed(args.seed + args.local_rank)
np.random.seed(args.seed + args.local_rank)
random.seed(args.seed + args.local_rank)
# set up distributed training
multi_gpu = not args.cpu and int(os.environ.get('WORLD_SIZE', 1)) > 1
if multi_gpu:
torch.cuda.set_device(args.local_rank)
distrib.init_process_group(backend='nccl', init_method='env://')
print_once(f'Inference with {distrib.get_world_size()} GPUs')
cfg = config.load(args.model_config)
config.apply_config_overrides(cfg, args)
symbols = helpers.add_ctc_blank(cfg['labels'])
use_dali = args.dali_device in ('cpu', 'gpu')
dataset_kw, features_kw = config.input(cfg, 'val')
measure_perf = args.steps > 0
# dataset
if args.transcribe_wav or args.transcribe_filelist:
if use_dali:
print("DALI supported only with input .json files; disabling")
use_dali = False
assert not args.pad_to_max_duration
assert not (args.transcribe_wav and args.transcribe_filelist)
if args.transcribe_wav:
dataset = SingleAudioDataset(args.transcribe_wav)
else:
dataset = FilelistDataset(args.transcribe_filelist)
data_loader = get_data_loader(dataset,
batch_size=1,
multi_gpu=multi_gpu,
shuffle=False,
num_workers=0,
drop_last=(True if measure_perf else False))
_, features_kw = config.input(cfg, 'val')
feat_proc = FilterbankFeatures(**features_kw)
elif use_dali:
# pad_to_max_duration is not supported by DALI - have simple padders
if features_kw['pad_to_max_duration']:
feat_proc = BaseFeatures(
pad_align=features_kw['pad_align'],
pad_to_max_duration=True,
max_duration=features_kw['max_duration'],
sample_rate=features_kw['sample_rate'],
window_size=features_kw['window_size'],
window_stride=features_kw['window_stride'])
features_kw['pad_to_max_duration'] = False
else:
feat_proc = None
data_loader = DaliDataLoader(
gpu_id=args.local_rank or 0,
dataset_path=args.dataset_dir,
config_data=dataset_kw,
config_features=features_kw,
json_names=args.val_manifests,
batch_size=args.batch_size,
pipeline_type=("train" if measure_perf else "val"), # no drop_last
device_type=args.dali_device,
symbols=symbols)
else:
dataset = AudioDataset(args.dataset_dir,
args.val_manifests,
symbols,
**dataset_kw)
data_loader = get_data_loader(dataset,
args.batch_size,
multi_gpu=multi_gpu,
shuffle=False,
num_workers=4,
drop_last=False)
feat_proc = FilterbankFeatures(**features_kw)
model = QuartzNet(encoder_kw=config.encoder(cfg),
decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
if args.ckpt is not None:
print(f'Loading the model from {args.ckpt} ...')
checkpoint = torch.load(args.ckpt, map_location="cpu")
key = 'ema_state_dict' if args.ema else 'state_dict'
state_dict = checkpoint[key]
model.load_state_dict(state_dict, strict=True)
model.to(device)
model.eval()
if feat_proc is not None:
feat_proc.to(device)
feat_proc.eval()
if args.amp:
model = model.half()
if args.torchscript:
greedy_decoder = GreedyCTCDecoder()
feat_proc, model, greedy_decoder = torchscript_export(
data_loader, feat_proc, model, greedy_decoder, args.output_dir,
use_amp=args.amp, use_conv_masks=True, model_toml=args.model_toml,
device=device, save=args.torchscript_export)
if multi_gpu:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
agg = {'txts': [], 'preds': [], 'logits': []}
dur = {'data': [], 'dnn': [], 'data+dnn': []}
looped_loader = chain.from_iterable(repeat(data_loader))
greedy_decoder = GreedyCTCDecoder()
sync = lambda: torch.cuda.synchronize() if device.type == 'cuda' else None
steps = args.steps + args.warmup_steps or len(data_loader)
with torch.no_grad():
for it, batch in enumerate(tqdm(looped_loader, initial=1, total=steps)):
if use_dali:
feats, feat_lens, txt, txt_lens = batch
if feat_proc is not None:
feats, feat_lens = feat_proc(feats, feat_lens)
else:
batch = [t.to(device, non_blocking=True) for t in batch]
audio, audio_lens, txt, txt_lens = batch
feats, feat_lens = feat_proc(audio, audio_lens)
sync()
t1 = time.perf_counter()
if args.amp:
feats = feats.half()
if model.encoder.use_conv_masks:
log_probs, log_prob_lens = model(feats, feat_lens)
else:
log_probs = model(feats, feat_lens)
preds = greedy_decoder(log_probs)
sync()
t2 = time.perf_counter()
# burn-in period; wait for a new loader due to num_workers
if it >= 1 and (args.steps == 0 or it >= args.warmup_steps):
dur['data'].append(t1 - t0)
dur['dnn'].append(t2 - t1)
dur['data+dnn'].append(t2 - t0)
if txt is not None:
agg['txts'] += helpers.gather_transcripts([txt], [txt_lens],
symbols)
agg['preds'] += helpers.gather_predictions([preds], symbols)
agg['logits'].append(log_probs)
if it + 1 == steps:
break
sync()
t0 = time.perf_counter()
# communicate the results
if args.transcribe_wav:
for idx, p in enumerate(agg['preds']):
print_once(f'Prediction {idx+1: >3}: {p}')
elif args.transcribe_filelist:
pass
elif not multi_gpu or distrib.get_rank() == 0:
wer, _ = process_evaluation_epoch(agg)
dllogger.log(step=(), data={'eval_wer': 100 * wer})
if args.save_predictions:
with open(args.save_predictions, 'w') as f:
f.write('\n'.join(agg['preds']))
if args.save_logits:
logits = torch.cat(agg['logits'], dim=0).cpu()
torch.save(logits, args.save_logits)
# report timings
if len(dur['data']) >= 20:
ratios = [0.9, 0.95, 0.99]
for stage in dur:
lat = durs_to_percentiles(dur[stage], ratios)
for k in [0.99, 0.95, 0.9, 0.5]:
kk = str(k).replace('.', '_')
dllogger.log(step=(), data={f'{stage.lower()}_latency_{kk}': lat[k]})
else:
print_once('Not enough samples to measure latencies.')
