PyTorch實例入門（1）：圖像分類

PyTorch的0.4版本帶來了不小的變化，其中我最喜歡的是：

1. Tensor和Variable這兩個類合併了。原來nn的input是一個variable，現在可以直接用tensor。這樣在語法上更簡潔易用，對初學者也更容易理解。
2. Windows support。官方支持了windows，作爲一個最近迴歸了windows的人很開心哈哈。

之前內存泄露的問題似乎也解決了，所以我又開心地從Tensorflow蹦回了PyTorch，順便寫點教程。先從最基本的開始，今天這篇文章講怎麼完成圖像分類的任務。閱讀前假設對神經網絡和Python有一定了解。

通常我們在使用PyTorch的時候會用到兩個包，一個是torch，一個是torchvision。其中torch是關於運算的包，torchvision則是打包了一些數據集，另外用torch實現了一些常見的神經網絡模型，比如ResNet。

我們使用CIFAR-10作爲數據集，包含了10個類別60000張圖片，每張圖片的大小爲32x32，其中訓練圖片50000張，測試圖片10000張。下圖是一些示例

torchvision中已經打包好了這個數據集，我們不用自己下載，直接如下調用就可以了。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
# cifar-10官方提供的數據集是用numpy array存儲的
# 下面這個transform會把numpy array變成torch tensor，然後把rgb值歸一到[0, 1]這個區間
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 在構建數據集的時候指定transform，就會應用我們定義好的transform
# root是存儲數據的文件夾，download=True指定如果數據不存在先下載數據
cifar_train = torchvision.datasets.CIFAR10(root='./data', train=True,
                                           download=True, transform=transform)
cifar_test = torchvision.datasets.CIFAR10(root='./data', train=False,
                                          transform=transform)

 load完了之後我們看看這兩個數據集的信息

print(cifar_train):
==>
Dataset CIFAR10
    Number of datapoints: 50000
    Split: train
    Root Location: ./data
    Transforms (if any): None
    Target Transforms (if any): None

print(cifar_test)
==>
Dataset CIFAR10
    Number of datapoints: 10000
    Split: test
    Root Location: ./data
    Transforms (if any): None
    Target Transforms (if any): None

# 數據其實是用numpy array存儲的，label是個list
print(cifar_train.train_data.shape)
print(type(cifar_train.train_labels), len(cifar_train.train_labels))
==>
(50000, 32, 32, 3)
<class 'list'> 50000

print(cifar_test.test_data.shape)
print(type(cifar_test.test_labels), len(cifar_test.test_labels))
==>
(10000, 32, 32, 3)
<class 'list'> 10000

在訓練的時候我們可以自己寫代碼手動遍歷數據集，指定batch和遍歷方法，不過PyTorch提供了一個DataLoader類來方便我們完成這些操作。

trainloader = torch.utils.data.DataLoader(cifar_train, batch_size=32, shuffle=True)
testloader = torch.utils.data.DataLoader(cifar_test, batch_size=32, shuffle=True)

現在我們來定義卷積神經網絡，爲了簡單起見，我們使用經典的LeNet，它包含兩個卷積層和三個全連接層，網絡結構如圖

在PyTorch中定義神經網絡非常簡單，第一步先是繼承nn.module這個類，然後定義如下兩個函數，一般的網絡這樣操作就足夠了

class LeNet(nn.Module):
    # 一般在__init__中定義網絡需要的操作算子，比如卷積、全連接算子等等
    def __init__(self):
        super(LeNet, self).__init__()
        # Conv2d的第一個參數是輸入的channel數量，第二個是輸出的channel數量，第三個是kernel size
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 由於上一層有16個channel輸出，每個feature map大小爲5*5，所以全連接層的輸入是16*5*5
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        # 最終有10類，所以最後一個全連接層輸出數量是10
        self.fc3 = nn.Linear(84, 10)
        self.pool = nn.MaxPool2d(2, 2)
    # forward這個函數定義了前向傳播的運算，只需要像寫普通的python算數運算那樣就可以了
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(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

PyTorch具有自動求導功能，我們不需要自己寫backward函數，所以很直觀方便，寫神經網絡的結構就像寫普通的數學運算公式一樣。定義好網絡之後我們就可以訓練了，訓練的代碼也非常簡單。首先，我們先構建一個網絡實例。由於需要用到GPU，所以先獲取device，然後再把網絡的參數複製到GPU上

# 如果你沒有GPU，那麼可以忽略device相關的代碼
device = torch.device("cuda:0")
net = LeNet().to(device)

然後我們需要定義Loss函數和優化方法，最簡單的就是使用SGD了。PyTorch都預先定義好了這些東西

# optim中定義了各種各樣的優化方法，包括SGD
import torch.optim as optim

# CrossEntropyLoss就是我們需要的損失函數
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

下面我們正式開始訓練

print("Start Training...")
for epoch in range(30):
    # 我們用一個變量來記錄每100個batch的平均loss
    loss100 = 0.0
    # 我們的dataloader派上了用場
    for i, data in enumerate(trainloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device) # 注意需要複製到GPU
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        loss100 += loss.item()
        if i % 100 == 99:
            print('[Epoch %d, Batch %5d] loss: %.3f' %
                  (epoch + 1, i + 1, loss100 / 100))
            loss100 = 0.0

print("Done Training!")

以上代碼中，核心的代碼是下面五行，我用註釋解釋了每一行的作用

# 首先要把梯度清零，不然PyTorch每次計算梯度會累加，不清零的話第二次算的梯度等於第一次加第二次的       
optimizer.zero_grad()
# 計算前向傳播的輸出
outputs = net(inputs)
# 根據輸出計算loss
loss = criterion(outputs, labels)
# 算完loss之後進行反向梯度傳播，這個過程之後梯度會記錄在變量中
loss.backward()
# 用計算的梯度去做優化
optimizer.step()

ok，訓練完了之後我們來檢測一下準確率，我們用訓練好的模型來預測test數據集

# 構造測試的dataloader
dataiter = iter(testloader)
# 預測正確的數量和總數量
correct = 0
total = 0
# 使用torch.no_grad的話在前向傳播中不記錄梯度，節省內存
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        # 預測
        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))

在我這裏最終訓練了30個epoch之後準確率爲64%。