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] # 丢弃一半的行和列

这个 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) # 正常工作。跟踪函数。 preprocessed_images = shrink(img_batch_2) # 正常工作。使用相同的具体函数。 </