python - 如何更新 tensorflow 2.0 中的镜像变量？

Question

我正在 tensorflow 2.0 版中构建模型(由于与我的 cuda 版本兼容，因此无法升级，我无权更改)。我正在使用 tf.strategy.MirroredStrategy()在 2 个 GPU 上训练我的模型。但是，我正在尝试实例化一个自定义密集层，其权重是不同密集层权重的转置。我的代码涉及这一行来构建自定义层:

from tensorflow.keras import backend as K

class DenseTied(Layer):
    # Really long class, full code can be found at link below

    def build(self, input_shape):
        self.kernel = K.transpose(self.tied_to.kernel)

然后我在模型中使用它，如下所示:

from tensorflow.keras.layers import Input, Dense

def build_model(input_shape):
    model_input = Input(shape=input_shape)
    dense1 = Dense(6144, activation='relu')
    dense_tied1 = DenseTied(49152, tied_to=dense1)

    x = dense1(model_input)
    model_output = dense_tied1(x)

    model = Model(model_input, model_output)

    model.compile(optimizer='adam', loss='mse')

    return model

尝试构建此模型时出现错误: AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute '_distribute_strategy' .
我已经追踪了一段时间，我指出问题出在

self.kernel = K.transpose(self.tied_to.kernel)

看来 self.tied_to.kernel是 <class 'tensorflow.python.distribute.values.MirroredVariable'> 类型但在调用 K.transpose() 之后在其上，结果输出的类型为 <class 'tensorflow.python.framework.ops.EagerTensor'> .我尝试按照说明 here但它没有用。我得到 AttributeError: 'MirroredStrategy' object has no attribute 'run'当它在文档中时。所以我想也许我的 Tensorflow 版本对于那种方法来说太旧了。
如何在 Tensorflow 2.0 中更新镜像变量？
此外，如果您想查看完整的自定义层代码，我正在尝试实现所描述的密集绑定(bind)层 here .

Answer 1

截至目前，该文档适用于 tensorflow 2.3。如果您使用的是 2.0，则应该是strategy.experimental_run_v2而不是 strategy.run .