This project implements a method for faster and more memory-efficient RNN-T loss computation, called `pruned rnnt`.
## 安装
组件支持组合
Note: There is also a fast RNN-T loss implementation in [k2](https://github.com/k2-fsa/k2) project, which shares the same code here. We make `fast_rnnt` a stand-alone project in case someone wants only this rnnt loss.
We first obtain pruning bounds for the RNN-T recursion using a simple joiner network that is just an addition of the encoder and decoder, then we use those pruning bounds to evaluate the full, non-linear joiner network.
The picture below display the gradients (obtained by `rnnt_loss_simple` with `return_grad=true`) of lattice nodes, at each time frame, only a small set of nodes have a non-zero gradient, which justifies the pruned RNN-T loss, i.e., putting a limit on the number of symbols per frame.
Python 3.10.12 (main, May 27 2025, 17:12:29) [GCC 11.4.0] on linux
pip install --verbose fast_rnnt
Type "help", "copyright", "credits" or "license" for more information.
>>> import fast_rnnt
>>> fast_rnnt.__version__
'1.2'
>>>
```
```
版本号与官方版本同步，查询该软件的版本号，例如0.3；
It will pass `-DCMAKE_BUILD_TYPE=Release -DFT_WITH_CUDA=OFF` to `cmake`.
## Known Issue
- 无
### Where to get help if I have problems with the installation ?
Please file an issue at <https://github.com/danpovey/fast_rnnt/issues>
and describe your problem there.
## Usage
### For rnnt_loss_simple
This is a simple case of the RNN-T loss, where the joiner network is just
addition.
Note: termination_symbol plays the role of blank in other RNN-T loss implementations, we call it termination_symbol as it terminates symbols of current frame.
```python
am=torch.randn((B,T,C),dtype=torch.float32)
lm=torch.randn((B,S+1,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss_simple(
lm=lm,
am=am,
symbols=symbols,
termination_symbol=termination_symbol,
boundary=boundary,
reduction="sum",
)
```
### For rnnt_loss_smoothed
The same as `rnnt_loss_simple`, except that it supports `am_only` & `lm_only` smoothing
that allows you to make the loss-function one of the form:
lm_only_scale * lm_probs +
am_only_scale * am_probs +
(1-lm_only_scale-am_only_scale) * combined_probs
where `lm_probs` and `am_probs` are the probabilities given the lm and acoustic model independently.
```python
am=torch.randn((B,T,C),dtype=torch.float32)
lm=torch.randn((B,S+1,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss_smoothed(
lm=lm,
am=am,
symbols=symbols,
termination_symbol=termination_symbol,
lm_only_scale=0.25,
am_only_scale=0.0
boundary=boundary,
reduction="sum",
)
```
### For rnnt_loss_pruned
`rnnt_loss_pruned` can not be used alone, it needs the gradients returned by `rnnt_loss_simple/rnnt_loss_smoothed` to get pruning bounds.
You can also find recipes [here](https://github.com/k2-fsa/icefall/tree/master/egs/librispeech/ASR/pruned_transducer_stateless) that uses `rnnt_loss_pruned` to train a model.
### For rnnt_loss
The `unprund rnnt_loss` is the same as `torchaudio rnnt_loss`, it produces same output as torchaudio for the same input.
```python
logits=torch.randn((B,S,T,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss(
logits=logits,
symbols=symbols,
termination_symbol=termination_symbol,
boundary=boundary,
reduction="sum",
)
```
## Benchmarking
The [repo](https://github.com/csukuangfj/transducer-loss-benchmarking) compares the speed and memory usage of several transducer losses, the summary in the following table is taken from there, you can check the repository for more details.
Note: As we declared above, `fast_rnnt` is also implemented in [k2](https://github.com/k2-fsa/k2) project, so `k2` and `fast_rnnt` are equivalent in the benchmarking.
|Name |Average step time (us) | Peak memory usage (MB)|
This project implements a method for faster and more memory-efficient RNN-T loss computation, called `pruned rnnt`.
Note: There is also a fast RNN-T loss implementation in [k2](https://github.com/k2-fsa/k2) project, which shares the same code here. We make `fast_rnnt` a stand-alone project in case someone wants only this rnnt loss.
## How does the pruned-rnnt work ?
We first obtain pruning bounds for the RNN-T recursion using a simple joiner network that is just an addition of the encoder and decoder, then we use those pruning bounds to evaluate the full, non-linear joiner network.
The picture below display the gradients (obtained by `rnnt_loss_simple` with `return_grad=true`) of lattice nodes, at each time frame, only a small set of nodes have a non-zero gradient, which justifies the pruned RNN-T loss, i.e., putting a limit on the number of symbols per frame.
It will pass `-DCMAKE_BUILD_TYPE=Release -DFT_WITH_CUDA=OFF` to `cmake`.
### Where to get help if I have problems with the installation ?
Please file an issue at <https://github.com/danpovey/fast_rnnt/issues>
and describe your problem there.
## Usage
### For rnnt_loss_simple
This is a simple case of the RNN-T loss, where the joiner network is just
addition.
Note: termination_symbol plays the role of blank in other RNN-T loss implementations, we call it termination_symbol as it terminates symbols of current frame.
```python
am=torch.randn((B,T,C),dtype=torch.float32)
lm=torch.randn((B,S+1,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss_simple(
lm=lm,
am=am,
symbols=symbols,
termination_symbol=termination_symbol,
boundary=boundary,
reduction="sum",
)
```
### For rnnt_loss_smoothed
The same as `rnnt_loss_simple`, except that it supports `am_only` & `lm_only` smoothing
that allows you to make the loss-function one of the form:
lm_only_scale * lm_probs +
am_only_scale * am_probs +
(1-lm_only_scale-am_only_scale) * combined_probs
where `lm_probs` and `am_probs` are the probabilities given the lm and acoustic model independently.
```python
am=torch.randn((B,T,C),dtype=torch.float32)
lm=torch.randn((B,S+1,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss_smoothed(
lm=lm,
am=am,
symbols=symbols,
termination_symbol=termination_symbol,
lm_only_scale=0.25,
am_only_scale=0.0
boundary=boundary,
reduction="sum",
)
```
### For rnnt_loss_pruned
`rnnt_loss_pruned` can not be used alone, it needs the gradients returned by `rnnt_loss_simple/rnnt_loss_smoothed` to get pruning bounds.
You can also find recipes [here](https://github.com/k2-fsa/icefall/tree/master/egs/librispeech/ASR/pruned_transducer_stateless) that uses `rnnt_loss_pruned` to train a model.
### For rnnt_loss
The `unprund rnnt_loss` is the same as `torchaudio rnnt_loss`, it produces same output as torchaudio for the same input.
```python
logits=torch.randn((B,S,T,C),dtype=torch.float32)
symbols=torch.randint(0,C,(B,S))
termination_symbol=0
boundary=torch.zeros((B,4),dtype=torch.int64)
boundary[:,2]=target_lengths
boundary[:,3]=num_frames
loss=fast_rnnt.rnnt_loss(
logits=logits,
symbols=symbols,
termination_symbol=termination_symbol,
boundary=boundary,
reduction="sum",
)
```
## Benchmarking
The [repo](https://github.com/csukuangfj/transducer-loss-benchmarking) compares the speed and memory usage of several transducer losses, the summary in the following table is taken from there, you can check the repository for more details.
Note: As we declared above, `fast_rnnt` is also implemented in [k2](https://github.com/k2-fsa/k2) project, so `k2` and `fast_rnnt` are equivalent in the benchmarking.
|Name |Average step time (us) | Peak memory usage (MB)|
