如何在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'>类型。我试着按照指示去做，但没有成功。我在文档中得到AttributeError: 'MirroredStrategy' object has no attribute 'run'。所以我想也许我的Tensorflow版本对于这个方法来说太旧了

如何在Tensorflow 2.0中更新镜像变量

另外，如果您想查看完整的自定义层代码，我将尝试实现here描述的密集绑定层

Answer 1

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