Introduction

The best way to run a TensorFlow SavedModel with different batch sizes is to export it with a flexible batch dimension in its input signature. If the SavedModel signature already allows None for the leading dimension, you can usually run batch size 1, 8, 32, or any other supported size without changing the model. If the SavedModel was exported with a fixed batch dimension, you generally need to re-export it rather than trying to patch the batch size at inference time.

The Batch Size Is Part of the Signature

A SavedModel is not only weights. It also includes callable signatures that define expected input names, dtypes, and shapes. TensorFlow's SavedModel tooling lets you inspect those shapes.

If the signature shape is something like (-1, 224, 224, 3), the leading dimension is flexible and different batch sizes are supported. If the shape is (32, 224, 224, 3), the exported signature expects batch size 32.

You can inspect the signature with the CLI:

saved_model_cli show --dir /path/to/model --tag_set serve --signature_def serving_default

That is the first thing to check before changing any inference code.

Export With a Flexible Batch Dimension

For custom modules or functions, define an input signature with None in the batch position.

1import tensorflow as tf
2
3class Doubler(tf.Module):
4    @tf.function(input_signature=[tf.TensorSpec([None, 4], tf.float32, name="x")])
5    def serve(self, x):
6        return {"y": x * 2.0}
7
8module = Doubler()
9tf.saved_model.save(module, "/tmp/doubler", signatures={"serving_default": module.serve})

Because the first dimension is None, the exported model can accept different batch sizes as long as the feature dimension still matches.

Load and Run With Different Batch Sizes

1import tensorflow as tf
2
3loaded = tf.saved_model.load("/tmp/doubler")
4infer = loaded.signatures["serving_default"]
5
6x1 = tf.ones([1, 4], dtype=tf.float32)
7x8 = tf.ones([8, 4], dtype=tf.float32)
8
9print(infer(x=x1)["y"])
10print(infer(x=x8)["y"])

The same SavedModel works for both batch sizes because the signature was exported that way.

Keras Models Usually Want the Same Principle

For Keras, the same rule applies. Build or export the model with a flexible leading dimension.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Input(shape=(4,)),
5    tf.keras.layers.Dense(2),
6])
7
8model.export("/tmp/keras_saved_model")

Keras input shapes normally leave the batch dimension flexible unless you deliberately constrain it. That makes many Keras SavedModels naturally compatible with multiple batch sizes.

What Happens If the Signature Is Fixed

If the SavedModel was exported from a traced function with a fixed input shape, loading it and calling it with a new batch size can fail because TensorFlow cannot find a matching concrete function.

That is why trying to "just feed a different batch size" sometimes works and sometimes fails. The deciding factor is the exported signature, not only the internal model layers.

If the signature is fixed, the best fix is usually to go back to the original model code and export a new SavedModel with a flexible leading dimension.

Performance Versus Compatibility

Supporting multiple batch sizes is a functional question. Performance is separate. A model that accepts multiple batch sizes may still have very different latency and throughput characteristics at batch size 1 versus batch size 64.

So the practical workflow is:

• export a flexible signature for compatibility
• benchmark different batch sizes for deployment
• choose a serving batch size based on latency and throughput needs

Do not confuse "the model accepts the batch size" with "this is the best batch size for production."

Common Pitfalls

The most common mistake is assuming batch size is always flexible just because the original training code used variable batches. The exported SavedModel signature is what matters at inference time.

Another issue is trying to patch the saved artifact after export instead of re-exporting from the original model with the correct signature.

Developers also sometimes inspect only the model code and not the SavedModel signature. Use saved_model_cli or inspect loaded.signatures so you know what was actually exported.

Finally, a flexible batch dimension does not guarantee good performance at every batch size. Test the sizes you intend to serve.

Summary

• The best solution is to export the SavedModel with a flexible leading batch dimension.
• Inspect the exported signature before assuming multiple batch sizes are supported.
• If the signature is fixed, re-export the model rather than trying to patch it later.
• Use saved_model_cli or loaded.signatures to verify input shapes.
• Separate compatibility with multiple batch sizes from performance tuning of batch size.