Pytorch单机多卡训练

import torch.nn

#原来单卡运行的GPU代码只需要进行少量改掉即可

#现在我有两块GPU,序号是0,1,然后导入以下代码
os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1"
device_ids = [0, 1]

#原来定义的model传入nn.DataParallel
model = nn.DataParallel(model,device_ids).cuda()

#之后model原来的参数就会全都转移到model.module这个模块下
"""
    model.G.train()#原来单卡就可以进行的操作
    model.module.G.train()#现在要改为这个代码,其他参数也是同理
"""



其他地方不用改动。

Pytorch单机多卡训练_第1张图片

运行代码之后,可以看到2个GPU都被使用了。

原来我得batchsize只能给8,跑一个epoch要90分钟以上

现在batchsize能给到14,跑一个epoch在60分钟左右。

但这是DP(DataParallel)模式,听说DDP模式下进行多卡训练,会更快。后边研究研究

你可能感兴趣的:(深度学习,深度学习,人工智能,python,并行训练,pytorch)