Pytorch多GPU训练过程

import torch#深度学习的pytoch平台
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5' #指定GPU编号
device = torch.device("cuda") #创建GPU对象

x_train = Variable(train,requires_grad=True).to(device=device,dtype=torch.float32) #把训练变量放到GPU

net = DNN(layers)
net = nn.DataParallel(net)
net.to(device=device)

net = nn.DataParallel(net)
net.to(device=device)

 if torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)

model = Model(input_size, output_size)  # 实例化模型对象
if torch.cuda.device_count() > 1:  # 检查电脑是否有多块GPU
    print(f"Let's use {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)  # 将模型对象转变为多GPU并行运算的模型

model.to(device)  # 把并行的模型移动到GPU上