DCGAN教程

如果您还没有这样做，我们建议您首先阅读入门指南，然后再进行本教程的学习。

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

运行原始DCGAN

请先使用原始代码浏览原始教程中的名人数据集。然后运行bash gan_baseline_run.sh。

启用DeepSpeed

代码可以在这里获取 here。

参数解析

应用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，并在DGX-2节点上的16个GPU上进行1个周期的训练，这导致了3倍的速度提升。结果总结如下：

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

  • Deepspeed 一个周期的总挂钟时间为128秒

###

更新: