def _ctc_loss_v3(labels, logits, label_length, logit_length, use_gpu): with test_util.device(use_gpu=use_gpu): sparse_labels = ctc_ops.dense_labels_to_sparse(labels, label_length) with backprop.GradientTape() as t: t.watch(logits) ref_loss = ctc_ops.ctc_loss_v3(labels=sparse_labels, logits=logits, label_length=label_length, logit_length=logit_length, blank_index=0) ref_grad = t.gradient(ref_loss, [logits]) return ref_loss, ref_grad
def _forwardAndBackward(self, sparse_labels, logits_time_major, seed): np.random.seed(seed) params = self._genInputParams(logits_time_major=logits_time_major, sparse_labels=sparse_labels) labels, logits, labels_lengths, logits_lengths = params output_shape = (labels_lengths.shape[0], ) upstream_gradients = self._randomFloats(output_shape) with backprop.GradientTape() as tape: tape.watch(logits) loss = ctc_ops.ctc_loss_v3(labels, logits, labels_lengths, logits_lengths, logits_time_major=logits_time_major, blank_index=0) gradient_injector_output = loss * upstream_gradients return loss, tape.gradient(gradient_injector_output, logits)