使用 tensorflow 签名和 tf.function 更新外部状态答案

【问题标题】：Update external state using tensorflow autograph and tf.function使用 tensorflow 签名和 tf.function 更新外部状态
【发布时间】：2019-09-12 16:58:30
【问题描述】：

我正在创建一个层，它是批量标准化的变体，并使用 tf.function 装饰器来加速它。但是，我收到一条错误消息，提示 autograph 不知道我是否要重用我正在尝试更新的变量。

class MyClass(tf.keras.layers.Layer):

    def build():
        self.foo = self.add_weight(...)

    @tf.function
    def call(inputs, training=None):
        lst = [K.moving_average_update(self.foo, .5, .999)]
        self.add_updates(lst)

这会引发类似于以下内容的错误：

ValueError: Variable my_class/foo/biased already exists, disallowed. Did you mean to set reuse=True or reuse=tf.AUTO_REUSE in VarScope?

告诉 autograph 在函数外重用状态的惯用方法是什么？

【问题讨论】：

您能否尝试将参数getter = lambda **kwargs: tf.get_variable('var_name', reuse=tf.AUTO_REUSE, **kwargs) 传递给self.add_weight？
这只是变量名，所以他们都应该这样做
@rvinas 我试了一下，但 get_variable 在 2.0 中已弃用
好吧 - 你可以通过 tf.Variable 代替吗？例如，创建一个带有参数**kwargs 的函数，返回一个在块with tf.variable_scope('scope', reuse=tf.AUTO_REUSE) 内创建的tf.Variable，然后将该函数传递给self.add_weight 的参数getter。

标签： python tensorflow tensorflow2.0

【解决方案1】：

为了重用变量self.foo，您可以在 tf.variable_scope 块中定义，设置参数reuse = tf.AUTO_REUSE。例如：

def build(self, input_shape):

    with tf.variable_scope('scope', reuse=tf.AUTO_REUSE):
        self.foo = tf.Variable(initial_value=YOUR_INITIAL_VALUE, name='var')

    trainable = None  # True or False
    if trainable:
        self._trainable_weights.append(self.foo)
    else:
        self._non_trainable_weights.append(self.foo)

注意：未经测试。

【讨论】：

很遗憾，2.0中没有variable_scope，也没有tf.AUTO_REUSE常量。
您能否尝试将tf.compat.v1 附加到这些函数中？（例如tf.compat.v1.AUTO_REUSE）？（见migration plan）
其实。它现在可以正常工作，无需任何修改。一旦我遇到了 2.0 的稳定版本，就可以工作。