DCGAN 教程

如果您尚未阅读，我们建议您在学习本教程之前，先阅读入门指南。

在本教程中，我们将使用自定义（用户定义）优化器和多引擎设置，将 DCGAN 模型移植到 DeepSpeed！

运行原始 DCGAN

请先使用原始代码学习 Celebrities 数据集的原始教程。然后运行 bash gan_baseline_run.sh。

启用 DeepSpeed

代码可从此处获取。

参数解析

应用 DeepSpeed 的第一步是使用 deepspeed.add_config_arguments() 函数为 DCGAN 模型添加配置参数，如下所示。

import deepspeed

def main():
    parser = get_argument_parser()
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    train(args)

初始化

我们使用 deepspeed.initialize 创建两个模型引擎（一个用于判别器网络，一个用于生成器网络，以及它们各自的优化器），如下所示。

    model_engineD, optimizerD, _, _ = deepspeed.initialize(args=args, model=netD, model_parameters=netD.parameters(), optimizer=optimizerD)
    model_engineG, optimizerG, _, _ = deepspeed.initialize(args=args, model=netG, model_parameters=netG.parameters(), optimizer=optimizerG)

请注意，DeepSpeed 会自动处理分布式训练方面，因此我们将 ngpu=0 设置为禁用 PyTorch 的默认数据并行模式。

判别器训练

我们对判别器的反向传播进行如下修改：

model_engineD.backward(errD_real)
model_engineD.backward(errD_fake)

这使得在优化器更新中包含了来自真实和虚假小批量的梯度。

生成器训练

我们对生成器的反向传播进行如下修改：

model_engineG.backward(errG)

注意：在使用梯度累积的情况下，生成器上的反向传播会导致梯度在判别器上累积，这是因为 errG 是通过判别器前向传播计算出来的，存在张量依赖关系；因此，在进行生成器反向传播之前，请将 netD 参数的 requires_grad 设置为 False。

配置

使用 DeepSpeed 的下一步是创建一个配置文件 JSON 文件（gan_deepspeed_config.json）。此文件提供了用户定义的 DeepSpeed 特定参数，例如批处理大小、优化器、调度器和其他参数。

{
  "train_batch_size" : 64,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.0002,
      "betas": [
        0.5,
        0.999
      ],
      "eps": 1e-8
    }
  },
  "steps_per_print" : 10
}

运行启用 DeepSpeed 的 DCGAN 模型

要开始使用 DeepSpeed 训练 DCGAN 模型，我们执行以下命令，该命令默认将使用所有检测到的 GPU。

deepspeed gan_deepspeed_train.py --dataset celeba --cuda --deepspeed_config gan_deepspeed_config.json --tensorboard_path './runs/deepspeed'

性能比较

我们使用总批处理大小为 64，并在 DGX-2 节点上的 16 个 GPU 上进行 1 个 epoch 的训练，这带来了 3 倍的加速。结果总结如下：

  • 基线模型在 1 个 epoch 的总运行时间为 393 秒

  • DeepSpeed 在 1 个 epoch 的总运行时间为 128 秒

###

更新：