如何在带input_signature的tf.function中传入Python非Tensor参数

核心思路 #

将mode作为Python字符串参数传入@tf.function，同时通过input_signature约束输入张量的通用形状（任意秩、任意维度大小）。这样既保证同一mode下任意形状的张量复用同一计算图，又能灵活切换填充模式，且不会在计算图中生成额外节点。

代码实现 #

import tensorflow as tf

@tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=None)])
def pad_last_dim(x, mode='CONSTANT'):
    # 获取输入张量的秩
    rank = tf.rank(x)
    # 构造padding：前rank-1维填充(0,0)，最后一维填充(1,2)
    padding = tf.concat([
        tf.zeros((rank - 1, 2), dtype=tf.int32),
        tf.constant([[1, 2]], dtype=tf.int32)
    ], axis=0)
    return tf.pad(x, padding, mode=mode)

关键细节说明 #

1. input_signature的作用：tf.TensorSpec(shape=None, dtype=None)表示接受任意形状、任意数据类型的张量，确保同一mode下，不管输入张量是2维、3维还是更高维，都不会触发重追踪，复用已生成的计算图。
2. mode参数的处理：直接传入Python字符串mode，TensorFlow会将其作为tf.pad算子的常量属性，而非计算图中的张量节点，避免额外开销。同时tf.function会根据不同的mode值缓存独立的计算图，每个mode仅追踪一次。
3. 动态构造padding：通过tf.rank获取输入张量的秩，动态生成适配任意秩的padding数组，确保最后一维始终应用(1,2)的填充规则。

验证效果 #

# 测试不同形状的张量，同一mode下无重追踪
x1 = tf.random.normal((3, 4))
x2 = tf.random.normal((2, 3, 5))
x3 = tf.random.normal((1, 2, 3, 4))

# 第一次调用会追踪图，后续同一mode调用复用
out1 = pad_last_dim(x1, mode='SYMMETRIC')
out2 = pad_last_dim(x2, mode='SYMMETRIC')
out3 = pad_last_dim(x3, mode='SYMMETRIC')

# 切换mode会生成新图，但仅追踪一次
out4 = pad_last_dim(x1, mode='REFLECT')

解决原有方案的痛点 #

• 替代硬编码mode：通过参数传入实现灵活切换填充模式。
• 避免字符串张量的额外节点：直接使用Python字符串作为算子属性，无冗余计算图节点。
• 防止被其他tf.function调用时重追踪：tf.function会缓存不同mode对应的图，被其他函数调用时，只要mode匹配就复用已有图，不会重新追踪。

内容的提问来源于stack exchange，提问作者kkm mistrusts SE