torch.nn 真的是什麼？

建立於：2018 年 12 月 26 日 | 最後更新：2025 年 1 月 24 日 | 最後驗證：2024 年 11 月 05 日

作者： Jeremy Howard, fast.ai。感謝 Rachel Thomas 和 Francisco Ingham。

我們建議將本教學作為筆記本執行，而不是腳本。要下載筆記本（.ipynb）檔案，請點擊頁面頂部的連結。

PyTorch 提供了設計優雅的模組和類別 torch.nn , torch.optim , Dataset , 和 DataLoader，以幫助你建立和訓練神經網路。為了充分利用它們的力量並為你的問題自定義它們，你需要真正理解它們到底在做什麼。為了培養這種理解，我們首先會在 MNIST 數據集上訓練基本的神經網路，而不使用這些模型的任何功能；我們最初只會使用最基本的 PyTorch 張量功能。然後，我們將逐步從 torch.nn、torch.optim、Dataset 或 DataLoader 中添加一個功能，展示每個部分的作用，以及它是如何使程式碼更簡潔或更靈活的。

本教學假設你已經安裝了 PyTorch，並且熟悉張量運算的基本知識。（如果你熟悉 Numpy 陣列運算，你會發現這裡使用的 PyTorch 張量運算幾乎相同）。

MNIST 數據設定

我們將使用經典的 MNIST 數據集，它包含手繪數字（介於 0 和 9 之間）的黑白圖像。

我們將使用 pathlib 處理路徑（Python 3 標準庫的一部分），並使用 requests 下載數據集。 我們只會在我們使用模組時導入它們，因此你可以清楚地看到每個點正在使用什麼。

from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

URL = "https://github.com/pytorch/tutorials/raw/main/_static/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)

此數據集採用 numpy 陣列格式，並已使用 pickle 儲存，pickle 是一種用於序列化數據的 Python 特定格式。

import pickle
import gzip

with gzip.open((PATH / FILENAME).as_posix(), "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

每張影像都是 28 x 28，並以長度為 784 (=28x28) 的扁平列儲存。讓我們看看一個；我們需要先將它重塑為 2d。

from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
# ``pyplot.show()`` only if not on Colab
try:
    import google.colab
except ImportError:
    pyplot.show()
print(x_train.shape)

(50000, 784)

PyTorch 使用 torch.tensor，而不是 numpy 陣列，所以我們需要轉換我們的數據。

import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)
n, c = x_train.shape
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]]) tensor([5, 0, 4,  ..., 8, 4, 8])
torch.Size([50000, 784])
tensor(0) tensor(9)

從頭開始的神經網路（沒有 torch.nn）

讓我們首先僅使用 PyTorch 張量運算建立一個模型。我們假設你已經熟悉神經網路的基本知識。（如果你不熟悉，你可以在 course.fast.ai 學習它們）。

PyTorch 提供了創建隨機或零填充張量的方法，我們將使用它們來創建我們簡單線性模型的權重和偏差。這些只是常規張量，但有一個非常特殊的添加：我們告訴 PyTorch 它們需要梯度。 這會導致 PyTorch 記錄對張量執行的所有操作，以便它可以自動計算反向傳播期間的梯度！

對於權重，我們在初始化之後設定 requires_grad，因為我們不希望該步驟包含在梯度中。（請注意，PyTorch 中尾隨的 _ 表示該操作是就地執行的。）

注意

我們在這裡使用 Xavier 初始化（通過乘以 1/sqrt(n)）來初始化權重。

import math

weights = torch.randn(784, 10) / math.sqrt(784)
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

由於 PyTorch 自動計算梯度的能力，我們可以將任何標準 Python 函數（或可調用物件）用作模型！ 因此，讓我們只編寫一個簡單的矩陣乘法和廣播加法來創建一個簡單的線性模型。 我們還需要一個啟動函數，所以我們將編寫 log_softmax 並使用它。 請記住：雖然 PyTorch 提供了許多預先編寫的損失函數、啟動函數等等，但你可以使用純 Python 輕鬆編寫自己的函數。 PyTorch 甚至會自動為你的函數創建快速加速器或向量化 CPU 程式碼。

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

在以上內容中，@ 代表矩陣乘法運算。我們會對一批資料（在此例中為 64 張圖片）呼叫我們的函數。這就是一次前向傳遞。請注意，由於我們從隨機權重開始，因此在這個階段我們的預測不會比隨機好。

bs = 64  # batch size

xb = x_train[0:bs]  # a mini-batch from x
preds = model(xb)  # predictions
preds[0], preds.shape
print(preds[0], preds.shape)

tensor([-2.5452, -2.0790, -2.1832, -2.6221, -2.3670, -2.3854, -2.9432, -2.4391,
        -1.8657, -2.0355], grad_fn=<SelectBackward0>) torch.Size([64, 10])

如你所見，preds 張量不僅包含張量值，還包含梯度函數。我們稍後會使用它來進行反向傳播。

讓我們實現負對數似然作為損失函數（同樣，我們可以只使用標準 Python）

def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll

讓我們檢查一下我們的隨機模型的損失，以便我們稍後查看在反向傳播之後是否有改善。

yb = y_train[0:bs]
print(loss_func(preds, yb))

tensor(2.4020, grad_fn=<NegBackward0>)

讓我們也實現一個函數來計算我們模型的準確度。對於每個預測，如果具有最大值的索引與目標值匹配，則預測是正確的。

def accuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

讓我們檢查一下我們的隨機模型的準確度，以便我們查看準確度是否隨著損失的改善而提高。

print(accuracy(preds, yb))

tensor(0.0938)

我們現在可以執行一個訓練迴圈。對於每次迭代，我們將

• 選擇一個小批次的資料（大小為 bs）

• 使用模型進行預測

• 計算損失

• loss.backward() 更新模型的梯度，在本例中為 weights 和 bias。

我們現在使用這些梯度來更新權重和偏差。我們在 torch.no_grad() 上下文管理器中執行此操作，因為我們不希望這些操作被記錄下來用於我們下一次的梯度計算。你可以在 這裡 閱讀更多關於 PyTorch 的 Autograd 如何記錄操作的資訊。

然後我們將梯度設定為零，以便為下一個迴圈做好準備。否則，我們的梯度會記錄已發生的所有操作的運行總計（即 loss.backward() 將梯度添加到已儲存的任何內容，而不是替換它們）。

提示

你可以使用標準的 python 除錯器來逐步執行 PyTorch 程式碼，從而可以在每個步驟檢查各種變數值。取消註解下面的 set_trace() 來試試看。

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        #         set_trace()
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()

就這樣：我們從頭開始創建並訓練了一個最小的神經網路（在本例中為邏輯迴歸，因為我們沒有隱藏層）！

讓我們檢查損失和準確度，並將其與我們之前獲得的結果進行比較。我們預期損失會減少，準確度會提高，而且它們確實如此。

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

tensor(0.0813, grad_fn=<NegBackward0>) tensor(1.)

使用 torch.nn.functional

我們現在將重構我們的程式碼，使其與之前執行相同的操作，只不過我們將開始利用 PyTorch 的 nn 類別，使其更加簡潔和靈活。從這裡的每一步開始，我們都應該使我們的程式碼變得：更短、更容易理解和/或更靈活。

第一個也是最簡單的步驟是通過用 torch.nn.functional（通常按照慣例匯入到命名空間 F 中）中的函數替換我們手寫的激活和損失函數來縮短我們的程式碼。這個模組包含 torch.nn 庫中的所有函數（而庫的其他部分包含類別）。除了各種各樣的損失和激活函數之外，你還會在這裡找到一些用於創建神經網路的便捷函數，例如池化函數。（也有用於執行卷積、線性層等的函數，但正如我們將看到的，通常最好使用庫的其他部分來處理這些函數。）

如果你使用的是負對數似然損失和對數 softmax 激活，那麼 Pytorch 提供了一個單一函數 F.cross_entropy，它結合了兩者。因此我們甚至可以從我們的模型中刪除激活函數。

import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

請注意，我們不再在 model 函數中呼叫 log_softmax。讓我們確認我們的損失和準確度與之前相同

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

tensor(0.0813, grad_fn=<NllLossBackward0>) tensor(1.)

使用 nn.Module 重構

接下來，我們將使用 nn.Module 和 nn.Parameter，以便獲得更清晰和更簡潔的訓練迴圈。我們繼承了 nn.Module（它本身是一個類別，並且能夠追蹤狀態）。在這種情況下，我們想要創建一個類別，該類別持有我們的權重、偏差和前向步驟的方法。nn.Module 具有許多屬性和方法（例如 .parameters() 和 .zero_grad()），我們將使用它們。

注意

nn.Module（大寫 M）是 PyTorch 特有的概念，並且是一個我們將大量使用的類別。nn.Module 不應與 Python 的（小寫 m）模組概念混淆，模組是可以匯入的 Python 程式碼檔案。

from torch import nn

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(784, 10) / math.sqrt(784))
        self.bias = nn.Parameter(torch.zeros(10))

    def forward(self, xb):
        return xb @ self.weights + self.bias

由於我們現在使用的是物件而不是僅僅使用函數，因此我們首先必須實例化我們的模型

model = Mnist_Logistic()

現在我們可以像以前一樣計算損失。請注意，nn.Module 物件的使用方式就好像它們是函數一樣（即它們是可呼叫的），但在幕後 Pytorch 會自動呼叫我們的 forward 方法。

print(loss_func(model(xb), yb))

tensor(2.3096, grad_fn=<NllLossBackward0>)

以前，對於我們的訓練迴圈，我們必須按名稱更新每個參數的值，並手動將每個參數的梯度歸零，如下所示

with torch.no_grad():
    weights -= weights.grad * lr
    bias -= bias.grad * lr
    weights.grad.zero_()
    bias.grad.zero_()

現在我們可以利用 model.parameters() 和 model.zero_grad()（它們都是由 PyTorch 為 nn.Module 定義的）來使這些步驟更加簡潔，並且更不容易忘記我們的某些參數，特別是如果我們有一個更複雜的模型

with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()

我們將把我們的小型訓練迴圈封裝在一個 fit 函數中，以便我們稍後可以再次執行它。

def fit():
    for epoch in range(epochs):
        for i in range((n - 1) // bs + 1):
            start_i = i * bs
            end_i = start_i + bs
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            pred = model(xb)
            loss = loss_func(pred, yb)

            loss.backward()
            with torch.no_grad():
                for p in model.parameters():
                    p -= p.grad * lr
                model.zero_grad()

fit()

讓我們仔細檢查一下我們的損失是否已減少

print(loss_func(model(xb), yb))

tensor(0.0821, grad_fn=<NllLossBackward0>)

使用 nn.Linear 重構

我們持續重構程式碼。不再手動定義和初始化 self.weights 和 self.bias，以及計算 xb  @ self.weights + self.bias，我們將改用 Pytorch 的 nn.Linear 類別來建立線性層，它會自動完成這些操作。Pytorch 有許多預定義的層，可以大大簡化我們的程式碼，而且通常也能加快速度。

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    def forward(self, xb):
        return self.lin(xb)

我們以與之前相同的方式實例化模型並計算損失。

model = Mnist_Logistic()
print(loss_func(model(xb), yb))

tensor(2.3313, grad_fn=<NllLossBackward0>)

我們仍然能夠使用與之前相同的 fit 方法。

fit()

print(loss_func(model(xb), yb))

tensor(0.0819, grad_fn=<NllLossBackward0>)

使用 torch.optim 重構

Pytorch 還有一個包含各種優化演算法的套件，torch.optim。 我們可以使用優化器的 step 方法來執行前向步驟，而不是手動更新每個參數。

這將讓我們取代之前手動編寫的優化步驟

with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()

而是只使用

opt.step()
opt.zero_grad()

(optim.zero_grad() 會將梯度重設為 0，我們需要在計算下一個小批次的梯度之前呼叫它。)

from torch import optim

我們將定義一個小函數來建立我們的模型和優化器，以便我們將來可以重複使用它。

def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, opt = get_model()
print(loss_func(model(xb), yb))

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

tensor(2.2659, grad_fn=<NllLossBackward0>)
tensor(0.0810, grad_fn=<NllLossBackward0>)

使用 Dataset 重構

PyTorch 有一個抽象的 Dataset 類別。 Dataset 可以是任何具有 __len__ 函數（由 Python 的標準 len 函數呼叫）和 __getitem__ 函數作為索引方式的東西。 本教學課程逐步介紹了創建自定義 FacialLandmarkDataset 類別作為 Dataset 子類的良好範例。

PyTorch 的 TensorDataset 是一個包裝張量的 Dataset。 通過定義長度和索引方式，這也讓我們能夠沿著張量的第一個維度進行迭代、索引和切片。 這將使我們更容易在訓練時在同一行中訪問自變數和應變數。

from torch.utils.data import TensorDataset

x_train 和 y_train 都可以合併到一個 TensorDataset 中，這樣可以更輕鬆地進行迭代和切片。

train_ds = TensorDataset(x_train, y_train)

以前，我們必須分別迭代 x 和 y 值的小批次

xb = x_train[start_i:end_i]
yb = y_train[start_i:end_i]

現在，我們可以一起執行這兩個步驟

xb,yb = train_ds[i*bs : i*bs+bs]

model, opt = get_model()

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        xb, yb = train_ds[i * bs: i * bs + bs]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

tensor(0.0826, grad_fn=<NllLossBackward0>)

使用 DataLoader 重構

PyTorch 的 DataLoader 負責管理批次。 您可以從任何 Dataset 建立 DataLoader。DataLoader 使迭代批次更容易。 無需使用 train_ds[i*bs : i*bs+bs]，DataLoader 會自動為我們提供每個小批次。

from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs)

以前，我們的迴圈像這樣迭代批次 (xb, yb)

for i in range((n-1)//bs + 1):
    xb,yb = train_ds[i*bs : i*bs+bs]
    pred = model(xb)

現在，我們的迴圈更加簡潔，因為 (xb, yb) 會自動從資料載入器載入

for xb,yb in train_dl:
    pred = model(xb)

model, opt = get_model()

for epoch in range(epochs):
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

tensor(0.0818, grad_fn=<NllLossBackward0>)

感謝 PyTorch 的 nn.Module、nn.Parameter、Dataset 和 DataLoader，我們的訓練迴圈現在顯著更小且更容易理解。 現在讓我們嘗試新增必要的基礎功能，以在實踐中創建有效的模型。

新增驗證

在第 1 節中，我們只是嘗試建立一個合理的訓練迴圈以用於我們的訓練資料。 實際上，您始終也應該有一個 驗證集，以便確定您是否過擬合。

攪亂訓練資料對於防止批次之間的相關性和過擬合是重要的。 另一方面，無論我們是否攪亂驗證集，驗證損失都將是相同的。 由於攪亂需要額外的時間，因此攪亂驗證資料沒有任何意義。

我們將使用比訓練集大兩倍的批次大小用於驗證集。 這是因為驗證集不需要反向傳播，因此佔用的記憶體更少（它不需要儲存梯度）。 我們利用這一點來使用更大的批次大小並更快地計算損失。

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)

valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=bs * 2)

我們將在每個 epoch 結束時計算並印出驗證損失。

(請注意，我們總是在訓練之前呼叫 model.train()，在推論之前呼叫 model.eval()，因為這些被諸如 nn.BatchNorm2d 和 nn.Dropout 之類的層使用，以確保這些不同階段的適當行為。)

model, opt = get_model()

for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

    model.eval()
    with torch.no_grad():
        valid_loss = sum(loss_func(model(xb), yb) for xb, yb in valid_dl)

    print(epoch, valid_loss / len(valid_dl))

0 tensor(0.3048)
1 tensor(0.2872)

建立 fit() 和 get_data()

我們現在將對自己進行一些重構。 由於我們經過了兩次相似的過程，分別計算訓練集和驗證集的損失，讓我們將其轉換為自己的函數 loss_batch，它計算一個批次的損失。

我們為訓練集傳入一個優化器，並使用它來執行反向傳播。 對於驗證集，我們不傳入優化器，因此該方法不執行反向傳播。

def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)

fit 執行必要的運算來訓練我們的模型，並計算每個 epoch 的訓練和驗證損失。

import numpy as np

def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)

get_data 傳回訓練集和驗證集的資料載入器。

def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

現在，我們獲取資料載入器並擬合模型的整個過程可以在 3 行程式碼中運行

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
model, opt = get_model()
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.2939354367017746
1 0.3258970756947994

您可以使用這 3 行基本程式碼來訓練各種模型。 讓我們看看是否可以使用它們來訓練卷積神經網路 (CNN)！

切換到 CNN

我們現在將使用三個卷積層建立我們的神經網路。 由於上一節中的任何函數都沒有對模型形式做出任何假設，因此我們將能夠使用它們來訓練 CNN 而無需進行任何修改。

我們將使用 PyTorch 預定義的 Conv2d 類別作為我們的卷積層。我們定義一個具有 3 個卷積層的 CNN。 每次卷積後都會接一個 ReLU。 最後，我們執行平均池化。（請注意，view 是 PyTorch 版本的 Numpy 的 reshape）

class Mnist_CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1)

    def forward(self, xb):
        xb = xb.view(-1, 1, 28, 28)
        xb = F.relu(self.conv1(xb))
        xb = F.relu(self.conv2(xb))
        xb = F.relu(self.conv3(xb))
        xb = F.avg_pool2d(xb, 4)
        return xb.view(-1, xb.size(1))

lr = 0.1

動量 (Momentum) 是隨機梯度下降的一種變體，它也會考慮先前的更新，通常會導致更快的訓練。

model = Mnist_CNN()
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3646130460739136
1 0.26228193019628526

使用 nn.Sequential

torch.nn 還有另一個方便的類別，我們可以使用它來簡化我們的程式碼：Sequential。 Sequential 物件會以循序方式執行其中包含的每個模組。 這是編寫神經網路的一種更簡單的方式。

為了利用這一點，我們需要能夠從給定的函數輕鬆定義一個自定義層。 例如，PyTorch 沒有 view 層，我們需要為我們的網路建立一個。 Lambda 將建立一個層，然後我們可以在使用 Sequential 定義網路時使用它。

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


def preprocess(x):
    return x.view(-1, 1, 28, 28)

使用 Sequential 建立的模型很簡單

model = nn.Sequential(
    Lambda(preprocess),
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AvgPool2d(4),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3330025281429291
1 0.22993727023601532

封裝 DataLoader

我們的 CNN 相當簡潔，但它僅適用於 MNIST，因為

• 它假設輸入是一個 28*28 的長向量

• 它假設最終的 CNN 網格大小為 4*4（因為這是我們使用的平均池化核心大小）

讓我們擺脫這兩個假設，以便我們的模型適用於任何 2d 單通道圖像。 首先，我們可以透過將資料預處理移至生成器中來刪除初始 Lambda 層

def preprocess(x, y):
    return x.view(-1, 1, 28, 28), y


class WrappedDataLoader:
    def __init__(self, dl, func):
        self.dl = dl
        self.func = func

    def __len__(self):
        return len(self.dl)

    def __iter__(self):
        for b in self.dl:
            yield (self.func(*b))

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

接下來，我們可以將 nn.AvgPool2d 替換為 nn.AdaptiveAvgPool2d，這允許我們定義我們想要的輸出張量的大小，而不是我們擁有的輸入張量。 因此，我們的模型將適用於任何大小的輸入。

model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 16, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.Conv2d(16, 10, kernel_size=3, stride=2, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d(1),
    Lambda(lambda x: x.view(x.size(0), -1)),
)

opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

讓我們試試看

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.3212135115623474
1 0.21439074140787123

使用您的 加速器 (Accelerator)

如果您幸運地可以使用 CUDA 等加速器（您可以從大多數雲端供應商以每小時約 0.50 美元的價格租用一個），則可以使用它來加速您的程式碼。 首先檢查您的加速器是否在 Pytorch 中正常運作

# If the current accelerator is available, we will use it. Otherwise, we use the CPU.
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

Using cuda device

讓我們更新 preprocess 以將批次移動到加速器

def preprocess(x, y):
    return x.view(-1, 1, 28, 28).to(device), y.to(device)


train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
train_dl = WrappedDataLoader(train_dl, preprocess)
valid_dl = WrappedDataLoader(valid_dl, preprocess)

最後，我們可以將我們的模型移動到加速器。

model.to(device)
opt = optim.SGD(model.parameters(), lr=lr, momentum=0.9)

您應該發現它現在運行得更快了

fit(epochs, model, loss_func, opt, train_dl, valid_dl)

0 0.1805140619158745
1 0.17031861956119537

總結

我們現在有一個通用的資料管道和訓練迴圈，您可以使用它來訓練許多類型的使用 Pytorch 的模型。 要了解訓練模型現在可以多麼簡單，請查看 mnist_sample notebook。

當然，您會想要新增許多東西，例如資料擴增、超參數調整、監控訓練、遷移學習等等。 這些功能可在 fastai 函式庫中使用，該函式庫是使用本教學中顯示的相同設計方法開發的，為希望進一步發展其模型的從業者提供了一個自然的下一步。

我們在本教程開始時承諾將透過範例解釋 torch.nn、torch.optim、Dataset 和 DataLoader。 所以讓我們總結一下我們所看到的

• torch.nn:

  • Module：建立一個可呼叫的物件，其行為類似於一個函數，但也可能包含狀態（例如神經網路層權重）。 它知道它包含哪些 Parameter，並且可以將它們的所有梯度歸零、循環遍歷它們以進行權重更新等。

  • Parameter：tensor 的包裝器，它告訴 Module 它具有需要在反向傳播期間更新的權重。 只有設定了 requires_grad 屬性的張量才會被更新

  • functional：一個模組（通常按照慣例匯入到 F 命名空間中），其中包含激活函數、損失函數等，以及諸如卷積層和線性層之類的無狀態版本的層。

• torch.optim：包含諸如 SGD 之類的優化器，這些優化器會在反向步驟期間更新 Parameter 的權重

• Dataset：具有 __len__ 和 __getitem__ 的物件的抽象介面，包括 Pytorch 提供的類別，例如 TensorDataset

• DataLoader：採用任何 Dataset 並建立一個迭代器，該迭代器會傳回資料批次。