DCGAN 教程

如果您还没有,我们建议您先阅读 入门 指南,然后再学习本教程。

在本教程中,我们将使用自定义(用户定义)优化器和多引擎设置将 DCGAN 模型移植到 DeepSpeed!

运行原始 DCGAN

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

启用 DeepSpeed

代码可以在 此处 获取。

参数解析

将 DeepSpeed 应用于 DCGAN 模型的第一步是向 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 个 epoch 的训练,这导致速度提高了 3 倍。结果摘要如下所示

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

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

###

更新: