gaoyiyeah / end2end Goto Github PK

View Code? Open in Web Editor NEW

This project forked from artbataev/end2end

0.0 2.0 0.0 5.86 MB

Losses and decoders for end-to-end ASR and OCR

Home Page: https://artbataev.github.io/end2end/

License: MIT License

CMake 1.78% Python 73.00% C++ 25.22%

end2end's Introduction

Losses and decoders for end-to-end Speech Recognition and Optical Character Recognition with PyTorch

The module focuses on experiments with CTC-loss (Connectionist Temporal Classification) and its modifications.

Under active development.

Documentation

artbataev.github.io/end2end/

Losses

CTC (C++, CPU)
CTC without blank (python / numba), needs refactoring
Dynamic segmentation with CTC (python), needs refactoring

Decoders

CTC Greedy Decoder (C++, CPU)
CTC Beam Search Decoder (C++, CPU)
CTC Beam Search Decoder with language model (C++, CPU)

How to install

Requirements:

Python 3.4+
PyTorch 1.0 and higher
numpy
numba

Install PyTorch 1.0 from pytorch.org, e.g.
```
pip install torch
```

Install tools to compile

sudo add-apt-repository ppa:ubuntu-toolchain-r-test
sudo apt-get install cmake libboost-all-dev

Install the module

pip install -v git+https://github.com/artbataev/end2end.git

git clone --recursive https://github.com/artbataev/end2end.git
cd end2end
python setup.py install
python -m tests.test_ctc
python -m tests.test_ctc_decoder

How to use

CTC Loss

import torch
from pytorch_end2end import CTCLoss

ctc_loss = CTCLoss(blank_idx=0, time_major=False, 
    reduce=True, size_average=True, after_logsoftmax=False)

batch_size = 4
alphabet_size = 28 # blank + 26 english characters + space

logits = torch.randn(batch_size, 50, alphabet_size).detach().requires_grad_()
targets = torch.randint(1, alphabet_size, (batch_size, 30), dtype=torch.long)
logits_lengths = torch.full((batch_size,), 50, dtype=torch.long)
targets_lengths = torch.randint(10, 30, (batch_size,), dtype=torch.long)

loss = ctc_loss(logits, targets, logits_lengths, targets_lengths)
loss.backward()

CTC Decoder

import torch
from pytorch_end2end import CTCDecoder

batch_size = 4
alphabet_size = 6
decoder = CTCDecoder(blank_idx=0, beam_width=100, time_major=False, 
    labels=["_", "a", "b", "c", "d", " "])

logits = torch.randn(batch_size, 50, alphabet_size).detach()
logits_lengths = torch.full((batch_size,), 50, dtype=torch.long)

decoded_targets, decoded_targets_lengths, decoded_sentences = decoder.decode(logits, logits_lengths)
for sentence in decoded_sentences:
    print(sentence)