Merge pull request #129055 from hasrat17:fix/ctc-loss-validate-label-class-ids

PiperOrigin-RevId: 1005376867

tensorflow/core/kernels/ctc_loss_op.cc

@@ -347,6 +347,21 @@ class CTCLossOpGPU : public OpKernel {
 
     absl::Span<const int32_t> labels_data(labels_values->flat<int32_t>().data(),
                                           num_indices);
+
+    // Validate label values before handing off to cuDNN, which does not
+    // produce a clear error for out-of-range class ids.
+    for (int i = 0; i < num_indices; ++i) {
+      const int32_t label_val = labels_data[i];
+      OP_REQUIRES(
+          ctx, label_val >= 0,
+          absl::InvalidArgumentError(absl::StrCat(
+              "All labels must be nonnegative integers, saw: ", label_val)));
+      OP_REQUIRES(ctx, label_val < num_classes,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "No label may be greater than num_classes. num_classes: ",
+                      num_classes, ", saw: ", label_val)));
+    }
+
     absl::Span<const int32_t> labels_lengths_data(labels_lengths.data(),
                                                   batch_size);
     absl::Span<const int32_t> input_lengths_data(

tensorflow/core/util/ctc/ctc_loss_calculator.h

@@ -324,8 +324,20 @@ absl::Status CTCLossCalculator<T>::PopulateLPrimes(
     // Convert label from DistBelief
     bool finished_sequence = false;
     for (int i = 0; i < label.size(); ++i) {
+      // Validate each label before using it so out-of-range values do not
+      // silently alias as sequence terminators.
+      if (label[i] < 0) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("All labels must be nonnegative integers, batch: ", b,
+                         " labels: ", absl::StrJoin(label, ",")));
+      } else if (label[i] >= num_classes) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("No label may be greater than num_classes. ",
+                         "num_classes: ", num_classes, ", batch: ", b,
+                         " labels: ", absl::StrJoin(label, ",")));
+      }
       if (i == 0 || !preprocess_collapse_repeated || label[i] != label[i - 1]) {
-        if (label[i] >= num_classes - 1) {
+        if (label[i] == num_classes - 1) {
           finished_sequence = true;
         } else {
           if (finished_sequence) {
@@ -342,18 +354,6 @@ absl::Status CTCLossCalculator<T>::PopulateLPrimes(
         }
       }
     }
-
-    for (int l_i : l) {
-      if (l_i < 0) {
-        return absl::InvalidArgumentError(
-            absl::StrCat("All labels must be nonnegative integers, batch: ", b,
-                         " labels: ", absl::StrJoin(l, ",")));
-      } else if (l_i >= num_classes) {
-        return absl::InvalidArgumentError(absl::StrCat(
-            "No label may be greater than num_classes. ", "num_classes: ",
-            num_classes, ", batch: ", b, " labels: ", absl::StrJoin(l, ",")));
-      }
-    }
     if (!ignore_longer_outputs_than_inputs) {
       // Make sure there is enough time to output the target indices.
       int time = seq_len(b) - output_delay_;

tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py

@@ -79,7 +79,7 @@ def _ctc_loss_v2(labels, inputs, sequence_length,
       logits_time_major=time_major)
 
 
-class CTCLossValidationTest(test.TestCase):
+class CTCLossValidationTest(test.TestCase, parameterized.TestCase):
 
   def _callCtcLoss(self, ctc_loss_fn, logits, blank_index=0):
     labels = constant_op.constant([[0]], dtype=dtypes.int32)
@@ -118,6 +118,222 @@ class CTCLossValidationTest(test.TestCase):
         with self.assertRaisesRegex(ValueError, r"must be in range \[-3, 3\)"):
           self._callCtcLoss(ctc_loss_fn, logits, blank_index=-4)
 
+  def _out_of_range_error(self, sparse, label):
+    """Returns the error regex expected for an out-of-range label id."""
+    if not sparse:
+      return r"All labels must be in the range \[0, num_classes\)"
+    if label < 0:
+      return r"All labels must be nonnegative integers"
+    return r"No label may be greater than num_classes"
+
+  def _make_labels(self, sparse, values):
+    """Creates a labels tensor for CTC loss testing."""
+    if sparse:
+      return sparse_tensor.SparseTensor(
+          indices=constant_op.constant(
+              [[0, i] for i in range(len(values))], dtype=dtypes.int64
+          ),
+          values=constant_op.constant(values, dtype=dtypes.int32),
+          dense_shape=constant_op.constant(
+              [1, len(values)], dtype=dtypes.int64
+          ),
+      )
+    return constant_op.constant([values], dtype=dtypes.int32)
+
+  def _call_ctc_loss_with_labels(
+      self, ctc_loss_fn, sparse, values, blank_index=-1
+  ):
+    """Invokes ctc_loss_fn with the provided label values and zero logits."""
+    logits = array_ops.zeros([3, 1, 10], dtype=dtypes.float32)
+    return ctc_loss_fn(
+        labels=self._make_labels(sparse, values),
+        logits=logits,
+        label_length=constant_op.constant([len(values)], dtype=dtypes.int32),
+        logit_length=constant_op.constant([3], dtype=dtypes.int32),
+        blank_index=blank_index,
+    )
+
+  @parameterized.product(
+      ctc_loss_fn=[ctc_ops.ctc_loss_v2, ctc_ops.ctc_loss_v3],
+      sparse=[True, False],
+      label=[-1, 10, 2147483647],
+  )
+  def testOutOfRangeLabelRejected(self, ctc_loss_fn, sparse, label):
+    """Tests that out-of-range label ids are rejected in eager mode."""
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        self._out_of_range_error(sparse, label),
+    ):
+      self.evaluate(
+          self._call_ctc_loss_with_labels(ctc_loss_fn, sparse, [label])
+      )
+
+  @parameterized.product(
+      ctc_loss_fn=[ctc_ops.ctc_loss_v2, ctc_ops.ctc_loss_v3],
+      sparse=[True, False],
+      label=[-1, 10, 2147483647],
+  )
+  def testOutOfRangeLabelRejectedInFunction(self, ctc_loss_fn, sparse, label):
+    """Tests that out-of-range label ids are rejected inside tf.function."""
+
+    @def_function.function
+    def test_fn():
+      return self._call_ctc_loss_with_labels(ctc_loss_fn, sparse, [label])
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        self._out_of_range_error(sparse, label),
+    ):
+      self.evaluate(test_fn())
+
+  @parameterized.product(
+      ctc_loss_fn=[ctc_ops.ctc_loss_v2, ctc_ops.ctc_loss_v3],
+      sparse=[True, False],
+  )
+  def testInRangeLabelsStillAccepted(self, ctc_loss_fn, sparse):
+    """Tests that valid labels are still accepted."""
+    loss = self.evaluate(
+        self._call_ctc_loss_with_labels(ctc_loss_fn, sparse, [0, 9])
+    )
+    self.assertAllEqual(loss.shape, [1])
+    self.assertTrue(np.all(np.isfinite(loss)))
+
+  def testDensePaddingNegativeValueIgnored(self):
+    """Padding with -1 past label_length must be ignored (negative OOB)."""
+    logits = array_ops.zeros([3, 1, 10], dtype=dtypes.float32)
+    # Reference: label [3], no padding.
+    loss_ref = self.evaluate(
+        ctc_ops.ctc_loss_v3(
+            labels=constant_op.constant([[3, 0]], dtype=dtypes.int32),
+            logits=logits,
+            label_length=constant_op.constant([1], dtype=dtypes.int32),
+            logit_length=constant_op.constant([3], dtype=dtypes.int32),
+            blank_index=-1,
+        )
+    )
+    # Same label with invalid padding value -1 past label_length.
+    loss_padded = self.evaluate(
+        ctc_ops.ctc_loss_v3(
+            labels=constant_op.constant([[3, -1]], dtype=dtypes.int32),
+            logits=logits,
+            label_length=constant_op.constant([1], dtype=dtypes.int32),
+            logit_length=constant_op.constant([3], dtype=dtypes.int32),
+            blank_index=-1,
+        )
+    )
+    self.assertAllClose(loss_ref, loss_padded)
+
+  def testDensePaddingPositiveOobValueIgnored(self):
+    """Padding with positive OOB ids (e.g. 10, 99) past label_length is ok."""
+    logits = array_ops.zeros([3, 1, 10], dtype=dtypes.float32)
+    loss_ref = self.evaluate(
+        ctc_ops.ctc_loss_v3(
+            labels=constant_op.constant([[3, 0]], dtype=dtypes.int32),
+            logits=logits,
+            label_length=constant_op.constant([1], dtype=dtypes.int32),
+            logit_length=constant_op.constant([3], dtype=dtypes.int32),
+            blank_index=-1,
+        )
+    )
+    for pad_val in [10, 99]:
+      loss_padded = self.evaluate(
+          ctc_ops.ctc_loss_v3(
+              labels=constant_op.constant([[3, pad_val]], dtype=dtypes.int32),
+              logits=logits,
+              label_length=constant_op.constant([1], dtype=dtypes.int32),
+              logit_length=constant_op.constant([3], dtype=dtypes.int32),
+              blank_index=-1,
+          )
+      )
+      self.assertAllClose(
+          loss_ref,
+          loss_padded,
+          msg=f"Padding value {pad_val} should be ignored",
+      )
+
+  @parameterized.product(
+      ctc_loss_fn=[ctc_ops.ctc_loss_v2, ctc_ops.ctc_loss_v3],
+  )
+  def testUnknownRankLogitsRejected(self, ctc_loss_fn):
+    """Tests that unknown-rank logits do not crash _get_dim and reject OOB.
+
+    Dense labels route through ctc_loss_dense, which calls _get_dim on both
+    logits and labels.  Tracing with TensorSpec(shape=None) gives logits an
+    unknown rank (labels keep a static shape), exercising the fallback in
+    _get_dim.
+    """
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]
+    )
+    def test_fn(logits):
+      return ctc_loss_fn(
+          labels=self._make_labels(False, [10]),
+          logits=logits,
+          label_length=constant_op.constant([1], dtype=dtypes.int32),
+          logit_length=constant_op.constant([3], dtype=dtypes.int32),
+          blank_index=-1,
+      )
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        self._out_of_range_error(False, 10),
+    ):
+      self.evaluate(test_fn(array_ops.zeros([3, 1, 10], dtype=dtypes.float32)))
+
+  @parameterized.product(
+      ctc_loss_fn=[ctc_ops.ctc_loss_v2, ctc_ops.ctc_loss_v3],
+  )
+  def testUnknownRankLogitsAccepted(self, ctc_loss_fn):
+    """Tests that valid labels with unknown-rank logits are accepted.
+
+    Uses blank_index=-1 so blank_index + _get_dim(logits, 2) is computed at
+    trace time, exercising blank_index_val under unknown rank.
+    """
+
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]
+    )
+    def test_fn(logits):
+      return ctc_loss_fn(
+          labels=self._make_labels(False, [1, 2]),
+          logits=logits,
+          label_length=constant_op.constant([2], dtype=dtypes.int32),
+          logit_length=constant_op.constant([3], dtype=dtypes.int32),
+          blank_index=-1,
+      )
+
+    loss = self.evaluate(
+        test_fn(array_ops.zeros([3, 1, 10], dtype=dtypes.float32))
+    )
+    self.assertAllEqual(loss.shape, [1])
+    self.assertTrue(np.all(np.isfinite(loss)))
+
+  def testDensePaddingMixedLengths(self):
+    """Padding past each row's label_length is ignored per row."""
+    logits = array_ops.zeros([4, 2, 10], dtype=dtypes.float32)
+    label_length = constant_op.constant([1, 2], dtype=dtypes.int32)
+    logit_length = constant_op.constant([4, 4], dtype=dtypes.int32)
+
+    def loss(values):
+      return ctc_ops.ctc_loss_v3(
+          labels=constant_op.constant(values, dtype=dtypes.int32),
+          logits=logits,
+          label_length=label_length,
+          logit_length=logit_length,
+          blank_index=-1,
+      )
+
+    self.assertAllClose(
+        self.evaluate(loss([[3, 0], [4, 5]])),
+        self.evaluate(loss([[3, 99], [4, 5]])),
+    )
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError,
+        r"All labels must be in the range \[0, num_classes\)",
+    ):
+      self.evaluate(loss([[3, 0], [4, 10]]))
+
 
 class CTCLossTest(test.TestCase):
 

tensorflow/python/ops/BUILD

@@ -1499,6 +1499,7 @@ py_library(
         ":array_ops",
         ":array_ops_gen",
         ":array_ops_stack",
+        ":control_flow_assert",
         ":ctc_ops_gen",
         ":custom_gradient",
         ":functional_ops",

tensorflow/python/ops/ctc_ops.py

@@ -18,7 +18,6 @@ import uuid
 
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
-
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import device
 from tensorflow.python.framework import dtypes
@@ -26,9 +25,9 @@ from tensorflow.python.framework import function
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import tensor_shape
-
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import control_flow_assert
 from tensorflow.python.ops import custom_gradient
 from tensorflow.python.ops import functional_ops
 from tensorflow.python.ops import gen_array_ops
@@ -749,10 +748,17 @@ def _ctc_loss_op_standard(labels, logits, logit_length, logits_time_major,
   part_after = logits[:, :, blank_index + 1:]
   part_blank = logits[:, :, blank_index:blank_index + 1]
   logits = array_ops.concat([part_before, part_after, part_blank], axis=2)
+  # Keep out-of-range ids unchanged so the kernel rejects them instead of
+  # aliasing num_classes onto the blank index num_classes - 1.
+  values = labels.values
+  shift = math_ops.logical_and(
+      values >= blank_index, values < _get_dim(logits, 2)
+  )
   labels = sparse_tensor.SparseTensor(
       labels.indices,
-      array_ops.where(labels.values < blank_index, labels.values,
-                      labels.values - 1), labels.dense_shape)
+      array_ops.where(shift, values - 1, values),
+      labels.dense_shape,
+  )
   return _ctc_loss_impl(
       labels=labels,
       inputs=logits,
@@ -767,10 +773,15 @@ def _ctc_loss_op_cudnn(labels, logits, logit_length, logits_time_major,
   part_after = logits[:, :, blank_index + 1:]
   part_blank = logits[:, :, blank_index:blank_index + 1]
   logits = array_ops.concat([part_blank, part_before, part_after], axis=2)
+  # Keep out-of-range ids unchanged so the kernel rejects them instead of
+  # shifting -1 to 0 (the blank index in cuDNN layout).
+  values = labels.values
+  shift = math_ops.logical_and(values >= 0, values < blank_index)
   labels = sparse_tensor.SparseTensor(
       labels.indices,
-      array_ops.where(labels.values < blank_index, labels.values + 1,
-                      labels.values), labels.dense_shape)
+      array_ops.where(shift, values + 1, values),
+      labels.dense_shape,
+  )
   return _ctc_loss_impl(
       labels=labels,
       inputs=logits,
@@ -783,6 +794,20 @@ def _ctc_loss_shape(op):
   return [op.inputs[2].get_shape(), op.inputs[0].get_shape()]
 
 
+def _assert_label_ids_in_range(labels, num_classes, valid_mask):
+  """Returns an Assert that unmasked labels are in [0, num_classes)."""
+  num_classes = math_ops.cast(num_classes, labels.dtype)
+  in_range = math_ops.logical_and(labels >= 0, labels < num_classes)
+  in_range = math_ops.logical_or(in_range, math_ops.logical_not(valid_mask))
+  return control_flow_assert.Assert(
+      math_ops.reduce_all(in_range),
+      [
+          "All labels must be in the range [0, num_classes), got num_classes =",
+          num_classes,
+      ],
+  )
+
+
 def _validate_ctc_loss_logits(logits, blank_index):
   rank = logits.shape.rank
   if rank is not None and rank < 3:
@@ -885,10 +910,17 @@ def ctc_loss_v2(labels,
           logits[:, :, blank_index:blank_index + 1],
       ],
                                 axis=2)
+      # Keep out-of-range ids unchanged so the kernel rejects them instead of
+      # aliasing num_classes onto the blank index num_classes - 1.
+      values = labels.values
+      shift = math_ops.logical_and(
+          values >= blank_index, values < _get_dim(logits, 2)
+      )
       labels = sparse_tensor.SparseTensor(
           labels.indices,
-          array_ops.where(labels.values < blank_index, labels.values,
-                          labels.values - 1), labels.dense_shape)
+          array_ops.where(shift, values - 1, values),
+          labels.dense_shape,
+      )
 
     return ctc_loss(
         labels=labels,
@@ -1146,24 +1178,43 @@ def ctc_loss_dense(labels,
     if not logits_time_major:
       logits = array_ops.transpose(logits, perm=[1, 0, 2])
 
+    # Only the first label_length entries per row are real labels; padding may
+    # hold arbitrary values.  Validate before the blank_index remap so that -1
+    # is checked before it could be mapped to 0.
+    label_mask = array_ops.sequence_mask(
+        label_length, maxlen=_get_dim(labels, 1)
+    )
+    assert_op = _assert_label_ids_in_range(
+        labels, _get_dim(logits, 2), valid_mask=label_mask
+    )
+    if assert_op is not None:
+      with ops.control_dependencies([assert_op]):
+        labels = array_ops.identity(labels)
+
     if blank_index != 0:
-      if blank_index < 0:
-        blank_index += _get_dim(logits, 2)
-      logits = array_ops.concat([
-          logits[:, :, blank_index:blank_index + 1],
-          logits[:, :, :blank_index],
-          logits[:, :, blank_index + 1:],
-      ],
-                                axis=2)
-      labels = array_ops.where(labels < blank_index, labels + 1, labels)
+      blank_index_val = (
+          blank_index + _get_dim(logits, 2) if blank_index < 0 else blank_index
+      )
+      logits = array_ops.concat(
+          [
+              logits[:, :, blank_index_val : blank_index_val + 1],
+              logits[:, :, :blank_index_val],
+              logits[:, :, blank_index_val + 1 :],
+          ],
+          axis=2,
+      )
+      labels = array_ops.where(labels < blank_index_val, labels + 1, labels)
+    else:
+      blank_index_val = blank_index
 
     args = [logits, labels, label_length, logit_length]
 
     if unique:
       unique_y, unique_idx = unique
       if blank_index != 0:
-        unique_y = array_ops.where(unique_y < blank_index, unique_y + 1,
-                                   unique_y)
+        unique_y = array_ops.where(
+            unique_y < blank_index_val, unique_y + 1, unique_y
+        )
         label_mask_len = math_ops.reduce_max(unique_idx, axis=1) + 1
         max_label_length = _get_dim(unique_y, 1)
         label_mask = array_ops.sequence_mask(label_mask_len, max_label_length)
@@ -1568,5 +1619,7 @@ def _scan(fn, elems, initial, reverse=False, inclusive=False, final_only=False):
 
 def _get_dim(tensor, i):
   """Get value of tensor shape[i] preferring static value if available."""
+  if tensor.shape.rank is None or tensor.shape.rank <= i:
+    return array_ops.shape(tensor)[i]
   return tensor_shape.dimension_value(
       tensor.shape[i]) or array_ops.shape(tensor)[i]