pytorch之GPU数据并行

使模型在gpu上运行

在原来的代码上修改了两处，如代码标注所示

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)#1.网络参数数据要是GPU格式
for epoch in range(2):  # loop over the dataset multiple times

running_loss = 0.0

for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
#2.输入数据以及标签数据都要是GPU格式
inputs=torch.FloatTensor(inputs)
inputs=inputs.to(device)
#labels=torch.FloatTensor(float(labels))
labels=labels.to(device)
# zero the parameter gradients,把梯度置零，也就是把loss关于weight的导数变成0.
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)#forward前向传播求出预测的值
loss = criterion(outputs, labels)#求los
loss.backward()#backward反向传播求梯度
optimizer.step()#optimize更新所有参数

# print statistics
running_loss += loss.item()
if i % 2000 == 1999:    # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0

print('Finished Training')

并行运行DataParallel

model = Model(input_size, output_size)#Model是定义的网络结构
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)#如果有多个GPU，用DataParallel进行GPU的并行计算
model.to(device)#model放入GPU
for data in rand_loader:
input = data.to(device)#输入数据放入GPU
output = model(input)
print("Outside: input size", input.size(),
"output_size", output.size())