Tensorflow-compatible Python library that provides second-order derivative for CTC loss function

Question

I need a TensorFlow-compatible Python library that provides a second-order derivative for Machine Learning CTC (Connectionist Temporal Classification) loss function.

Namely, a library that provides an implementation of the ctc_loss function in Python API such that, in contrast to tf.nn.ctc_loss, the second gradient must be calculable analytically with reasonable performance like this:

with tf.GradientTape() as tape1: 
    with tf.GradientTape() as tape2:
        loss = ctc_loss(logits=logits, ...)
    gradient = tape2.gradient(loss, sources=logits)
hessian = tape1.gradient(gradient, sources=logits)

This may be needed, for example, for the influence function.

Answer 1

github.com/alexeytochin/tf_seq2seq_losses (disclaimer: this is my own project) is a library that offers such a function.

To use it, install tf-seq2seq-losses pip package

pip install tf-seq2seq-losses

Example usage:

import tensorflow as tf
from tf_seq2seq_losses import classic_ctc_loss

batch_size = 1
num_tokens = 3
logit_length = 5
logits=tf.zeros(shape=[batch_size, logit_length, num_tokens], dtype=tf.float32)
with tf.GradientTape(persistent=True) as tape1: 
    tape1.watch([logits])
    with tf.GradientTape() as tape2:
        tape2.watch([logits])
        loss = classic_ctc_loss(
            labels=tf.constant([[1, 2, 2, 1]], dtype=tf.int32),
            logits=logits,
            label_length=tf.constant([4], dtype=tf.int32),
            logit_length=tf.constant([logit_length], dtype=tf.int32),
            blank_index=0,
        )
    gradient = tape2.gradient(loss, sources=logits)
hessian = tape1.batch_jacobian(gradient, source=logits, experimental_use_pfor=False)

License: Apache License, Version 2.0