forked from SeanNaren/deepspeech.pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
85 lines (71 loc) · 3.74 KB
/
test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import argparse
import json
import torch
from torch.autograd import Variable
from data.data_loader import SpectrogramDataset, AudioDataLoader
from decoder import ArgMaxDecoder
from model import DeepSpeech
parser = argparse.ArgumentParser(description='DeepSpeech prediction')
parser.add_argument('--sample_rate', default=16000, type=int, help='Sample rate')
parser.add_argument('--labels_path', default='labels.json', help='Contains all characters for prediction')
parser.add_argument('--model_path', default='models/deepspeech_final.pth.tar',
help='Path to model file created by training')
parser.add_argument('--audio_path', default='audio.wav',
help='Audio file to predict on')
parser.add_argument('--window_size', default=.02, type=float, help='Window size for spectrogram in seconds')
parser.add_argument('--window_stride', default=.01, type=float, help='Window stride for spectrogram in seconds')
parser.add_argument('--window', default='hamming', help='Window type for spectrogram generation')
parser.add_argument('--cuda', default=True, type=bool, help='Use cuda to train model')
parser.add_argument('--val_manifest', metavar='DIR',
help='path to validation manifest csv', default='data/val_manifest.csv')
parser.add_argument('--batch_size', default=20, type=int, help='Batch size for training')
parser.add_argument('--num_workers', default=4, type=int, help='Number of workers used in dataloading')
args = parser.parse_args()
if __name__ == '__main__':
package = torch.load(args.model_path)
model = DeepSpeech(rnn_hidden_size=package['hidden_size'], nb_layers=package['hidden_layers'],
num_classes=package['nout'])
if args.cuda:
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(package['state_dict'])
model.eval()
with open(args.labels_path) as label_file:
labels = str(''.join(json.load(label_file)))
decoder = ArgMaxDecoder(labels)
audio_conf = dict(sample_rate=args.sample_rate,
window_size=args.window_size,
window_stride=args.window_stride,
window=args.window)
test_dataset = SpectrogramDataset(audio_conf=audio_conf, manifest_filepath=args.val_manifest, labels=labels,
normalize=True)
test_loader = AudioDataLoader(test_dataset, batch_size=args.batch_size,
num_workers=args.num_workers)
total_cer, total_wer = 0, 0
for i, (data) in enumerate(test_loader):
inputs, targets, input_percentages, target_sizes = data
inputs = Variable(inputs)
# unflatten targets
split_targets = []
offset = 0
for size in target_sizes:
split_targets.append(targets[offset:offset + size])
offset += size
if args.cuda:
inputs = inputs.cuda()
out = model(inputs)
out = out.transpose(0, 1) # TxNxH
seq_length = out.size(0)
sizes = Variable(input_percentages.mul_(int(seq_length)).int())
decoded_output = decoder.decode(out.data, sizes)
target_strings = decoder.process_strings(decoder.convert_to_strings(split_targets))
wer, cer = 0, 0
for x in range(len(target_strings)):
wer += decoder.wer(decoded_output[x], target_strings[x]) / float(len(target_strings[x].split()))
cer += decoder.cer(decoded_output[x], target_strings[x]) / float(len(target_strings[x]))
total_cer += cer
total_wer += wer
wer = total_wer / len(test_loader.dataset)
cer = total_cer / len(test_loader.dataset)
print('Validation Summary \t'
'Average WER {wer:.0f}\t'
'Average CER {cer:.0f}\t'.format(wer=wer * 100, cer=cer * 100))