如何在 TensorFlow 2.0 中为多个 trainable_weights 运行 apply_gradients

问题描述

我已经用 tensorflow2.0 构建了一个类似深度学习模型的编码器-解码器，出于某种原因，编码器和解码器是两个独立的模型，我想使用 optimizer.apply_gradients() 函数优化这两个模型. TensorFlow document 没有为此案例提供太多信息。它表明grads_and_vars应该是“（梯度，变量）对列表”。我尝试了两种方法，如下所示：

第一种方法： optimizer.apply_gradients(zip(grads,self._encoder.trainable_weights,self._decoder.trainable_weights))

方法二： optimizer.apply_gradients([zip(grads,self._encoder.trainable_weights),zip(grads,self._decoder.trainable_weights)])

两者都不起作用。正确的做法是什么？

解决方法

您可以尝试类似以下代码的训练循环，引用自 DCGAN tutorial。

with tf.GradientTape() as enc_tape,tf.GradientTape() as dec_tape:
    feature_space = encoder(images,training=True)
    recovered_image = decoder(feature_space,training=True)

    #Some custom loss function...
    loss = tf.keras.losses.MSE(recovered_image,images)

    gradients_of_encoder = enc_tape.gradient(loss,encoder.trainable_variables)
    gradients_of_decoder = dec_tape.gradient(loss,decoder.trainable_variables)

enc_optimizer.apply_gradients(zip(gradients_of_encoder,encoder.trainable_variables))
dec_optimizer.apply_gradients(zip(gradients_of_decoder,decoder.trainable_variables))