Raise error when batch size does not match the first dimension of the shape.

PiperOrigin-RevId: 381062461

Author: tf-model-analysis-team

tensorflow_model_analysis/model_util.py

@@ -857,6 +857,12 @@ def maybe_expand_dims(arr):
     def to_dense(t):
       return tf.sparse.to_dense(t) if isinstance(t, tf.SparseTensor) else t
 
+    def check_shape(t, batch_size, key=None):
+      if t.shape[0] != batch_size:
+        raise ValueError(
+            'First dimension does not correspond with batch size. '
+            f'Batch size: {batch_size}, Dimensions: {t.shape}, Key: {key}.')
+
     result = copy.copy(batched_extract)
     record_batch = batched_extract[constants.ARROW_RECORD_BATCH_KEY]
     serialized_examples = batched_extract[constants.INPUT_KEY]
@@ -895,16 +901,26 @@ def to_dense(t):
               outputs = signature(inputs)
           else:
             outputs = signature(tf.constant(inputs, dtype=tf.string))
+
+          dense_outputs = {}
+          if isinstance(outputs, dict):
+            for k, v in outputs.items():
+              dense_outputs[k] = to_dense(v)
+              check_shape(dense_outputs[k], record_batch.num_rows, key=k)
+          else:
+            dense_outputs = to_dense(outputs)
+            check_shape(dense_outputs, record_batch.num_rows)
+
           for i in range(record_batch.num_rows):
-            if isinstance(outputs, dict):
+            if isinstance(dense_outputs, dict):
               output = {
-                  k: maybe_expand_dims(to_dense(v)[i].numpy())
-                  for k, v in outputs.items()
+                  k: maybe_expand_dims(v[i].numpy())
+                  for k, v in dense_outputs.items()
               }
             else:
               output = {
                   signature_name:
-                      maybe_expand_dims(np.asarray(to_dense(outputs))[i])
+                      maybe_expand_dims(np.asarray(dense_outputs)[i])
               }
             if result[extracts_key][i] is None:
               result[extracts_key][i] = collections.defaultdict(dict)

tensorflow_model_analysis/model_util_test.py

@@ -172,6 +172,39 @@ def custom_multi_output(features):
       tf.saved_model.save(model, export_path, signatures=signatures)
     return export_path
 
+  def createModelWithInvalidOutputShape(self):
+    input1 = tf.keras.layers.Input(shape=(1,), name='input_1')
+    input2 = tf.keras.layers.Input(shape=(1,), name='input_2')
+    inputs = [input1, input2]
+    input_layer = tf.keras.layers.concatenate(inputs)
+    output_layer = tf.keras.layers.Dense(
+        2, activation=tf.nn.sigmoid, name='output')(
+            input_layer)
+    # Flatten the layer such that the first dimension no longer corresponds
+    # with the batch size.
+    reshape_layer = tf.keras.layers.Lambda(
+        lambda x: tf.reshape(x, [-1]), name='reshape')(
+            output_layer)
+    model = tf.keras.models.Model(inputs, reshape_layer)
+
+    @tf.function
+    def serving_default(serialized_tf_examples):
+      parsed_features = tf.io.parse_example(
+          serialized_tf_examples, {
+              'input_1': tf.io.FixedLenFeature([1], dtype=tf.float32),
+              'input_2': tf.io.FixedLenFeature([1], dtype=tf.float32)
+          })
+      return model(parsed_features)
+
+    input_spec = tf.TensorSpec(shape=(None,), dtype=tf.string, name='examples')
+    signatures = {
+        'serving_default': serving_default.get_concrete_function(input_spec),
+    }
+
+    export_path = tempfile.mkdtemp()
+    model.save(export_path, save_format='tf', signatures=signatures)
+    return export_path
+
   def createModelWithMultipleMixedInputs(self, save_as_keras):
     dense_input = tf.keras.layers.Input(
         shape=(2,), name='input_1', dtype=tf.int64)
@@ -821,6 +854,50 @@ def check_result(got):
 
       util.assert_that(result, check_result, label='result')
 
+  def testModelSignaturesDoFnError(self):
+    export_path = self.createModelWithInvalidOutputShape()
+    signature_names = {constants.PREDICTIONS_KEY: {'': [None]}}
+    eval_shared_models = {
+        '':
+            self.createTestEvalSharedModel(
+                eval_saved_model_path=export_path,
+                tags=[tf.saved_model.SERVING])
+    }
+    model_specs = [config.ModelSpec()]
+    eval_config = config.EvalConfig(model_specs=model_specs)
+    schema = self.createDenseInputsSchema()
+    tfx_io = tf_example_record.TFExampleBeamRecord(
+        physical_format='text',
+        schema=schema,
+        raw_record_column_name=constants.ARROW_INPUT_COLUMN)
+    tensor_adapter_config = tensor_adapter.TensorAdapterConfig(
+        arrow_schema=tfx_io.ArrowSchema(),
+        tensor_representations=tfx_io.TensorRepresentations())
+
+    examples = [
+        self._makeExample(input_1=1.0, input_2=2.0),
+        self._makeExample(input_1=3.0, input_2=4.0),
+        self._makeExample(input_1=5.0, input_2=6.0),
+    ]
+
+    with self.assertRaisesRegex(
+        ValueError, 'First dimension does not correspond with batch size.'):
+      with beam.Pipeline() as pipeline:
+        # pylint: disable=no-value-for-parameter
+        _ = (
+            pipeline
+            | 'Create' >> beam.Create([e.SerializeToString() for e in examples])
+            | 'BatchExamples' >> tfx_io.BeamSource(batch_size=3)
+            | 'ToExtracts' >> beam.Map(_record_batch_to_extracts)
+            | 'ModelSignatures' >> beam.ParDo(
+                model_util.ModelSignaturesDoFn(
+                    eval_config=eval_config,
+                    eval_shared_models=eval_shared_models,
+                    signature_names=signature_names,
+                    default_signature_names=None,
+                    prefer_dict_outputs=False,
+                    tensor_adapter_config=tensor_adapter_config)))
+
   def testHasRubberStamp(self):
     # Model agnostic.
     self.assertFalse(model_util.has_rubber_stamp(None))