DCGAN 教程

如果您还没有，我们建议您在逐步完成本教程之前，先阅读入门指南。

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

运行原始 DCGAN

请先通过原始教程了解使用原始代码的名人数据集。然后运行bash gan_baseline_run.sh。

启用 DeepSpeed

代码可以从这里获得。

参数解析

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

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，并在 16 个 GPU 上对 DGX-2 节点进行 1 个 epoch 的训练，这导致速度提高了 3 倍。结果摘要如下所示

  • 1 个 epoch 的基线总挂钟时间为 393 秒

  • 1 个 epoch 的 DeepSpeed 总挂钟时间为 128 秒

###

更新：