Pytorch实现cifar-10图像分类

导入必要的包

import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

通过transform 实现对数据进行处理

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

数据加载：

trainset = torchvision.datasets.CIFAR10(root=r'./data',
                                        train=True,download=False,transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size= 4, shuffle=False, num_workers=0)

testset = torchvision.datasets.CIFAR10(root=r'./data',
                                       train=False,download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=0)

classes = ('plane', 'car', 'bird','cat','deer',
           'dog', 'frog', 'horse', 'ship', 'truck')

cifar-10数据可以事先下载好，放到 data文件夹下。 若网速够快，可以在线下载：

设置： download=True 

展示一下我们加载的cifar-10数据：

def imshow(img):
    img = img/2 +0.5    # unnormalize
    npimg =img.numpy()
    plt.imshow(np.transpose(npimg, (1,2,0)))
    plt.show()

dataiter = iter(trainloader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))

print(' '.join('%5s' %classes[labels[j]] for j in range(4)))

定义完了模型，这里只是设置了2层卷积+2层池化+3层全连接的网络结构

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)   # 6*16*5 全连接线性化
        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

设置损失函数和优化器：

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

定义训练轮次，进行训练：

epochs = 5
for epoch in range(epochs):
    running_loss = 0.0
    for i , data in enumerate(trainloader, 0):
        inputs, labels = data
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward +backward +optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 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 # make loss zeros

保存训练好的模型：

# save model
torch.save(net.state_dict(), './save_model') # only save weights

加载模型：

#load model
model = Net()
model.load_state_dict(torch.load('./save_model'))

进行预测：

outputs = model(images)
_, predicted = torch.max(outputs, 1)    # [3, 9, 9, 4]
print('Predicted:', ' '.join('%5s '% classes[labels[j]] for j in range(4)))

对整个测试集进行预测，计算准确率

# test the correct
correct = 0
total = 0
with torch.no_grad():   # 实现一定速度的提升，并节省一半的显存，因为其不需要保存梯度
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted ==labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d%%' %( 100*correct / total))
# Accuracy of the network on the 10000 test images: 60%

使用GPU、或多个GPU进行网络训练时，设置：

# 如何在GPU 上训练
device= torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)   # could look up cuda

model.to(device)
inputs, labels = inputs.to(device), labels.to(device)

# 用 DataParallel 使用多个GPU
# Pytorch 默认只会使用一个GPU，可以通过DataParallel 让你的模型并行运行
if torch.cuda.device_count()>1:
    print("Let's use ", torch.cuda.device_count(), "GPUs" )
    model = nn.DataParallel(model)
model.to(device)