变量.使用Tensorflow 2在多GPU上分配(值)

Question

我有一个在单个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上运行时，任务并没有完成。它适用于两个训练步骤，然后在整个时期保持不变。

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

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

将变量分配为会更好吗

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

Answer 1

根据tensorflow问题进行简单修复

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