75、TensorFlow函数与图的深入解析

TensorFlow函数与图的深入解析

1. 限制TF函数的输入签名

在某些情况下，你可能希望将TF函数限制为特定的输入签名。例如，假设你知道只会使用28×28像素的图像批次来调用TF函数，但这些批次的大小可能会有很大差异。你可能不希望TensorFlow为每个批次大小生成不同的具体函数，或者依赖它自行确定何时使用 None 。在这种情况下，你可以像这样指定输入签名：

import tensorflow as tf

@tf.function(input_signature=[tf.TensorSpec([None, 28, 28], tf.float32)])
def shrink(images):
    return images[:, ::2, ::2] # drop half the rows and columns

这个TF函数将接受任何形状为 [*, 28, 28] 的float32张量，并且每次都会重用相同的具体函数：

img_batch_1 = tf.random.uniform(shape=[100, 28, 28])
img_batch_2 = tf.random.uniform(shape=[50, 28, 28])
preprocessed_images = shrink(img_batch_1) # Works fine. Traces the function.
preprocessed_images = shrink(img_ba