Package: tf-seq2seq-losses
tf-seq2seq-lossesTensorflow implementations forConnectionist Temporal Classification(CTC) loss in TensorFlow.InstallationTested with Python 3.7.$pipinstalltf-seq2seq-lossesWhy1. FasterOfficial CTC loss implementationtf.nn.ctc_lossis unreasonably slow.
The proposed implementation is approximately 30 times faster as it follows form the benchmark:nameforward timegradient calculation timetf.nn.ctc_loss13.2 ± 0.0210.4 ± 3.classic_ctc_loss0.138 ± 0.0060.28 ± 0.01simple_ctc_loss0.0531 ± 0.0030.119 ± 0.004(Tested on single GPU: GeForce GTX 970,  Driver Version: 460.91.03, CUDA Version: 11.2). Seebenchmark.pyfor the experimental setup. To reproduce this benchmark use$pytest-olog_cli=true--log-level=INFOtests/benchmark.pyfrom the project root directory.
Hereclassic_ctc_lossis the standard version of CTC loss
that corresponds to the decoding with repeated tokens collapse likea_bb_ccc_c   ->   abcc(equivalent totensorflow.nn.ctc_loss).
The loss functionsimple_ctc_lossis a simplified version corresponding to the trivial decoding rule, for example,a_bb_ccc_c   ->   abbcccc(simple blank removing).2. Supports second order derivativeimporttensorflowastffromtf_seq2seq_lossesimportclassic_ctc_lossbatch_size=2num_tokens=3logit_length=5labels=tf.constant([[1,2,2,1],[1,2,1,0]],dtype=tf.int32)label_length=tf.constant([4,3],dtype=tf.int32)logits=tf.zeros(shape=[batch_size,logit_length,num_tokens],dtype=tf.float32)logit_length=tf.constant([5,4],dtype=tf.int32)withtf.GradientTape(persistent=True)astape1:tape1.watch([logits])withtf.GradientTape()astape2:tape2.watch([logits])loss=tf.reduce_sum(classic_ctc_loss(labels=labels,logits=logits,label_length=label_length,logit_length=logit_length,blank_index=0,))gradient=tape2.gradient(loss,sources=logits)hessian=tape1.batch_jacobian(gradient,source=logits,experimental_use_pfor=False)# shape = [2, 5, 3, 5, 3]As well as for the first order gradient the TensorFlow autogradient is not used here.
Instead, an approach similar tothis oneis used with complexitywhereis the sequence length (the complexity for the gradient is).3. Numerically stableProposed implementation is more numerically stable. For example, it calculates resonable output forlogitsof order1e+10and even for-tf.inf.If logit length is too short to predict label output the probability of expected prediction is zero.
Thus, the loss output istf.inffor this sample but not707.13184fortf.nn.ctc_loss.4. No C++ compilationThis is a pure Python/TensorFlow implementation. We do not have to build or compile any C++/CUDA stuff.UsageThe interface is identical totensorflow.nn.ctc_losswithlogits_time_major=False.importtensorflowastffromtf_seq2seq_lossesimportclassic_ctc_lossbatch_size=1num_tokens=3# = 2 tokens + blanklogit_length=5loss=classic_ctc_loss(labels=tf.constant([[1,2,2,1]],dtype=tf.int32),logits=tf.zeros(shape=[batch_size,logit_length,num_tokens],dtype=tf.float32),label_length=tf.constant([4],dtype=tf.int32),logit_length=tf.constant([logit_length],dtype=tf.int32),blank_index=0,)Under the roofTensorFlow operations sich astf.while_loopandtf.TensorArray.
The bottleneck is the iterations over the logit length in order to calculate
α and β
(see, for example, the originalpaper). Expected gradient GPU calculation time is linear over logit length.Known ProbelemsWarning:AutoGraph could not transform <function classic_ctc_loss at ...> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux,export AUTOGRAPH_VERBOSITY=10) and attach the full output.observed for TensorFlow version 2.4.1
has no effect for performance. It is caused byUnionin type annotations.tf.jacobianandtf.batch_jacobianfor second derivative ofclassic_ctc_losswithexperimental_use_pfor=Falseintf.GradientTapecauses an unexpectedUnimplementedError.
for TensorFlow version 2.4.1. This can be avoided byexperimental_use_pfor=Trueor by usingClassicCtcLossData.hessaindirectly withouttf.GradientTape.Future plansAdd decoding (inference)Add rnn-t loss.Add m-wer loss.