如何最好地将@tf.function 装饰器用于类方法？答案

【问题标题】：How to best use @tf.function decorator for class methods?如何最好地将@tf.function 装饰器用于类方法？
【发布时间】：2021-01-15 11:30:18
【问题描述】：

用一个最小的例子来说明我的问题，假设我想创建一个类

class test_a:
    def __init__(self, X):
        self.X = X
    def predict(self, a):
        return a * self.X

重要的是，如果我将新的X 分配给test_a 的实例，predict() 函数应该会改变。在这个例子中它工作正常：

X = tf.ones((1, 1))
a = test_a(X)
y = tf.ones((1, 1))
a.predict(y) # output [[1.]]

# now I want to change the value of a.X
Xnew = 2 * tf.ones((1, 1))
a.X = Xnew
a.predict(y) # output [[2.]], as desired.

现在假设我想使用 @tf.function 装饰器来加速 predict()。

class test_b:
    def __init__(self, X):
        self.X = X
        
    @tf.function
    def predict(self, a):
        return a * self.X

现在出现以下不良行为：

X = tf.ones((1, 1))
b = test_b(X)
y = tf.ones((1, 1))
b.predict(y) # output [[1.]]

# now I want to change the value of b.X
Xnew = 2 * tf.ones((1, 1))
b.X = Xnew
b.predict(y) # output is still [[1.]], but I would like it to be [[2.]]

到目前为止，我唯一的想法是拥有一个方法_predict(X, a)，然后我可以对其进行装饰，然后在（未装饰的）方法predict(self, a) 中调用_predict(self.X, a)。任何帮助如何更好地做到这一点将不胜感激。

【问题讨论】：

我认为self.X 应该是tf.Variable(trainable=False)。你必须使用tf.Variable.assign() 来改变它。避免python副作用
非常感谢，此更改使代码按预期工作！但是，我注意到我收到以下警告：""" WARNING:tensorflow:7 out of the last 11 calls to 触发了 tf.function retracing。跟踪很昂贵，并且过多的跟踪可能是由于（1）在循环中重复创建@tf.function，（2）传递不同形状的张量，（3）传递Python对象而不是张量。“””这是否意味着每次我分配时都会创建一个新图self.X 的新价值？会不会有问题？
每次调用具有新形状的模型时都会构建新图形
非常感谢您的解释，这对我有帮助。这个建议引起的一个问题是我希望 X 的形状有可能发生变化。使用 self.X.assign(...) 当形状不匹配时出现错误。如何使用 assign 更改 X 的形状？
我认为 tf.function 不可能。我什至不确定是否可以在急切模式下进行

标签： python python-3.x tensorflow tensorflow2.0

【解决方案1】：

为了避免 python 副作用 - self.X 应该是 tf.Variable(trainable=False)。而且你必须使用tf.Variable.assign()来改变它。

【讨论】：