我有一个可以在单个GPU上完美运行的模型，如下所示：

alpha = tf.Variable(alpha,name='ws_alpha',trainable=False,dtype=tf.float32,aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,)

...
class CustomModel(tf.keras.Model):


    @tf.function
    def train_step(inputs):
        ...
        alpha.assign_add(increment)

...


model.fit(dataset,epochs=10)

但是，当我在多个GPU上运行时，分配工作尚未完成。它适用于两个训练步骤，然后在整个时期保持不变。

alpha是两层的加权和，例如out = a*Layer1 + (1-a)*Layer2。它不是可训练的参数，但类似于step_count变量。

有人在Tensorflow 2的多GPU设置中分配单个值的经验吗？

最好将变量分配为：

with tf.device("CPU:0"):
    alpha = tf.Variable()

？