Package: optimized-transducer
IntroductionThis project implements the optimization techniques proposed inImproving RNN Transducer Modeling for End-to-End Speech Recognitionto reduce the memory consumption for computing transducer loss.HINT: You can find ASR training code using this repo
inhttps://github.com/k2-fsa/icefall. You can also find
decoding code inicefall.How does it differ from the RNN-T loss from torchaudioIt produces same output astorchaudiofor the same input, sooptimized_transducershould be equivalent totorchaudio.functional.rnnt_loss().This project is more memory efficient (Seehttps://github.com/csukuangfj/transducer-loss-benchmarkingfor benchmark results)Also,torchaudioaccepts only output fromnn.Linear, but
we also support output fromlog-softmax(You can set the optionfrom_log_softmaxtoTruein this case).It also supports amodifiedversion of transducer. Seebelowfor what
the meaning ofmodified transduceris.How does it differ fromwarp-transducerIt borrows the methods of computing alpha and beta fromwarp-transducer. Therefore,optimized_transducerproduces the samealphaandbetaaswarp-transducerfor the same input.However,warp-transducerproduces different gradients for CPU and CUDA
when using the same input. Seehttps://github.com/HawkAaron/warp-transducer/issues/93.
I also created acolab notebookto reproduce that issue.This project produces consistent gradient on CPU and CUDA for the same input, just like
whattorchaudiois doing. (We borrow the gradient computation formula fromtorchaudio).optimized_transduceruses less memory than that ofwarp-transducer(Seehttps://github.com/csukuangfj/transducer-loss-benchmarkingfor benchmark results).It also supports amodifiedversion of transducer. Seebelowfor what
the meaning ofmodified transduceris.Modified TransducerInmodified transducer, we limit the maximum number of symbols per frame to 1. The following
figure compares the formula for forward and backward procedures between standard transducer
and modified transducer.Note: Modified transducer is proposed independently by@danpovey.
We were later informed that the idea already existed inRecurrent Neural Aligner: An Encoder-Decoder Neural Network Model for
Sequence to Sequence MappingInstallationYou can install it viapip:pip install optimized_transducerTo check thatoptimized_transducerwas installed successfully, please runpython3 -c "import optimized_transducer; print(optimized_transducer.__version__)"which should print the version of the installedoptimized_transducer, e.g.,1.2.Installation FAQWhat operating systems are supported ?It has been tested on Ubuntu 18.04. It should also work on macOS and other unixes systems.
It may work on Windows, though it is not tested.How to display installation log ?Usepip install --verbose optimized_transducerHow to reduce installation time ?Useexport OT_MAKE_ARGS="-j"
pip install --verbose optimized_transducerIt will pass-jtomake.Which version of PyTorch is supported ?It has been tested on PyTorch >= 1.5.0. It may work on PyTorch < 1.5.0How to install a CPU version ofoptimized_transducer?Useexport OT_CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release -DOT_WITH_CUDA=OFF"
export OT_MAKE_ARGS="-j"
pip install --verbose optimized_transducerIt will pass-DCMAKE_BUILD_TYPE=Release -DOT_WITH_CUDA=OFFtocmake.What Python versions are supported ?Python >= 3.6 is known to work. It may work for Python 2.7, though it is not tested.Where to get help if I have problems with the installation ?Please file an issue athttps://github.com/csukuangfj/optimized_transducer/issuesand describe your problem there.Usageoptimized_transducerexpects that the output shape of the joint network isNOT(N, T, U, V), but is(sum_all_TU, V), which is a concatenation
of 2-D tensors:(T_1 * U_1, V),(T_2 * U_2, V), ...,(T_N, U_N, V).Note:(T_1 * U_1, V)is just the reshape of a 3-D tensor(T_1, U_1, V).Suppose your original joint network looks somewhat like the following:encoder_out=torch.rand(N,T,D)# from the encoderdecoder_out=torch.rand(N,U,D)# from the decoder, i.e., the prediction networkencoder_out=encoder_out.unsqueeze(2)# Now encoder out is (N, T, 1, D)decoder_out=decoder_out.unsqueeze(1)# Now decoder out is (N, 1, U, D)x=encoder_out+decoder_out# x is of shape (N, T, U, D)activation=torch.tanh(x)logits=linear(activation)# linear is an instance of `nn.Linear`.loss=torchaudio.functional.rnnt_loss(logits=logits,targets=targets,logit_lengths=logit_lengths,target_lengths=target_lengths,blank=blank_id,reduction="mean",)You need to change it to the following:encoder_out=torch.rand(N,T,D)# from the encoderdecoder_out=torch.rand(N,U,D)# from the decoder, i.e., the prediction networkencoder_out_list=[encoder_out[i,:logit_lengths[i],:]foriinrange(N)]decoder_out_list=[decoder_out[i,:target_lengths[i]+1,:]foriinrange(N)]x=[e.unsqueeze(1)+d.unsqueeze(0)fore,dinzip(encoder_out_list,decoder_out_list)]x=[p.reshape(-1,D)forpinx]x=torch.cat(x)activation=torch.tanh(x)logits=linear(activation)# linear is an instance of `nn.Linear`.loss=optimized_transducer.transducer_loss(logits=logits,targets=targets,logit_lengths=logit_lengths,target_lengths=target_lengths,blank=blank_id,reduction="mean",from_log_softmax=False,)Caution: We usedfrom_log_softmax=Falsein the above example sincelogitsis the output ofnn.Linear.Hint: Iflogitsis the output oflog-softmax, you should usefrom_log_softmax=True.In most cases, you should pass the output ofnn.Linearto compute the loss, i.e.,
usefrom_log_softmax=False, to save memory.If you want to do some operations on the output oflog-softmaxbefore feeding it
tooptimized_transducer.transducer_loss(),from_log_softmax=Trueis helpful in
this case. But be aware that this will increase the memory usage.To use themodifiedtransducer, pass an additional argumentone_sym_per_frame=Truetooptimized_transducer.transducer_loss().For more usages, please refer tohttps://github.com/csukuangfj/optimized_transducer/blob/master/optimized_transducer/python/optimized_transducer/transducer_loss.pyhttps://github.com/csukuangfj/optimized_transducer/blob/master/optimized_transducer/python/tests/test_cuda.pyhttps://github.com/csukuangfj/optimized_transducer/blob/master/optimized_transducer/python/tests/test_compute_transducer_loss.pyhttps://github.com/csukuangfj/optimized_transducer/blob/master/optimized_transducer/python/tests/test_max_symbol_per_frame.pyFor developersAs a developer, you don't need to usepip install optimized_transducer.
To make development easier, you can usegit clone https://github.com/csukuangfj/optimized_transducer.git
cd optimized_transducer
mkdir build
cd build
cmake -DOT_BUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Release ..
make -j
export PYTHONPATH=$PWD/../optimized_transducer/python:$PWD/lib:$PYTHONPATHI usually create a filepath.shinside thebuilddirectory, containingexport PYTHONPATH=$PWD/../optimized_transducer/python:$PWD/lib:$PYTHONPATHso what you need to do iscd optimized_transducer/build
source path.sh

# Then you are ready to run Python tests
python3 optimized_transducer/python/tests/test_compute_transducer_loss.py

# You can also use "import optimized_transducer" in your Python projectsTo run all Python tests, usecd optimized_transducer/build
ctest --output-on-failureAlternatively one can "make" all available testsmake -j test