假设我有2个GPU，并且我想使用tf.data.Dataset API在急切的执行设置中在同一数据集上同时在这两个GPU上训练两个模型。我想知道如何将相同的数据批量同时输入两个模型，并在两个GPU上并行训练它们。
我正在使用tf-1.13。 （也欢迎使用tf-2.0的答案）。
我尝试了以下方法，但是模型是一个接一个地分批训练的，但不能同时进行。 注意：我已经使用tf.keras定义了模型
训练单个模型
@tf.function # in case of tf-2.0
def train_step(inputs, labels): # (inputs,labels) batch
with tf.GradientTape() as g:
outputs = model(inputs, training=True)
loss = loss_function(labels, outputs)
gradients = g.gradient(loss, modelvariables)
optimizer.apply_gradients(zip(gradients, model.variables))
训练两个模型
@tf.function # in case of tf-2.0
def train_step(inputs, labels): # all inputs in batches
with tf.GradientTape() as g1, tf.GradientTape() as g2:
outputs1 = model1(inputs, training=True)
outputs2 = model2(inputs, training=True)
loss1 = loss_function1(labels, outputs1)
loss2 = loss_function2(labels, outputs2)
gradients1 = g1.gradient(loss1, model1.variables)
gradients2 = g2.gradient(loss2, model2.variables)
optimizer1.apply_gradients(zip(gradients1, model1.variables))
optimizer2.apply_gradients(zip(gradients2, model2.variables))
# train_step = tf.contrib.eager.defun(train_step) # tf-1.13
编辑：我知道我们可以使用@tf.function装饰器将函数编译为tf图，以提高性能。但是，如何将不同的GPU（或GPU的一部分）分配给使用相同数据集的不同模型。