diff --git a/tensorflow/core/kernels/ctc_loss_op.cc b/tensorflow/core/kernels/ctc_loss_op.cc index ec03efd9660..4baa1d49703 100644 --- a/tensorflow/core/kernels/ctc_loss_op.cc +++ b/tensorflow/core/kernels/ctc_loss_op.cc @@ -347,6 +347,21 @@ class CTCLossOpGPU : public OpKernel { absl::Span labels_data(labels_values->flat().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 labels_lengths_data(labels_lengths.data(), batch_size); absl::Span input_lengths_data( diff --git a/tensorflow/core/util/ctc/ctc_loss_calculator.h b/tensorflow/core/util/ctc/ctc_loss_calculator.h index 3ed3f07d609..5286afd819f 100644 --- a/tensorflow/core/util/ctc/ctc_loss_calculator.h +++ b/tensorflow/core/util/ctc/ctc_loss_calculator.h @@ -324,8 +324,20 @@ absl::Status CTCLossCalculator::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::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_; diff --git a/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py b/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py index 3aa788ade97..327e378c905 100644 --- a/tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py +++ b/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): diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 19c59a1efa4..d9d8e404e35 100644 --- a/tensorflow/python/ops/BUILD +++ b/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", diff --git a/tensorflow/python/ops/ctc_ops.py b/tensorflow/python/ops/ctc_ops.py index ec284ae6d83..050f6f492fb 100644 --- a/tensorflow/python/ops/ctc_ops.py +++ b/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]