Package: fast-rnnt
This project implements a method for faster and more memory-efficient RNN-T loss computation, calledpruned rnnt.Note: There is also a fast RNN-T loss implementation ink2project, which shares the same code here. We makefast_rnnta 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 byrnnt_loss_simplewithreturn_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.This picture is taken fromhereInstallationYou can install it viapip:pip install fast_rnntYou can also install from source:git clone https://github.com/danpovey/fast_rnnt.git
cd fast_rnnt
python setup.py installTo check thatfast_rnntwas installed successfully, please runpython3 -c "import fast_rnnt; print(fast_rnnt.__version__)"which should print the version of the installedfast_rnnt, e.g.,1.0.How to display installation log ?Usepip install --verbose fast_rnntHow to reduce installation time ?Useexport FT_MAKE_ARGS="-j"
pip install --verbose fast_rnntIt will pass-jtomake.Which version of PyTorch is supported ?It has been tested on PyTorch >= 1.5.0.Note: The cuda version of the Pytorch should be the same as the cuda version in your environment,
or it will cause a compilation error.How to install a CPU version offast_rnnt?Useexport FT_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release -DFT_WITH_CUDA=OFF"
export FT_MAKE_ARGS="-j"
pip install --verbose fast_rnntIt will pass-DCMAKE_BUILD_TYPE=Release -DFT_WITH_CUDA=OFFtocmake.Where to get help if I have problems with the installation ?Please file an issue athttps://github.com/danpovey/fast_rnnt/issuesand describe your problem there.UsageFor rnnt_loss_simpleThis 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.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=0boundary=torch.zeros((B,4),dtype=torch.int64)boundary[:,2]=target_lengthsboundary[:,3]=num_framesloss=fast_rnnt.rnnt_loss_simple(lm=lm,am=am,symbols=symbols,termination_symbol=termination_symbol,boundary=boundary,reduction="sum",)For rnnt_loss_smoothedThe same asrnnt_loss_simple, except that it supportsam_only&lm_onlysmoothing
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_probswherelm_probsandam_probsare the probabilities given the lm and acoustic model independently.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=0boundary=torch.zeros((B,4),dtype=torch.int64)boundary[:,2]=target_lengthsboundary[:,3]=num_framesloss=fast_rnnt.rnnt_loss_smoothed(lm=lm,am=am,symbols=symbols,termination_symbol=termination_symbol,lm_only_scale=0.25,am_only_scale=0.0boundary=boundary,reduction="sum",)For rnnt_loss_prunedrnnt_loss_prunedcan not be used alone, it needs the gradients returned byrnnt_loss_simple/rnnt_loss_smoothedto get pruning bounds.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=0boundary=torch.zeros((B,4),dtype=torch.int64)boundary[:,2]=target_lengthsboundary[:,3]=num_frames# rnnt_loss_simple can be also rnnt_loss_smoothedsimple_loss,(px_grad,py_grad)=fast_rnnt.rnnt_loss_simple(lm=lm,am=am,symbols=symbols,termination_symbol=termination_symbol,boundary=boundary,reduction="sum",return_grad=True,)s_range=5# can be other valuesranges=fast_rnnt.get_rnnt_prune_ranges(px_grad=px_grad,py_grad=py_grad,boundary=boundary,s_range=s_range,)am_pruned,lm_pruned=fast_rnnt.do_rnnt_pruning(am=am,lm=lm,ranges=ranges)logits=model.joiner(am_pruned,lm_pruned)pruned_loss=fast_rnnt.rnnt_loss_pruned(logits=logits,symbols=symbols,ranges=ranges,termination_symbol=termination_symbol,boundary=boundary,reduction="sum",)You can also find recipesherethat usesrnnt_loss_prunedto train a model.For rnnt_lossTheunprund rnnt_lossis the same astorchaudio rnnt_loss, it produces same output as torchaudio for the same input.logits=torch.randn((B,S,T,C),dtype=torch.float32)symbols=torch.randint(0,C,(B,S))termination_symbol=0boundary=torch.zeros((B,4),dtype=torch.int64)boundary[:,2]=target_lengthsboundary[:,3]=num_framesloss=fast_rnnt.rnnt_loss(logits=logits,symbols=symbols,termination_symbol=termination_symbol,boundary=boundary,reduction="sum",)BenchmarkingTherepocompares 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_rnntis also implemented ink2project, sok2andfast_rnntare equivalent in the benchmarking.NameAverage step time (us)Peak memory usage (MB)torchaudio60144712959.2fast_rnnt(unpruned)27440715106.5fast_rnnt(pruned)381122647.8optimized_transducer56768410903.1warprnnt_numba22934013061.8warp-transducer21077213061.8