如何使用 Tensorflow 2/Keras 保存和恢复训练具有多个模型部分的 GAN

Question

我目前正在尝试添加一个功能来中断和恢复对从这个示例代码创建的 GAN 的训练：https://machinelearningmastery.com/how-to-develop-an-auxiliary-classifier-gan-ac-gan-from-scratch-with-keras/

我设法让它以一种方式工作，我将整个复合 GAN 的权重保存在 summarize_performance 函数中，该函数每 10 个 epoch 触发一次，如下所示：

# save all weights
filename3 = 'weights_%08d.h5' % (step+1)
gan_model.save_weights(filename3)
print('>Saved: %s and %s and %s' % (filename1, filename2, filename3))

它加载在我添加到名为 load_model 的程序开头的函数中，它采用正常构建的 gan 架构，但将其权重更新为最新值，如下所示：

#load model from file and return startBatch number
def load_model(gan_model):
   start_batch = 0
   files = glob.glob("./weights_0*.h5")
   if(len(files) > 0 ):
       most_recent_file = files[len(files)-1]
       gan_model.load_weights(most_recent_file)
       #TODO: breaks if using more than 8 digits for batches
       startBatch = int(most_recent_file[10:18])
       if (start_batch != 0):
           print("> found existing weights; starting at batch %d" % start_batch)
   return start_batch

其中 start_batch 被传递到训练函数以跳过已经完成的时期。

虽然这种减重方法确实“有效”，但我仍然认为我这里的方法是错误的，因为我发现重量数据显然不包括 GAN 的优化器状态，因此训练不会继续就像它没有被打断一样。

我发现保存进度同时保存优化器状态的方法显然是通过保存整个模型而不仅仅是权重来完成的

这里我运行遇到了一个问题，因为在 GAN 中我不只有一个模型可以训练，但我有 3 个模型：

生成器模型g_model
判别器模型d_model
和复合GAN模型gan_model

它们都是相互联系、相互依赖的。如果我采用天真的方法并分别保存和恢复这些零件模型中的每一个，我最终将拥有 3 个独立的脱节模型而不是 GAN

有没有办法保存和恢复整个 GAN，让我恢复训练，就好像没有发生中断一样？

Answer 1

如果您想恢复整个 GAN，可以考虑使用 tf.train.Checkpoint：

### In your training loop

checkpoint_dir = '/checkpoints'
checkpoint = tf.train.Checkpoint(gan_optimizer=gan_optimizer,
                            discriminator_optimizer=discriminator_optimizer,
                                  generator=generator,
                                  discriminator=discriminator
                                  gan_model = gan_model)
  
ckpt_manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)
if ckpt_manager.latest_checkpoint:
    checkpoint.restore(ckpt_manager.latest_checkpoint)  
    print ('Latest checkpoint restored!!')

....
....


if (epoch + 1) % 40 == 0:
    ckpt_save_path = ckpt_manager.save()
    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,ckpt_save_path))

### After x number of epochs, just save your generator model for inference.

generator.save('your_model.h5')

你也可以考虑完全去掉复合模型。 Here 是我的意思的一个例子。

如何使用 Tensorflow 2/Keras 保存和恢复训练具有多个模型部分的 GAN

How to save and resume training a GAN with multiple model parts with Tensorflow 2/ Keras

python

deep-learning

keras

tensorflow

generative-adversarial-network