infer.py 5.2 KB
Newer Older
1
"""Inferer for DeepSpeech2 model."""
2 3 4
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
X
Xinghai Sun 已提交
5

6
import argparse
7
import distutils.util
8
import multiprocessing
9 10
import paddle.v2 as paddle
from data_utils.data import DataGenerator
11
from model import DeepSpeech2Model
12
from error_rate import wer
13
import utils
14

15
parser = argparse.ArgumentParser(description=__doc__)
16
parser.add_argument(
X
Xinghai Sun 已提交
17
    "--num_samples",
Y
Yibing Liu 已提交
18
    default=10,
X
Xinghai Sun 已提交
19
    type=int,
20
    help="Number of samples for inference. (default: %(default)s)")
21
parser.add_argument(
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40
    "--num_conv_layers",
    default=2,
    type=int,
    help="Convolution layer number. (default: %(default)s)")
parser.add_argument(
    "--num_rnn_layers",
    default=3,
    type=int,
    help="RNN layer number. (default: %(default)s)")
parser.add_argument(
    "--rnn_layer_size",
    default=512,
    type=int,
    help="RNN layer cell number. (default: %(default)s)")
parser.add_argument(
    "--use_gpu",
    default=True,
    type=distutils.util.strtobool,
    help="Use gpu or not. (default: %(default)s)")
41 42
parser.add_argument(
    "--num_threads_data",
43
    default=multiprocessing.cpu_count(),
44 45
    type=int,
    help="Number of cpu threads for preprocessing data. (default: %(default)s)")
Y
Yibing Liu 已提交
46 47 48 49 50
parser.add_argument(
    "--num_processes_beam_search",
    default=multiprocessing.cpu_count(),
    type=int,
    help="Number of cpu processes for beam search. (default: %(default)s)")
51 52 53 54 55 56
parser.add_argument(
    "--specgram_type",
    default='linear',
    type=str,
    help="Feature type of audio data: 'linear' (power spectrum)"
    " or 'mfcc'. (default: %(default)s)")
57 58 59 60 61
parser.add_argument(
    "--trainer_count",
    default=8,
    type=int,
    help="Trainer number. (default: %(default)s)")
62
parser.add_argument(
63 64
    "--mean_std_filepath",
    default='mean_std.npz',
65 66
    type=str,
    help="Manifest path for normalizer. (default: %(default)s)")
67
parser.add_argument(
68
    "--decode_manifest_path",
Y
Yibing Liu 已提交
69
    default='datasets/manifest.test',
70 71
    type=str,
    help="Manifest path for decoding. (default: %(default)s)")
72
parser.add_argument(
73
    "--model_filepath",
Y
Yibing Liu 已提交
74
    default='checkpoints/params.latest.tar.gz',
75 76
    type=str,
    help="Model filepath. (default: %(default)s)")
77 78
parser.add_argument(
    "--vocab_filepath",
79
    default='datasets/vocab/eng_vocab.txt',
80 81
    type=str,
    help="Vocabulary filepath. (default: %(default)s)")
Y
Yibing Liu 已提交
82 83
parser.add_argument(
    "--decode_method",
Y
Yibing Liu 已提交
84
    default='beam_search',
Y
Yibing Liu 已提交
85
    type=str,
Y
Yibing Liu 已提交
86 87
    help="Method for ctc decoding: best_path or beam_search. (default: %(default)s)"
)
Y
Yibing Liu 已提交
88 89
parser.add_argument(
    "--beam_size",
90
    default=500,
Y
Yibing Liu 已提交
91 92 93
    type=int,
    help="Width for beam search decoding. (default: %(default)d)")
parser.add_argument(
Y
Yibing Liu 已提交
94 95
    "--num_results_per_sample",
    default=1,
Y
Yibing Liu 已提交
96
    type=int,
Y
Yibing Liu 已提交
97 98 99
    help="Number of output per sample in beam search. (default: %(default)d)")
parser.add_argument(
    "--language_model_path",
Y
Yibing Liu 已提交
100
    default="lm/data/common_crawl_00.prune01111.trie.klm",
Y
Yibing Liu 已提交
101
    type=str,
Y
Yibing Liu 已提交
102
    help="Path for language model. (default: %(default)s)")
Y
Yibing Liu 已提交
103 104
parser.add_argument(
    "--alpha",
105
    default=0.26,
Y
Yibing Liu 已提交
106 107 108 109
    type=float,
    help="Parameter associated with language model. (default: %(default)f)")
parser.add_argument(
    "--beta",
110
    default=0.1,
Y
Yibing Liu 已提交
111 112
    type=float,
    help="Parameter associated with word count. (default: %(default)f)")
113 114 115 116 117 118
parser.add_argument(
    "--cutoff_prob",
    default=0.99,
    type=float,
    help="The cutoff probability of pruning"
    "in beam search. (default: %(default)f)")
119 120 121
args = parser.parse_args()


122
def infer():
Y
Yibing Liu 已提交
123
    """Inference for DeepSpeech2."""
124
    data_generator = DataGenerator(
125
        vocab_filepath=args.vocab_filepath,
126
        mean_std_filepath=args.mean_std_filepath,
127
        augmentation_config='{}',
128
        specgram_type=args.specgram_type,
129
        num_threads=args.num_threads_data)
130
    batch_reader = data_generator.batch_reader_creator(
131 132
        manifest_path=args.decode_manifest_path,
        batch_size=args.num_samples,
Y
Yibing Liu 已提交
133
        min_batch_size=1,
134
        sortagrad=False,
135
        shuffle_method=None)
136
    infer_data = batch_reader().next()
137

138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
    ds2_model = DeepSpeech2Model(
        vocab_size=data_generator.vocab_size,
        num_conv_layers=args.num_conv_layers,
        num_rnn_layers=args.num_rnn_layers,
        rnn_layer_size=args.rnn_layer_size,
        pretrained_model_path=args.model_filepath)
    result_transcripts = ds2_model.infer_batch(
        infer_data=infer_data,
        decode_method=args.decode_method,
        beam_alpha=args.alpha,
        beam_beta=args.beta,
        beam_size=args.beam_size,
        cutoff_prob=args.cutoff_prob,
        vocab_list=data_generator.vocab_list,
        language_model_path=args.language_model_path,
        num_processes=args.num_processes_beam_search)
154

155 156 157
    target_transcripts = [
        ''.join([data_generator.vocab_list[token] for token in transcript])
        for _, transcript in infer_data
Y
Yibing Liu 已提交
158
    ]
159 160 161 162
    for target, result in zip(target_transcripts, result_transcripts):
        print("\nTarget Transcription: %s\nOutput Transcription: %s" %
              (target, result))
        print("Current wer = %f" % wer(target, result))
163 164 165


def main():
166
    utils.print_arguments(args)
167
    paddle.init(use_gpu=args.use_gpu, trainer_count=args.trainer_count)
168
    infer()
169 170 171 172


if __name__ == '__main__':
    main()