如何在TensorFlow中禁用tf.function的函数重跟踪（retracing）

针对性代码调整 #

代码里的重跟踪触发源有两个：循环内每轮新建tf.Variable对象、GradientTape.jacobian()默认pfor路径的内部tf.function封装。对应修改如下：

1. 替换循环内的Variable新建逻辑 #

每轮迭代不要重新实例化tf.Variable，改为对初始变量做原地赋值，切断新建对象触发的重跟踪链路：

# 替换原循环内的新建Variable代码
# 原：var = tf.Variable([var_next[0, 0], var_next[0, 1], var_next[0, 2], var_next[0, 3]])
# 改：
var.assign(tf.reshape(var_next, [4]))

2. 关闭jacobian计算的pfor自动图封装 #

警告提示的pfor..f函数是jacobian()默认并行路径下的内部图封装逻辑，给计算参数调整即可关闭自动追踪：

# 原：h = tape2.jacobian(g, var)
# 改：
h = tape2.jacobian(g, var, experimental_use_pfor=False)

全局禁用重跟踪的配置 #

如果需要彻底关闭这类自动重跟踪行为，可按需选择配置：

• 强制全局eager模式运行，完全关闭自动图追踪、重跟踪流程，适配小规模自定义优化循环的运行场景，在脚本最开头加入：

  tf.config.run_functions_eagerly(True)
  

• 仅过滤日志层面的重跟踪警告，不改动原有运行逻辑：

  import logging
  logging.getLogger("tensorflow").setLevel(logging.ERROR)
  

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