TensorFlow py_func是否支持多参数？若支持请提供传参示例

当然支持！在TensorFlow 1.x中，tf.py_func完全可以向你指定的自定义Python函数传递多个参数。甚至在TensorFlow 2.x里，虽然官方更推荐使用tf.numpy_function替代py_func，但两者的多参数传递逻辑是一致的。

基础示例（TensorFlow 1.x） #

先定义一个接受多个参数的普通Python函数，再用py_func把它包装成TensorFlow操作：

import tensorflow as tf

# 自定义Python函数：接受两个张量对应的numpy数组，返回计算结果
def custom_add(a, b):
    # 这里可以写任意Python逻辑，比如计算两数之和再加5
    return a + b + 5

# 创建两个输入张量
tensor_a = tf.constant([1, 2, 3], dtype=tf.int32)
tensor_b = tf.constant([4, 5, 6], dtype=tf.int32)

# 使用py_func传递多参数
# inp参数是一个张量列表，列表元素会按顺序传入custom_add函数
result = tf.py_func(
    func=custom_add,
    inp=[tensor_a, tensor_b],
    Tout=tf.int32  # 指定输出张量的数据类型
)

# 运行会话计算结果
with tf.Session() as sess:
    print(sess.run(result))  # 输出: [10 12 14]

TensorFlow 2.x 替代方案（tf.numpy_function） #

如果你用的是TF2.x，建议优先用tf.numpy_function，用法几乎一样：

import tensorflow as tf

def custom_add(a, b):
    return a + b + 5

tensor_a = tf.constant([1, 2, 3], dtype=tf.int32)
tensor_b = tf.constant([4, 5, 6], dtype=tf.int32)

result = tf.numpy_function(
    func=custom_add,
    inp=[tensor_a, tensor_b],
    Tout=tf.int32
)

# TF2.x无需会话，直接转换为numpy数组查看结果
print(result.numpy())  # 输出: [10 12 14]

关键说明 #

• inp参数是张量列表，列表中的每个元素会按顺序作为参数传递给自定义函数。比如inp=[x, y, z]就会调用func(x, y, z)。
• 自定义函数的输入会自动转换为numpy数组（而不是TensorFlow张量），函数返回的numpy数组也会被自动转换为TensorFlow张量。
• 必须通过Tout指定输出张量的类型，如果有多个输出，Tout可以是类型列表（比如Tout=[tf.int32, tf.float32]）。

内容的提问来源于stack exchange，提问作者Kousuke Kuzuoka