PyTorch-CTC is an implementation of CTC (Connectionist Temporal Classification) beam search decoding for PyTorch. C++ code borrowed liberally from TensorFlow with some improvements to increase flexibility.
The library is largely self-contained and requires only PyTorch and CFFI. Building the C++ library requires gcc. KenLM language modeling support is also optionally included, and enabled by default.
# get the code
git clone --recursive https://github.com/ryanleary/pytorch-ctc.git
cd pytorch-ctc
# install dependencies (PyTorch and CFFI)
pip install -r requirements.txt
python setup.py install
# If you do NOT require kenlm, the `--recursive` flag is not required on git clone
# and `--exclude-kenlm` should be appended to the `python setup.py install` command
pytorch-ctc includes a CTC beam search decoder with multiple scorer implementations. A scorer
is a function that the decoder calls to condition the probability of a given beam based on its state.
Two Scorer implementations are currently implemented for pytorch-ctc.
Scorer: is a NO-OP and enables the decoder to do a vanilla beam decode
scorer = Scorer()
KenLMScorer: conditions beams based on the provided KenLM binary language model.
scorer = KenLMScorer(labels, lm_path, trie_path, blank_index=0, space_index=28)
where:
labels
is a string of output labels given in the same order as the output layerlm_path
path to a binary KenLM language model for decodingtrie_path
path to a Trie containing the lexicon (see generate_lm_trie)blank_index
is used to specify which position in the output distribution represents theblank
classspace_index
is used to specify which position in the output distribution represents the word separator class
The KenLMScorer
may be further configured with weights for the language model contribution to the score (lm_weight
), as well as word and valid word bonuses (to offset decreasing probability as a function of sequence length).
scorer.set_lm_weight(2.1)
scorer.set_word_weight(1.1)
scorer.set_valid_word_weight(1.5)
decoder = CTCBeamDecoder(scorer, labels, top_paths=3, beam_width=20,
blank_index=0, space_index=28, merge_repeated=False)
where:
scorer
is an instance of a concrete implementation of theBaseScorer
classlabels
is a string of output labels given in the same order as the output layertop_paths
is used to specify how many hypotheses to returnbeam_width
is the number of beams to evaluate in a given stepblank_index
is used to specify which position in the output distribution represents theblank
classspace_index
is used to specify which position in the output distribution represents the word separator classmerge_repeated
if True will collapse repeated characters
output, score, out_seq_len = decoder.decode(probs, sizes=None)
where:
probs
is a FloatTensor of log-probabilities with shape(seq_len, batch_size, num_classes)
seq_len
is an optional IntTensor of integer sequence lengths with shape(batch_size)
and returns:
output
is an IntTensor of character classes of shape(top_paths, batch_size, seq_len)
score
is a FloatTensor of log-probabilities representing the likelihood of the transcription with shape(top_paths, batch_size)
out_seq_len
is an IntTensor containing the length of the output sequence with shape(top_paths, batch_size)
The CTCBeamDecoder
may be further configured with weights for the label size (label_size
), and label margin ('label_margin'). These parameters helps to reduce
the computation time.
Label selection size controls how many items in each beam are passed through to the beam scorer. Only items with top N input scores are considered. Label selection margin controls the difference between minimal input score (versus the best scoring label) for an item to be passed to the beam scorer. This margin is expressed in terms of log-probability. Default is to do no label selection.
decoder.set_label_selection_parameters(label_size=0, label_margin=6)
generate_lm_dict(dictionary_path, kenlm_path, output_path, labels, blank_index, space_index)
A vocabulary trie is required for the KenLM Scorer. The trie is created from a lexicon specified as a newline separated text file of words in the vocabulary.
Thanks to ebrevdo for the original TensorFlow CTC decoder implementation, timediv for his KenLM extension, and SeanNaren for his assistance.