使用tf.py_func循环调用lambda函数时出现意外输出及行为差异问题

首先我得明确你遇到的核心问题：在TensorFlow中用tf.py_func结合循环共享变量时，它的执行逻辑和纯Numpy的动态计算完全不一样，这本质上是静态计算图 vs 动态即时执行的差异导致的，再加上tf.py_func本身的特性限制。

先回顾你的纯Numpy代码，逻辑直观且符合预期：

import numpy as np
def my_func(x, k): return np.tile(x,k)
x = np.ones((1), np.int64)
for i in range(1,3):
    x = my_func(x, i)
print(x)  # 输出 [1 1]

每次循环迭代时，i的取值是确定的（第一次1，第二次2），Numpy会立即计算出新的x，整个过程是逐行执行的动态逻辑。

为什么tf.py_func会出问题？ #

假设你写的TensorFlow代码大概是这样的：

import tensorflow as tf
import numpy as np

def my_func(x, k):
    return np.tile(x, k)

x = tf.constant(np.ones((1), np.int64))
for i in range(1,3):
    # 用lambda捕获循环变量i
    x = tf.py_func(lambda x: my_func(x, i), [x], tf.int64)

with tf.Session() as sess:
    print(sess.run(x))  # 输出可能是 [1 1 1 1]，不符合预期

问题出在两个关键点：

1. 静态图的闭包变量捕获：TensorFlow在图构建阶段（循环执行时）不会立即执行tf.py_func里的Python逻辑，只是把操作加入计算图。lambda捕获的是变量i的引用，而非当前迭代的数值。当最终运行sess.run()时，i已经循环到了2（循环结束后的最终值），所以两次tf.py_func调用都会用i=2，导致结果超出预期。
2. tf.py_func的形状丢失：tf.py_func默认不会保留张量的形状信息，这也可能引发后续操作的意外行为，但这不是本次问题的核心。

解决方法 #

方法1：把循环变量转换成TensorFlow张量作为输入

不要用闭包捕获i，而是将i转为常量张量，作为tf.py_func的显式输入，这样每次迭代的i值会被固化到计算图中：

import tensorflow as tf
import numpy as np

def my_func(x, k):
    return np.tile(x, k)

x = tf.constant(np.ones((1), np.int64))
for i in range(1,3):
    # 将i转为张量，作为py_func的第二个输入
    k_tensor = tf.constant(i, dtype=tf.int64)
    x = tf.py_func(my_func, [x, k_tensor], tf.int64)
    # 手动设置形状，避免py_func丢失形状信息
    x.set_shape(tf.TensorShape([None]))

with tf.Session() as sess:
    print(sess.run(x))  # 输出 [1 1]，符合预期

方法2：优先使用TensorFlow原生操作

tf.py_func是兼容Python逻辑的兜底方案，存在无法自动微分、依赖Python运行时等限制。如果逻辑能用原生TF操作实现，一定要优先选择，比如tf.tile：

import tensorflow as tf

x = tf.constant([1], dtype=tf.int64)
for i in range(1,3):
    x = tf.tile(x, [i])

with tf.Session() as sess:
    print(sess.run(x))  # 输出 [1 1]

原生操作完全贴合TensorFlow的静态图逻辑，不会出现行为偏差。

方法3：用TensorFlow 2.x的Eager模式

如果你用的是TF2.x，默认的Eager执行模式和Numpy的动态逻辑一致，无需构建静态图，行为会和预期完全匹配：

import tensorflow as tf
import numpy as np

def my_func(x, k):
    # 把张量转成numpy数组处理
    return np.tile(x.numpy(), k)

x = tf.constant(np.ones((1), np.int64))
for i in range(1,3):
    # TF2.x用tf.py_function替代旧的tf.py_func
    x = tf.py_function(my_func, [x, i], tf.int64)

print(x.numpy())  # 输出 [1 1]

总结 #

tf.py_func在静态图模式下和Numpy的行为差异，核心是静态图构建逻辑与动态即时执行的本质区别。如果必须用Python自定义逻辑，要么把变量显式传入tf.py_func，要么切换到TF2的Eager模式；能不用tf.py_func就尽量用原生操作，既高效又能避免奇怪的行为。

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