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 秒
###