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())