带有GradientTape的健身和自定义训练功能中的GPU使用率非常不同

Question

当我使用Tensorflow keras中的数据生成器函数使用model.fit训练模型时，我发现GPU利用率可能接近100％。这是因为我训练的模型在并行化方面是一个很好的模型。

尽管如此，如果我构建了一个火车功能（类似于以下功能），我注意到利用率却大不相同：

@tf.function
def train_step(input, output):
  with tf.GradientTape() as tape:
    predictions = model(input)
    loss = loss_object(output, predictions)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  train_loss(loss)

详细地说，我注意到它从0％下降到几乎100％，然后在两次更新之间再次下降到0％。我猜在两次更新之间，首先需要将数据从CPU传输到GPU。这解释了为什么它会从100％变为0％的原因。

但这给我带来了一个问题：如何使用tf.GradientTape使它更适合自定义训练功能，以使GPU几乎一直都是100％（与model.fit一样）？