input_signature and tf.TensorSpec: Control Tracing in TensorFlow

The Story — The Airport Bag Sizer

In the last lessons you saw that every new shape can cause a new trace. This lesson shows how to control tracing with an explicit signature. It also shows how the same signature becomes the public interface of a saved model.

tf.TensorSpec — A Description of a Tensor

A TensorSpec holds a shape, a dtype and an optional name. It has no values. None in the shape means "any size on this axis".

import tensorflow as tf

spec = tf.TensorSpec(shape=[None, 3], dtype=tf.float32, name="features")
print(spec)

for t in [tf.zeros([5, 3]), tf.zeros([1, 3]), tf.zeros([5, 4]),
          tf.zeros([5, 3], tf.int32), tf.zeros([3])]:
    print(f"{str(t.shape):8s} {t.dtype.name:8s} fits? {spec.is_compatible_with(t)}")

print(tf.TensorSpec.from_tensor(tf.zeros([2, 2])))     # copy a tensor's spec
print(tf.TensorSpec(shape=None, dtype=tf.float32))     # unknown rank

input_signature — One Graph for Every Batch Size

@tf.function(input_signature=[tf.TensorSpec(shape=[None, 3], dtype=tf.float32)])
def row_sums(x):
    print("  tracing with", x.shape)
    return tf.reduce_sum(x, axis=1)

for batch in [1, 8, 64, 5]:
    row_sums(tf.ones([batch, 3]))
print("traces:", row_sums.experimental_get_tracing_count())

Four different batch sizes, one trace. The graph was built for shape (None, 3) and works for all of them.

Anything That Does Not Fit Is Rejected

for bad in [tf.ones([2, 4]), tf.ones([2, 3], tf.int32), 3.0]:
    try:
        row_sums(bad)
    except TypeError as e:
        print("TypeError:", str(e).split("`")[1])

Try It — Signature Matcher

Several Inputs and Nested Inputs

Give one spec per argument, in order. Dicts and tuples work too — the signature mirrors the structure of the arguments.

@tf.function(input_signature=[
    tf.TensorSpec([None, 3], tf.float32, name="x"),
    tf.TensorSpec([None], tf.int32, name="labels"),
])
def batch_info(x, labels):
    return tf.shape(x)[0], tf.reduce_max(labels)

print([int(t) for t in batch_info(tf.ones([4, 3]), tf.constant([2, 0, 1, 2]))])

@tf.function(input_signature=[{
    "img":   tf.TensorSpec([None, 28, 28], tf.float32),
    "label": tf.TensorSpec([None], tf.int64),
}])
def count_images(batch):
    return tf.shape(batch["img"])[0]

print(count_images({"img": tf.zeros([6, 28, 28]), "label": tf.zeros([6], tf.int64)}).numpy())

Where Signatures Really Matter

1. The Last Batch Is Smaller

100 samples in batches of 32 gives shapes 32, 32, 32 and then 4. Without a signature, that last batch triggers a second trace.

ds = tf.data.Dataset.range(100).map(lambda i: tf.fill([3], tf.cast(i, tf.float32))).batch(32)

@tf.function
def total_plain(x):
    return tf.reduce_sum(x)

@tf.function(input_signature=[tf.TensorSpec([None, 3], tf.float32)])
def total_sig(x):
    return tf.reduce_sum(x)

for x in ds:
    total_plain(x); total_sig(x)

print("batch shapes :", [tuple(x.shape) for x in ds])
print("traces plain :", total_plain.experimental_get_tracing_count())
print("traces sig   :", total_sig.experimental_get_tracing_count())

2. Saving a Model for Serving

When you save a model, the signature becomes its API. Servers, phones and other languages call the model through it, with no Python code.

class Linear(tf.Module):
    def __init__(self):
        self.w = tf.Variable([[1.0], [2.0], [3.0]])

    @tf.function(input_signature=[tf.TensorSpec([None, 3], tf.float32, name="x")])
    def predict(self, x):
        return x @ self.w

tf.saved_model.save(Linear(), "/tmp/linear_model")
loaded = tf.saved_model.load("/tmp/linear_model")

print(loaded.predict(tf.ones([2, 3])).numpy().tolist())
print("signatures:", list(loaded.signatures.keys()))

Golden Rules