对于需要多次调用的步骤,我们可以使用装饰器 @tf.function 来缓存这些步骤,从而避免不必要的重建。
示例代码:
import tensorflow as tf
@tf.function
def my_func(x):
print("Running computation")
return x * x
x = tf.constant(2)
y = tf.constant(3)
res1 = my_func(x) # Running computation
res2 = my_func(x) # No computation
res3 = my_func(y) # Running computation
运行结果:
Running computation