在具有多个 Keras 模型的 TF2 自定义训练循环中应用梯度的正确方法

Question

我正在使用涉及多个 Keras 模型的 GradientTape 实现自定义训练循环。我有 3 个网络，model_a、model_b 和 model_c。我创建了一个列表来保存他们的 trainbale_weights 为：

trainables = list() 
trainables.append(model_a.trainable_weights) # CovNet 
trainables.append(model_b.trainable_weights) # CovNet 
trainables.append(model_c.trainable_weights) # Fully Connected Network

然后我计算损失并尝试将梯度应用为：

loss = 0.
optimizer = tf.keras.optimizers.Adam()
for _, (x, y) in enumerate(train_dataset):
   with tf.GradientTape() as tape:
     y = ...
     loss = ... # custom loss function!
gradients = tape.gradient(loss, trainables)
optimizer.apply_gradients(zip(gradients, trainables))

但我收到以下错误，我不确定错误在哪里：

AttributeError: 'list' object has no attribute '_in_graph_mode'

如果我遍历梯度和可训练对象，然后应用梯度，它会起作用，但我不确定这是否是正确的方法。

for i in range(len(gradients)):
   optimizer.apply_gradients(zip(gradients[i], trainables[i]))

Answer 1

问题是 tape.gradient 期望 trainables 是可训练变量的平面列表而不是列表列表。您可以通过将所有可训练权重连接到一个平面列表中来解决此问题：

trainables = model_a.trainable_weights + model_b.trainable_weights + model_c.trainable_weights

在具有多个 Keras 模型的 TF2 自定义训练循环中应用梯度的正确方法

Correct way to apply gradients in TF2 custom training loop with multiple Keras models

keras

tensorflow

python-3.6

tensorflow2.0