Pytorch小試牛刀：CIFAR-10分類

下面我們來嘗試實現對CIFAR-10數據集的分類，步驟如下:

1.使用torchvision加載並預處理CIFAR-10數據集
2.定義網絡
3.定義損失函數和優化器
4.訓練網絡並更新網絡參數
5.測試網絡

CIFAR-10數據加載及預處理

CIFAR-101是一個常用的彩色圖片數據集，它有10個類別: 'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'。每張圖片都是33232，也即3-通道彩色圖片，分辨率32*32。

import torchvision as tv
import torchvision.transforms as transforms
from torchvision.transforms import ToPILImage
import torch as t

show = ToPILImage()  # 可以把Tensor轉成Image，方便可視化
# 第一次運行程序torchvision會自動下載CIFAR-10數據集，
# 大約100M，需花費一定的時間，
# 如果已經下載有CIFAR-10，可通過root參數指定

# 定義對數據的預處理
transform = transforms.Compose([
        transforms.ToTensor(), # 轉為Tensor
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), # 歸一化
                             ])

# 訓練集
trainset = tv.datasets.CIFAR10(
                    root='D:\\dtset',
                    train=True,
                    download=True,
                    transform=transform)

trainloader = t.utils.data.DataLoader(
                    trainset,
                    batch_size=4,
                    shuffle=True,
                    num_workers=2)

# 測試集
testset = tv.datasets.CIFAR10(
                    'D:\\dtset',
                    train=False,
                    download=True,
                    transform=transform)

testloader = t.utils.data.DataLoader(
                    testset,
                    batch_size=4,
                    shuffle=False,
                    num_workers=2)

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

下載地址： https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
Dataset對象是一個數據集，可以按下標訪問，返回形如(data, label)的數據。

(data, label) = trainset[100]
print(classes[label])

# (data + 1) / 2是為了還原被歸一化的數據
show((data + 1) / 2).resize((100, 100))

Dataloader是一個可迭代的對象，它將dataset返回的每一條數據拼接成一個batch，並提供多線程加速優化和數據打亂等操作。當程序對dataset的所有數據遍歷完一遍之后，相應的對Dataloader也完成了一次迭代。

dataiter = iter(trainloader)
images, labels = dataiter.next() # 返回4張圖片及標簽
print(' '.join('%11s'%classes[labels[j]] for j in range(4)))
show(tv.utils.make_grid((images+1)/2)).resize((400,100))

定義網絡

拷貝上面的LeNet網絡，修改self.conv1第一個參數為3通道，因CIFAR-10是3通道彩圖。

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

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5) 
        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 = F.max_pool2d(F.relu(self.conv1(x)), (2, 2)) 
        x = F.max_pool2d(F.relu(self.conv2(x)), 2) 
        x = x.view(x.size()[0], -1) 
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)        
        return x


net = Net()
print(net)

定義損失函數和優化器(loss和optimizer)

from torch import optim
criterion = nn.CrossEntropyLoss() # 交叉熵損失函數
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

訓練網絡

所有網絡的訓練流程都是類似的，不斷地執行如下流程：

• 輸入數據
• 前向傳播+反向傳播
• 更新參數

t.set_num_threads(8)
for epoch in range(2):  
    
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        
        # 輸入數據
        inputs, labels = data
        
        # 梯度清零
        optimizer.zero_grad()
        
        # forward + backward 
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()   
        
        # 更新參數 
        optimizer.step()
        
        # 打印log信息
        # loss 是一個scalar,需要使用loss.item()來獲取數值，不能使用loss[0]
        running_loss += loss.item()
        if i % 2000 == 1999: # 每2000個batch打印一下訓練狀態
            print('[%d, %5d] loss: %.3f' \
                  % (epoch+1, i+1, running_loss / 2000))
            running_loss = 0.0
print('Finished Training')

此處僅訓練了2個epoch（遍歷完一遍數據集稱為一個epoch），來看看網絡有沒有效果。將測試圖片輸入到網絡中，計算它的label，然后與實際的label進行比較。

dataiter = iter(testloader)
images, labels = dataiter.next() # 一個batch返回4張圖片
print('實際的label: ', ' '.join(\
            '%08s'%classes[labels[j]] for j in range(4)))
show(tv.utils.make_grid(images / 2 - 0.5)).resize((400,100))

接着計算網絡預測的label：

# 計算圖片在每個類別上的分數
outputs = net(images)
# 得分最高的那個類
_, predicted = t.max(outputs.data, 1)

print('預測結果: ', ' '.join('%5s'\
            % classes[predicted[j]] for j in range(4)))

已經可以看出效果，准確率100%，但這只是一部分的圖片，再來看看在整個測試集上的效果。

correct = 0 # 預測正確的圖片數
total = 0 # 總共的圖片數


# 由於測試的時候不需要求導，可以暫時關閉autograd，提高速度，節約內存
with t.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = t.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum()

print('10000張測試集中的准確率為: %d %%' % (100 * correct // total))

10000張測試集中的准確率為: 54 %
訓練的准確率遠比隨機猜測(准確率10%)好，證明網絡確實學到了東西。

在GPU上訓練

就像之前把Tensor從CPU轉到GPU一樣，模型也可以類似地從CPU轉到GPU。

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

net.to(device)
images = images.to(device)
labels = labels.to(device)
output = net(images)
loss= criterion(output,labels)

loss

如果發現在GPU上並沒有比CPU提速很多，實際上是因為網絡比較小，GPU沒有完全發揮自己的真正實力。

對PyTorch的基礎介紹至此結束。總結一下，主要包含以下內容。

Tensor: 類似Numpy數組的數據結構，與Numpy接口類似，可方便地互相轉換。
autograd/Variable: 為tensor提供自動求導功能。
nn: 專門為神經網絡設計的接口，提供了很多有用的功能(神經網絡層，損失函數，優化器等)。
神經網絡訓練: 以CIFAR-10分類為例演示了神經網絡的訓練流程，包括數據加載、網絡搭建、訓練及測試。
通過本節的學習，相信讀者可以體會出PyTorch具有接口簡單、使用靈活等特點。從下一章開始將深入系統地講解PyTorch的各部分知識。