pytorch用gpu单机多卡并行训练

在代码的最顶上加上这个，值写你想要并行训练的gpu编号。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

将model传入DataParallel方法，并使用.cuda()

model = torch.nn.DataParallel(model)
model = model.cuda()

把数据样本都用to传入到gpu中。

device = torch.device('cuda')
sequences = sequences.to(device)
labels = labels.to(device)

然后就ok啦

标签：torch,多卡,DataParallel,pytorch,cuda,device,gpu,model
来源： https://blog.csdn.net/weixin_40982849/article/details/121118077