評價此頁

torch.nn 到底是什麼?#

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

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

我們建議將此教程作為 notebook 執行,而不是指令碼。要下載 notebook (.ipynb) 檔案,請單擊頁面頂部的連結。

PyTorch 提供了設計優雅的模組和類 torch.nntorch.optimDatasetDataLoader 來幫助您建立和訓練神經網路。為了充分利用它們的功能併為您的特定問題進行定製,您需要真正理解它們的工作原理。為了培養這種理解,我們將首先在 MNIST 資料集上訓練一個基礎神經網路,而不使用這些模型中的任何功能;一開始我們將只使用最基本的 PyTorch 張量功能。然後,我們將一次逐步新增 torch.nntorch.optimDatasetDataLoader 的一項功能,精確地展示每個部分的作用,以及它如何使程式碼更簡潔或更靈活。

本教程假定您已安裝 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)
nn tutorial
(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([-1.8047, -2.2880, -2.9946, -2.7956, -2.5301, -2.2179, -2.2883, -2.5213,
        -2.1877, -1.9631], 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.3240, grad_fn=<NegBackward0>)

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

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

讓我們檢查一下我們隨機模型的準確率,這樣我們就可以看到隨著損失的改善,我們的準確率是否也在提高。

print(accuracy(preds, yb))
tensor(0.1094)

我們現在可以執行訓練迴圈。在每次迭代中,我們將

  • 選擇一個數據的小批次(大小為 bs

  • 使用模型進行預測

  • 計算損失

  • loss.backward() 更新模型梯度,在本例中是 weightsbias

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

然後,我們將梯度設定為零,以便為下一次迴圈做好準備。否則,我們的梯度將記錄所有已發生操作的累積總數(即 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.0836, 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.0836, grad_fn=<NllLossBackward0>) tensor(1.)

使用 nn.Module 進行重構#

接下來,我們將使用 nn.Modulenn.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

由於我們現在使用的是一個物件而不是一個函式,所以我們首先必須例項化我們的模型

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

print(loss_func(model(xb), yb))
tensor(2.3371, 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.0837, grad_fn=<NllLossBackward0>)

使用 nn.Linear 進行重構#

我們繼續重構我們的程式碼。而不是手動定義和初始化 self.weightsself.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.3447, grad_fn=<NllLossBackward0>)

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

fit()

print(loss_func(model(xb), yb))
tensor(0.0817, 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()

而是隻使用

(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.2724, grad_fn=<NllLossBackward0>)
tensor(0.0809, grad_fn=<NllLossBackward0>)

使用 Dataset 進行重構#

PyTorch 有一個抽象的 Dataset 類。Dataset 可以是任何具有 __len__ 函式(由 Python 的標準 len 函式呼叫)和 __getitem__ 函式作為索引方式的物件。本教程詳細介紹了建立一個自定義 FacialLandmarkDataset 類作為 Dataset 子類的精彩示例。

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

from torch.utils.data import TensorDataset

x_trainy_train 都可以組合到一個 TensorDataset 中,該資料集將更容易迭代和切片。

以前,我們必須單獨遍歷 xy 值的最小批次

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.0814, grad_fn=<NllLossBackward0>)

使用 DataLoader 進行重構#

PyTorch 的 DataLoader 負責管理批次。您可以從任何 Dataset 建立 DataLoaderDataLoader 使迭代批次更加容易。您不再需要使用 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.0804, grad_fn=<NllLossBackward0>)

得益於 PyTorch 的 nn.Modulenn.ParameterDatasetDataLoader,我們的訓練迴圈現在明顯更小,更容易理解。現在,讓我們嘗試新增建立有效模型在實踐中所需的 J 些基本功能。

新增驗證#

在第一部分中,我們只是試圖建立一個合理的訓練迴圈來用於我們的訓練資料。實際上,您總是應該有一個驗證集,以識別您是否過擬合。

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

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

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.BatchNorm2dnn.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.3900)
1 tensor(0.3259)

建立 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.3494310810804367
1 0.28792409490942955

您可以使用這 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

動量是隨機梯度下降的一個變體,它還考慮了先前的更新,並且通常能帶來更快的訓練速度。

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.4556122406005859
1 0.2508379440307617

使用 nn.Sequential#

torch.nn 還有一個方便的類可供我們使用來簡化程式碼:SequentialSequential 物件按順序執行其中包含的每個模組。這是編寫神經網路的一種更簡單的方法。

為了利用這一點,我們需要能夠輕鬆地從給定函式定義一個自定義層。例如,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.3620248727560043
1 0.22110013802051545

包裝 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.3840307270050049
1 0.2389558298230171

使用您的 加速器#

如果您有幸能夠使用 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.1930989332795143
1 0.19160217237174512

結束語#

現在我們有了一個通用的資料管道和訓練迴圈,您可以使用它來訓練各種 Pytorch 模型。為了瞭解訓練模型現在可以多麼簡單,請檢視 mnist_sample notebook

當然,您會想要新增許多功能,例如資料增強、超引數調整、監控訓練、遷移學習等。這些功能在 fastai 庫中都可用,該庫採用了本教程中所示的相同設計方法進行開發,為希望進一步改進模型的從業者提供了自然的下一步。

我們一開始就承諾,將透過示例解釋 torch.nntorch.optimDatasetDataLoader。所以,讓我們總結一下我們所看到的

  • torch.nn:

    • Module:建立一個可呼叫的物件,其行為類似於函式,但也可以包含狀態(例如神經網路層權重)。它知道它包含哪些 Parameter(s),並且可以清零所有梯度,迴圈遍歷它們進行權重更新等。

    • Parameter:一個張量的包裝器,它告訴 Module 它具有在反向傳播期間需要更新的權重。只有 requires_grad 屬性設定為 True 的張量才會被更新。

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

  • torch.optim:包含最佳化器,例如 SGD,它們在反向傳播步驟期間更新 Parameter 的權重。

  • Dataset:一個抽象介面,包含具有 __len____getitem__ 的物件,包括 PyTorch 提供的類,例如 TensorDataset

  • DataLoader:接收任何 Dataset 並建立一個返回資料批次的迭代器。

指令碼總執行時間: (0 分鐘 25.360 秒)