Skip to main content

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

Project description

Documentation Status

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 want to make efficient in the short term 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

A lot of the code is still unfinished (Sep 11, 2020). We finished the CPU versions of many algorithms and this code is in k2/csrc/host/; however, after that we figured out how to implement things on the GPU and decided to change the interfaces so the CPU and GPU code had a more unified interface. Currently in k2/csrc/ we have more GPU-oriented implementations (although these algorithms will also work on CPU). We had almost finished the Python wrapping for the older code, in the k2/python/ subdirectory, but we decided not to release code with that wrapping because it would have had to be reworked to be compatible with our GPU algorithms. Instead we will use the interfaces drafted in k2/csrc/ e.g. the Context object (which encapsulates things like memory managers from external toolkits) and the Tensor object which can be used to wrap tensors from external toolkits; and wrap those in Python (using pybind11). The code in host/ will eventually be either deprecated, rewritten or wrapped with newer-style interfaces.

Plans for initial release

We hope to get the first version working in early October. The current short-term aim is to finish the GPU implementation of pruned composition of a normal FSA with a dense FSA, which is the same as decoder search in speech recognition and can be used to implement CTC training and lattice-free MMI (LF-MMI) training. The proof-of-concept that we will release initially is something that's like CTC but allowing more general supervisions (general FSAs rather than linear sequences). This will work on GPU. The same underlying code will support LF-MMI so that would be easy to implement soon after. We plan to put example code in a separate repository.

Plans after initial release

We will then gradually implement more algorithms in a way that's compatible with the interfaces in k2/csrc/. Some of them will be CPU-only to start with. The idea is to eventually have very rich capabilities for operating on collections of sequences, including methods to convert from a lattice to a collection of linear sequences and back again (for purposes of neural language model rescoring, neural confidence estimation and the like).

Quick start

Want to try it out without installing anything? We have setup a Google Colab.

Caution: k2 is not nearly ready for actual use! We are still coding the core algorithms, and hope to have an early version working by early October.

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.0.dev20210617-py38-none-any.whl (55.4 MB view details)

Uploaded Python 3.8

k2-1.0.dev20210617-py37-none-any.whl (55.4 MB view details)

Uploaded Python 3.7

k2-1.0.dev20210617-py36-none-any.whl (55.4 MB view details)

Uploaded Python 3.6

k2-1.0.dev20210617-cp38-cp38-macosx_10_15_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.8 macOS 10.15+ x86-64

k2-1.0.dev20210617-cp37-cp37m-macosx_10_15_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.7m macOS 10.15+ x86-64

k2-1.0.dev20210617-cp36-cp36m-macosx_10_15_x86_64.whl (1.4 MB view details)

Uploaded CPython 3.6m macOS 10.15+ x86-64

File details

Details for the file k2-1.0.dev20210617-py38-none-any.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-py38-none-any.whl
  • Upload date:
  • Size: 55.4 MB
  • Tags: Python 3.8
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.8.10

File hashes

Hashes for k2-1.0.dev20210617-py38-none-any.whl
Algorithm Hash digest
SHA256 d4f2fb7afaae0d70ed6c928348ed3da925462bc4fc01b971e9cecef6f191baca
MD5 7f97d586f39feebaf069a70a8e0ef7b3
BLAKE2b-256 27be9641d2ba7539bea68eadddee82ebf10c55201638bb4347c551d8768bd97b

See more details on using hashes here.

File details

Details for the file k2-1.0.dev20210617-py37-none-any.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-py37-none-any.whl
  • Upload date:
  • Size: 55.4 MB
  • Tags: Python 3.7
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.7.10

File hashes

Hashes for k2-1.0.dev20210617-py37-none-any.whl
Algorithm Hash digest
SHA256 a85e024d2fc67f03b4da61e8545954c17c7b8c6ee6ba586573ac67a4d3b99379
MD5 bc055fe53380d4a0ec8dcb7f42af0302
BLAKE2b-256 41c4d92d2ed680c3459f7807eb77cbcc0d02e2ad5d783311fdd487a8ff44d37e

See more details on using hashes here.

File details

Details for the file k2-1.0.dev20210617-py36-none-any.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-py36-none-any.whl
  • Upload date:
  • Size: 55.4 MB
  • Tags: Python 3.6
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.6.13

File hashes

Hashes for k2-1.0.dev20210617-py36-none-any.whl
Algorithm Hash digest
SHA256 732223450c06991961d90dd32897c1cd4ca60c9c032cfd189eaaab910b589997
MD5 d05a4a1ad4be75108b5cf24a515bd3d6
BLAKE2b-256 4467fa9a7a88c9022854f47a40064d47df121e8a24625462d0ca77d477ad2e61

See more details on using hashes here.

File details

Details for the file k2-1.0.dev20210617-cp38-cp38-macosx_10_15_x86_64.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-cp38-cp38-macosx_10_15_x86_64.whl
  • Upload date:
  • Size: 1.4 MB
  • Tags: CPython 3.8, macOS 10.15+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.8.10

File hashes

Hashes for k2-1.0.dev20210617-cp38-cp38-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 b772357d61f1fcfc27d68ad03bb6b4376a32981a64643fcdc347f65c92f2dbf6
MD5 70e38ba1f1cc8aab0dc6f1a56f0bdcaa
BLAKE2b-256 9bba10479e30b25f27a7fcf348d7951dc0db9da3b938c60097adc04cf38f5179

See more details on using hashes here.

File details

Details for the file k2-1.0.dev20210617-cp37-cp37m-macosx_10_15_x86_64.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-cp37-cp37m-macosx_10_15_x86_64.whl
  • Upload date:
  • Size: 1.4 MB
  • Tags: CPython 3.7m, macOS 10.15+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.7.10

File hashes

Hashes for k2-1.0.dev20210617-cp37-cp37m-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 0bd09135ff532ca7bc88a1a7e73100f9b2c092898cbf2e6bb7a19e8c6d26a46c
MD5 1cd39e1669e56c8af71d796c4a9b37c9
BLAKE2b-256 c20ea513a5a3c88aba485488584583497f9c82180b68b4b745d98168d10cdf55

See more details on using hashes here.

File details

Details for the file k2-1.0.dev20210617-cp36-cp36m-macosx_10_15_x86_64.whl.

File metadata

  • Download URL: k2-1.0.dev20210617-cp36-cp36m-macosx_10_15_x86_64.whl
  • Upload date:
  • Size: 1.4 MB
  • Tags: CPython 3.6m, macOS 10.15+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.5.0 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.61.1 CPython/3.6.13

File hashes

Hashes for k2-1.0.dev20210617-cp36-cp36m-macosx_10_15_x86_64.whl
Algorithm Hash digest
SHA256 50475c7f879a502a2bd1e2fb8c6ee1bbb064dd440f668139a005f1144d62b259
MD5 fe45079989717c8dce9f7fb3be9e36c4
BLAKE2b-256 cd80d1c61c2ac11c6f1e512952ea038ed25c18ba3b87166dd5db2e732575cce7

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