Skip to main content

FSA/FST algorithms, intended to (eventually) be interoperable with PyTorch and similar

Project description

k2

The vision of k2 is to be able to seamlessly integrate Finite State Automaton (FSA) and Finite State Transducer (FST) algorithms into autograd-based machine learning toolkits like PyTorch and TensorFlow. For speech recognition applications, this should make it easy to interpolate and combine various training objectives such as cross-entropy, CTC and MMI and to jointly optimize a speech recognition system with multiple decoding passes including lattice rescoring and confidence estimation. We hope k2 will have many other applications as well.

One of the key algorithms that we have implemented is pruned composition of a generic FSA with a "dense" FSA (i.e. one that corresponds to log-probs of symbols at the output of a neural network). This can be used as a fast implementation of decoding for ASR, and for CTC and LF-MMI training. This won't give a direct advantage in terms of Word Error Rate when compared with existing technology; but the point is to do this in a much more general and extensible framework to allow further development of ASR technology.

Implementation

A few key points on our implementation strategy.

Most of the code is in C++ and CUDA. We implement a templated class Ragged, which is quite like TensorFlow's RaggedTensor (actually we came up with the design independently, and were later told that TensorFlow was using the same ideas). Despite a close similarity at the level of data structures, the design is quite different from TensorFlow and PyTorch. Most of the time we don't use composition of simple operations, but rely on C++11 lambdas defined directly in the C++ implementations of algorithms. The code in these lambdas operate directly on data pointers and, if the backend is CUDA, they can run in parallel for each element of a tensor. (The C++ and CUDA code is mixed together and the CUDA kernels get instantiated via templates).

It is difficult to adequately describe what we are doing with these Ragged objects without going in detail through the code. The algorithms look very different from the way you would code them on CPU because of the need to avoid sequential processing. We are using coding patterns that make the most expensive parts of the computations "embarrassingly parallelizable"; the only somewhat nontrivial CUDA operations are generally reduction-type operations such as exclusive-prefix-sum, for which we use NVidia's cub library. Our design is not too specific to the NVidia hardware and the bulk of the code we write is fairly normal-looking C++; the nontrivial CUDA programming is mostly done via the cub library, parts of which we wrap with our own convenient interface.

The Finite State Automaton object is then implemented as a Ragged tensor templated on a specific data type (a struct representing an arc in the automaton).

Autograd

If you look at the code as it exists now, you won't find any references to autograd. The design is quite different to TensorFlow and PyTorch (which is why we didn't simply extend one of those toolkits). Instead of making autograd come from the bottom up (by making individual operations differentiable) we are implementing it from the top down, which is much more efficient in this case (and will tend to have better roundoff properties).

An example: suppose we are finding the best path of an FSA, and we need derivatives. We implement this by keeping track of, for each arc in the output best-path, which input arc it corresponds to. (For more complex algorithms an arc in the output might correspond to a sum of probabilities of a list of input arcs). We can make this compatible with PyTorch/TensorFlow autograd at the Python level, by, for example, defining a Function class in PyTorch that remembers this relationship between the arcs and does the appropriate (sparse) operations to propagate back the derivatives w.r.t. the weights.

Current state of the code

We have wrapped all the C++ code to Python with pybind11 and have finished the integration with PyTorch.

We are currently writing speech recognition recipes using k2, which are hosted in a separate repository. Please see https://github.com/k2-fsa/icefall.

Plans after initial release

We are currently trying to make k2 ready for production use (see the branch v2.0-pre).

Quick start

Want to try it out without installing anything? We have setup a Google Colab. You can find more Colab notebooks using k2 in speech recognition at https://icefall.readthedocs.io/en/latest/recipes/librispeech/conformer_ctc.html.

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

k2-1.24.1-cp310-cp310-macosx_10_15_x86_64.whl (2.4 MB view details)

Uploaded CPython 3.10 macOS 10.15+ x86-64

k2-1.24.1-cp39-cp39-macosx_10_15_x86_64.whl (2.4 MB view details)

Uploaded CPython 3.9 macOS 10.15+ x86-64

k2-1.24.1-cp38-cp38-macosx_10_15_x86_64.whl (2.4 MB view details)

Uploaded CPython 3.8 macOS 10.15+ x86-64

k2-1.24.1-cp37-cp37m-macosx_10_15_x86_64.whl (2.4 MB view details)

Uploaded CPython 3.7m macOS 10.15+ x86-64

File details

Details for the file k2-1.24.1-cp310-cp310-macosx_10_15_x86_64.whl.

File metadata

File hashes

Hashes for k2-1.24.1-cp310-cp310-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 d17d0c7115eb500f51025cd95048ae15b75d12a603776e0ad2d588622ac09bc6
MD5 bc3eeb75b9381121cdb297385c6a408c
BLAKE2b-256 42b59fc13bcab6fa65ee3e9faf6b9d033518838a3daf2a8ba9280f230c49bba8

See more details on using hashes here.

File details

Details for the file k2-1.24.1-cp39-cp39-macosx_10_15_x86_64.whl.

File metadata

File hashes

Hashes for k2-1.24.1-cp39-cp39-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 b2c5bed70032d731ccdf48b8b7056febcb70311866552305a762a6f96292b139
MD5 22b7d8f01e8c38ed6c6f4e92a648458a
BLAKE2b-256 389f879e1f47428b829526d3bd0afc9e6cfc461044dadc7b0a549d503b8cd3d1

See more details on using hashes here.

File details

Details for the file k2-1.24.1-cp38-cp38-macosx_10_15_x86_64.whl.

File metadata

File hashes

Hashes for k2-1.24.1-cp38-cp38-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 36e76e57f535fef4e2371844bc7c9eae533efa9cc05dc1fae8f6bc265993c258
MD5 beb1354087f34d772f5f4fe75453385a
BLAKE2b-256 642126f06140a67404b9aad6d0f9349db659c7deec4744774ad67ab7a0ff76ec

See more details on using hashes here.

File details

Details for the file k2-1.24.1-cp37-cp37m-macosx_10_15_x86_64.whl.

File metadata

File hashes

Hashes for k2-1.24.1-cp37-cp37m-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 a182c34c8b2973bd5042d9de9a14d5364ba8d6d6a82875602270ff76635de35b
MD5 74eb625486b1803534ea6e3c6e0b278e
BLAKE2b-256 c291d7e0c2f2411f7566dd0c919bc857d5e88fde33750bd379a53f908f2670c5

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page