python - 如何使用 @tf.function 进入循环和可变变量 tensorflow 2.0

Question

我有两个循环，在第二个循环中我有一个@tf.function，我需要更新 mutable_variable 的值 T2 次，而不是将值返回 0 并重复 T1 次。

我的问题是代码很慢。甚至 mutable_variable 中的操作也很简单。

for x in range T1:
 mutable_variable=0  
for i in range T2:

   somefunction(mutable_variable)
   print(mutable_variable)


@tf.function
def somefunction(mutable_variable)
   mutable_variable=mutable_variable+update_value_of(mutable_variable)

我收到此警告：我该如何解决？

警告：tensorflow：在最后 11 次调用中，有 11 次调用 <function main..train_step at 0x7f1f503b12f0> 触发了 tf.function 回溯。跟踪是昂贵的，过多的跟踪可能是由于（1）在循环中重复创建 @tf.function，（2）传递不同形状的张量，（3）传递 Python 对象而不是张量。对于 (1)，请在循环之外定义您的 @tf.function。对于 (2)，@tf.function 具有 Experimental_relax_shapes=True 选项，可以放宽可以避免不必要的回溯的参数形状。(3)请参考https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args和https://www.tensorflow.org/api_docs/python/tf/function更多细节。更新记录：12 of 3912021-07-21 21:40:59.243953: W tensorflow/core/grappler/optimizers/loop_optimizer.cc:906] 使用控制输入跳过合并节点的循环优化：rot90/Assert/AssertGuard/branch_executed/_8

标签： pythontensorflowtensorflow2.0