pytorch 数据并行实验

实验数据集为CIFAR10数据集，分别在CPU、单GPU、四块GPU下进行测试。
软件版本：
python：3.5
torch：1.0
torchvision：0.2.0
cuda：9.0

1. CPU情况下：

import torch
from torchvision import datasets,transforms
from torch.autograd import Variable
import torchvision
import time
import torch.nn as nn
import torch.nn.functional as F

transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])

data_train=datasets.CIFAR10(root="./data/",transform=transform,train=True,download=True)
data_test=datasets.CIFAR10(root="./data/",transform=transform,train=False,download=True)

data_loader_train=torch.utils.data.DataLoader(dataset=data_train,batch_size=64,shuffle=True)
data_loader_test=torch.utils.data.DataLoader(dataset=data_test,batch_size=64,shuffle=True)



import numpy as np




# get some random training images
dataiter = iter(data_loader_train)
images, labels = dataiter.next()



class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

n_epochs=50
time_open=time.time()
for epoch in range(n_epochs):  # loop over the dataset multiple times

    running_loss = 0.0
    running_correct=0
    print("Epoch {}/{}".format(epoch,n_epochs))
    print("-"*10)
    for data in data_loader_train:
        # get the inputs
        x_train, y_train = data
        x_train, y_train = Variable(x_train),Variable(y_train)
        outputs = net(x_train)
        _,pred=torch.max(outputs.data,1)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize