torch.nn 究竟是什麼？

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

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

建議以 notebook 形式而非指令碼形式執行本教程。要下載 notebook (.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（一種 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)

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

• 選擇一個 mini-batch 的資料（大小為 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 庫中的所有函式（而庫的其他部分包含類）。除了各種損失函式和啟用函式外，你還可以在這裡找到一些方便建立神經網路的函式，例如池化函式。（還有用於卷積、線性層等的函式，但正如我們將看到的，這些通常使用庫的其他部分來處理會更好）。

如果你使用負對數似然損失和 log 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）module 的概念混淆，後者是一個可以匯入的 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，我們需要在計算下一個 minibatch 的梯度之前呼叫它。）

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__ 函式（作為索引訪問方式）的物件。本教程提供了一個很好的示例，介紹瞭如何建立一個繼承自 Dataset 的自定義 FacialLandmarkDataset 類。

PyTorch 的 TensorDataset 是一個封裝張量的 Dataset。透過定義長度和索引方式，這也提供了一種沿張量第一個維度進行迭代、索引和切片的方法。這將使我們在訓練時更容易在同一行中訪問自變數和因變數。

from torch.utils.data import TensorDataset

x_train 和 y_train 都可以合併到一個 TensorDataset 中，這樣迭代和切片會更容易。

train_ds = TensorDataset(x_train, y_train)

之前，我們必須單獨迭代 x 和 y 值的 minibatches

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 負責管理 batch。你可以從任何 Dataset 建立一個 DataLoader。DataLoader 使迭代 batch 更加容易。不再需要使用 train_ds[i*bs : i*bs+bs]，DataLoader 會自動提供每個 minibatch。

from torch.utils.data import DataLoader

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

之前，我們的迴圈是這樣迭代 batch (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，我們的訓練迴圈現在變得非常小且更容易理解。現在讓我們嘗試新增在實踐中建立有效模型所需的基本功能。

新增驗證

在第一節中，我們只是試圖建立一個合理的訓練迴圈用於訓練資料。實際上，你始終應該有一個驗證集，以判斷是否發生了過擬合。

打亂訓練資料非常重要，以防止批次之間存在相關性和過擬合。另一方面，無論我們是否打亂驗證集，驗證損失都將是相同的。由於打亂資料會花費額外的時間，因此打亂驗證資料沒有意義。

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

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 執行必要的運算來訓練模型並計算每個訓練週期的訓練損失和驗證損失。

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 返回用於訓練集和驗證集的資料載入器 (dataloaders)。

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.29393543709516523
1 0.3258970961511135

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

切換到 CNN

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

我們將使用 PyTorch 預定義的 Conv2d 類作為我們的卷積層。我們定義一個具有 3 個卷積層的 CNN。每個卷積後面都跟著一個 ReLU。最後，我們執行平均池化（average pooling）。(注意，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 物件按順序執行其中包含的每個模組。這是一種編寫神經網路更簡單的方式。

為了利用這一點，我們需要能夠輕鬆地從給定的函式定義一個**自定義層 (custom layer)**。例如，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.AdaptiveAvgPool2d 替換 nn.AvgPool2d，這允許我們定義我們想要的*輸出*張量大小，而不是我們擁有的*輸入*張量大小。因此，我們的模型將適用於任何大小的輸入。

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.1786795982837677
1 0.16877726690769196

結束語

現在我們有了一個通用的資料管道和訓練迴圈，您可以使用它來訓練多種型別的 PyTorch 模型。想看看現在訓練模型有多簡單，請檢視 mnist_sample notebook。

當然，還有許多您想要新增的功能，例如資料增強、超引數調優、訓練監控、遷移學習等等。這些功能在 fastai 庫中都可用，該庫是按照本教程中展示的設計方法開發的，為希望進一步提升模型的實踐者提供了自然的下一步。

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

• torch.nn:

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

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

  • functional：一個模組（通常按慣例匯入到 F 名稱空間中），包含啟用函式、損失函式等，以及卷積層和線性層等無狀態版本的層。

• torch.optim：包含最佳化器，例如 SGD，它們在反向傳播步驟中更新 Parameter 的權重

• Dataset：具有 __len__ 和 __getitem__ 方法的物件的抽象介面，包括 PyTorch 提供的類，如 TensorDataset

• DataLoader：接受任何 Dataset 並建立一個迭代器，該迭代器返回批次資料。