【Pytorch 深度學習筆記】The PyTorch nn module

哈囉大家好我是 LukeTseng,感謝您點進本篇筆記,該篇筆記主要配合讀本 《Deep Learning with pytorch》 進行學習,另外透過網路資料作為輔助。本系列筆記是我本人奠基深度學習基礎知識的開始,若文章有誤煩請各位指正,謝謝!

本篇為 《Deep Learning with pytorch》 這本書 6.2 The PyTorch nn module 章節內容的相關筆記。

6.2 The PyTorch nn module

在先前的章節中,作者用一個簡單的溫度計校正問題來解釋機器學習機制,這個問題其實只要一個簡單的線性函數即可解決。

雖然直接將原本的線性模型替換成 NN,在解決這個簡單問題上看似是殺雞用牛刀了,但這是一個非常關鍵奠基的步驟。

在簡單的問題上導入 NN 模組,可以幫助我們熟悉背後的運作機制,為未來解決與擴展更複雜的問題奠定基礎。

為了讓我們能輕鬆的建構起一個 NN,PyTorch 提供了一個專門的子模組稱為 torch.nn。這個模組包含了建立各種 NN 架構所需的所有基礎建構區塊。

模組(Modules)與層(Layers)的關聯

在 PyTorch 的術語中,這些建構神經網路的基礎區塊被統一稱為「模組」(modules),對於熟悉其他深度學習框架的開發者來說,PyTorch 中的模組其實就等同於其他框架中常說的「層」(layers)。

在實作上,所有的 PyTorch 模組都是繼承自 nn.Module 基礎類別的 Python 類別。

模組的兩大核心屬性

一個衍生自 nn.Module 的類別,可包含以下兩種重要的屬性:

  1. 參數(Parameters):模組可以擁有一個或多個 Parameter 實例作為屬性,這些參數本質上就是張量(tensors),它們的數值會在 NN 的訓練過程中透過優化演算法不斷更新與調整(可以把它們想像成先前線性模型中的 weight ww 和 bias bb)。
  2. 子模組(Submodules): 模組也可以將其他繼承自 nn.Module 的子類別作為自己的屬性,當一個模組內部包含其他子模組時,父模組將能夠自動追蹤並管理這些子模組內部的所有參數。

重要的實作規範:避免將子模組隱藏在一般容器中

書中特別提出了一個開發者常犯的嚴重錯誤:子模組必須被定義為頂層屬性(top-level attributes),絕不能將它們埋在標準的 Python 列表(list)或字典(dict)實例中。

如果將子模組放在一般的 list 或 dict 中,PyTorch 的優化器(optimizer)會無法定位到這些子模組,進而導致它們的參數無法被追蹤與更新。

如果模型架構確實要用 list 或 dict 來動態管理多個子模組,PyTorch 提供了專門的替代這兩個的資料結構:

  • nn.ModuleList
  • nn.ModuleDict

用這些專門的類別,才能確保 Optimizer 正確獲取所有的參數。

nn.Linear

nn.Linearnn.Module 的子類別,會對輸入做 affine transformation(仿射變換),即在做類似 y=wx+by = wx + b 的運算。

nn.Linear 是透過其內部的 weightbias 這兩個參數屬性來完成運算的,與書中作者先前在溫度計實驗中手動實作的線性模型中,在功能上完全等價。

基於這個等價關係,我們可直接從此出發,將先前的線性回歸程式碼轉換為使用 PyTorch nn 模組的標準形式。

6.2.1 Using call rather than forward

書中寫了以下範例:

1
2
3
4
import torch.nn as nn

linear_model = nn.Linear(1, 1)
linear_model(t_un_val)

這個範例示範 linear_model 可像函式一樣被呼叫,這是因為所有 PyTorch 提供的 nn.Module 子類別都有定義 __call__ 方法。

換句話說,雖然 linear_model 是一個物件,不是普通的函式,但用 PyTorch 的我們可以寫 output = linear_model(input),因為背後實際上等價於 Python 呼叫:output = linear_model.__call__(input)

也就是說 PyTorch 已經先幫我們做好這些複雜的 function call 了,從此簡化我們在訓練上的流程。

__call__forward 的關係

當用一組參數呼叫 nn.Module 的實例時,最後會呼叫同樣參數的 forward 方法。

forward 負責真正的 forward computation(前向計算,比較常稱前向傳播 forward propagation),也就是模型從輸入算出輸出的主要運算。

__call__ 會在呼叫 forward 的前後做一些重要工作。

概念上可以理解成 y = model(x) 的背後原理大致上是 y = model.__call__(x),而 __call__ 裡面會再做 result = self.forward(x) 這個運算。

因此真正的模型架構仍然寫在 forward() 裡,但實際使用模型時要透過 model(x) 進入。

為什麼不直接寫 model.forward(x)

書中給出下例做對比:

1
2
y = model(x)          # 正確
y = model.forward(x) # 不建議

直接呼叫 forward() 在技術上可能會產生和 __call__() 一樣的輸出,但不應該在使用者的程式碼裡面這樣做。

原因是 __call__() 不只是單純轉呼叫 forward()

書中在此簡化展示了 Module.__call__() 的底層代碼邏輯來對這個問題做說明,它會先執行 forward pre-hooks,接著呼叫 forward(),然後再執行 forward hooks,簡化流程如下:

1
2
3
4
5
6
7
8
9
10
def __call__(self, *input, **kwargs):
for hook in self._forward_pre_hooks.values():
hook(self, input)

result = self.forward(*input, **kwargs)

for hook in self._forward_hooks.values():
hook_result = hook(self, input, result)

return result

model(x) 的完整流程用簡單的條列呈現如下:

  1. model.__call__(x)
  2. 執行 forward 前的 hook
  3. 呼叫 model.forward(x)
  4. 執行 forward 後的 hook
  5. 回傳結果

所以說如果直接寫 model.forward(x),流程就變成直接跳到 model.forward(x),變得只做 forward 運算,而跳過 __call__ 裡面的其他機制。

所以 forward() 要怎麼用?

正確的寫法是定義 forward(),而不要直接呼叫 forward(),範例如下:

1
2
3
4
5
6
7
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(1, 1)

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

之後再用這個 class 建 model:

1
2
model = MyModel()
y = model(x)

,就不要再去寫 y = model.forward(x) 了。

小結

  • forward():寫模型運算邏輯的地方。
  • model(x):實際使用模型的入口。

可以將 forward() 想成餐廳廚房裡「煮菜」的動作,而 __call__() 是「點餐的流程」。

如果走正常的流程(像是客人點餐 → 櫃台記錄 → 廚房煮菜 → 櫃台出餐 → 系統留下紀錄)就不會有太大問題(如同正常寫 model(x))。但直接衝進廚房叫廚師趕快把菜煮好,菜可能一樣會出來,但櫃台沒有紀錄、優惠沒有套用、訂單狀態沒有更新、後續統計也可能漏掉(如同寫 model.forward(x)),輸出看起來可能沒錯,但 PyTorch 的內部機制可能會被繞過。

6.2.2 Returning to the linear model

這一節的目的是把前面第 5 章手寫的線性模型,改寫成 PyTorch nn 模組的標準形式,前面寫過以下這樣的程式碼:

1
2
def model(t_u, w, b):
return w * t_u + b

其中 wb 是我們自己建立的參數,但到了這一節開始改用 PyTorch 提供的 nn.Linear 來表示同一件事。

重點是模型本質沒有變,仍然是線性模型,改變的是工程上的寫法,從「手動管理參數」轉成「由 nn.Module 幫你管理參數」。

nn.Linear(1, 1) 的意思

接著書中撰寫了以下範例:

1
2
3
4
import torch.nn as nn

linear_model = nn.Linear(1, 1)
linear_model(t_un_val)

nn.Linear 的建構子主要接收三個參數:

  1. 輸入特徵數
  2. 輸出特徵數
  3. 是否使用 bias;其中 bias 預設為 True。

因此 nn.Linear(1, 1) 的意思是模型吃進 1 個輸入特徵,輸出 1 個預測值,且預設包含 bias。

另外這樣寫也等同於 y=wx+by = wx + b (在 1 個輸入特徵,輸出 1 個預測值,且包含 bias 這樣的條件下才成立)。

input feature(輸入特徵)與 output feature(輸出特徵)?

在上述例子裡,feature 數量就是 module 輸入 tensor 與輸出 tensor 的大小,所以是 1 和 1(一個 input feature 跟一個 output feature)。

如果輸入同時使用「溫度」與「氣壓」這兩個特徵,則輸入特徵(input feature)數就會是 2;而輸出如果仍然只預測一個數值,輸出特徵數(output feature)就是 1。

nn.Linear(1, 1) 代表有一個 input feature,跟一個 output feature。

如果改成 nn.Linear(2, 1) 的話就是有兩個 input feature 跟一個 output feature。

nn.Linear(1, 1) 裡面有什麼參數?

由於這個模型只有 1 個輸入特徵、1 個輸出特徵,所以它只需要一個 weight 和一個 bias,用 Python 表示如下:

1
2
linear_model.weight
linear_model.bias

觀察模型參數,會看到 weight 是一個 Parameter,bias 也是一個 Parameter,且都有 requires_grad=True

如果用 jupyter notebook 下去跑的話可能會長這樣(這裡沿用書本範例):

1
2
3
4
5
linear_model.weight
# tensor([[-0.0674]], requires_grad=True)

linear_model.bias
# tensor([0.7488], requires_grad=True)

這個意思是表示目前模型長成這樣:$$y = -0.0674x + 0.7488$$

這目前還只是初始化後的隨機值(指 weightbias),還沒有經過訓練,訓練的目的就是讓這個 weightbias 慢慢調整,使模型輸出接近真實答案。

另外要特別注意:這裡的 Parameter 不是普通 tensor,它是 PyTorch 用來表示「模型可訓練參數」的特殊 tensor,只要它被放在 nn.Module 裡,PyTorch 就能在訓練時追蹤它、計算梯度,並交給 optimizer 更新。(簡單來說就是基本上都交給 PyTorch 去跑了)

單筆輸入的情形

1
2
x = torch.ones(1)
linear_model(x)

Output:

1
tensor([0.6814], grad_fn=<AddBackward0>)

這段程式可執行,但作者提醒雖然 PyTorch 允許這樣寫,但這其實不是最標準的輸入維度。

因為 nn.Module 和它的子類別設計上通常要一次處理多筆 samples,因此輸入 tensor 的第 0 維應該代表 batch size。

也就是說,標準的輸入形狀不應像上述例子寫的 tensor([0.6814], grad_fn=<AddBackward0>),而要像是 torch.Size([B, 1]) 這樣的格式。

其中 B 是 batch size,代表一次丟進模型的資料筆數。

Batch input:一次處理多筆資料

在此節相當重要,因為 PyTorch 的神經網路層幾乎都預設支援 batch。

nn 裡的任何 module 都是設計來一次處理多筆輸入的,如果要讓 nn.Linear 一次處理 10 筆資料,就可以建立一個大小為 B×NinB \times N_{in} 的輸入 tensor,其中 BB 是 batch size,NinN_{in} 是輸入特徵數。

例如下例:

1
2
x = torch.ones(10, 1)
linear_model(x)

這裡的 shape 是 [10, 1],所代表的意思是 10 筆資料,且每筆資料有 1 個特徵。

而輸出結果為 [10, 1],意思是 10 筆資料,且每筆資料輸出 1 個預測值。

因此 nn.Linear(1, 1) 的輸入輸出關係可以寫成如下:

  • 輸入 shape:[B, 1]
  • 輸出 shape:[B, 1]

如果寫成一般式,像 nn.Linear(Nin, Nout) 的話,輸出輸入關係就會成這樣:

  • 輸入 shape:[B, Nin]
  • 輸出 shape:[B, Nout]

batch 概念也適用於影像

【Pytorch 深度學習筆記】The PyTorch nn module - batch 概念也適用於影像

Three RGB images batched together and fed into a neural network. The output is a batch of three vectors of size 4. Reproduced from Eli Stevens, Luca Antiga, and Thomas Viehmann, Deep Learning with PyTorch, Fig. 6.3, Manning Publications, 2020.

書中 Figure 6.7 用影像資料說明 batch,對 RGB 圖片來說,一張圖片通常有 channel、height、width,因此一批圖片的輸入形狀會是 B×C×H×WB \times C \times H \times W

其中:

  • BB 代表 batch size。
  • CC 代表 channel 數量。(例如 RGB 是 3,ARGB 是 4)
  • HH 代表 Height,高度。
  • WW 代表 Width,寬度。

圖中示範 batch size 為 3,也就是一次把 3 張 RGB 圖片送進神經網路,最後輸出 3 個向量,每個向量大小為 4。

這張圖想表達的意思是,無論是簡單一維溫度資料,還是複雜的影像資料,PyTorch 都習慣把第 0 維當作 batch 維度。

基本上就一個重點:模型並非只處理一筆資料,而是處理一批資料。

為什麼要用 batch?

使用 batch 有很多原因,第一個原因是硬體效率,尤其是 GPU。

GPU 是高度平行化的硬體,如果一次只丟一筆資料,很多運算單元會沒利用到,在那邊站著發呆,如果一次丟入一批資料時,計算可以分散到更多運算單元上,因此效率更高。

第二個原因是有些進階模型會使用整個 batch 的統計資訊,而 batch 越大,統計估計通常越穩定。

像是如果要影印 100 張講義,通常不會每次只放一張進影印機、等它印完再放下一張,而是會把一疊紙一起放進去,讓機器連續處理。因此 batch 的概念跟影印機的概念類似,一次處理一批資料,硬體利用率更高。

batch 優化

1
2
3
4
5
t_c = [0.5, 14.0, 15.0, 28.0, 11.0, 8.0, 3.0, -4.0, 6.0, 13.0, 21.0]
t_u = [35.7, 55.9, 58.2, 81.9, 56.3, 48.9, 33.9, 21.8, 48.4, 60.4, 68.4]
t_c = torch.tensor(t_c).unsqueeze(1)
t_u = torch.tensor(t_u).unsqueeze(1)
t_u.shape

Output:

1
torch.Size([11, 1])

這是之前第 5 章寫過的溫度資料範例,t_ut_c 原本是 1D tensor,其大小為 torch.Size([11]),即 11 筆資料排成一列,前面手寫模型時可以這樣做,是因為用的是 w * x + b 這種線性模型。其中 wb 是 scalar(純量),PyTorch 可以靠 broadcasting 自動對整個 1D tensor 做運算。

這種寫法之所以能成立,是因為只有單一的輸入特徵,但如果有兩個輸入特徵,就必須多一個維度,把資料整理成 row 是樣本、column 是特徵的矩陣。

所以如果要換到 nn.Linear 時,需要把 torch.Size([11]) 改成 torch.Size([11, 1]) 這樣的形式,在書中用 .unsqueeze(1) 這個方法來做。

:::info
.unsqueeze(1) 的 1 是指在內層多增加一個維度的意思,反之 .unsqueeze(0) 就是在外層。
:::

nn.Linear 取代手寫模型

講完這麼多,終於可以來用 PyTorch 裡面的 nn module 了。

接著書中更新訓練程式,將原本手寫的線性模型改成 linear_model = nn.Linear(1, 1),然後建立 optimizer:

1
2
3
4
optimizer = optim.SGD(
linear_model.parameters(),
lr=1e-2
)

這裡跟以前不一樣的地方在於,以前我們要自己建立參數,例如 params = torch.tensor([1.0, 0.0], requires_grad=True),然後把 params 傳給 optimizer,而現在參數已由 linear_model 持有,所以只要呼叫 linear_model.parameters() 這行就可以取得模型中所有需要訓練的參數。

這就是 nn.Module 的設計,它讓模型除了可以拿來做運算函式以外,同時也是一個參數的容器。

parameters()

1
linear_model.parameters()

這個結果會得到一個 generator,轉成 list 後,可看到裡面有兩個 Parameter:一個是 weight,一個是 bias,可能的結果如下:

1
2
3
4
5
6
7
[
Parameter containing:
tensor([[0.7398]], requires_grad=True),

Parameter containing:
tensor([0.7974], requires_grad=True)
]

這表示 optimizer 可直接拿到所有需要更新的參數。

另外書中也說 parameters() 會遞迴進入 module 的 submodules,回傳所有找到的參數,並整理成平坦的列表。

這對於之後較複雜的神經網路來說很重要,因為一個模型可能包含很多層,例如 Linear、Conv2d、BatchNorm 等,只要它們是 nn.Module 的子模組,model.parameters() 就能自動收集全部可訓練參數。

optimizer 如何更新 nn.Linear 的參數?

optimizer 取得的是一串 requires_grad=True 的 tensor,而所有 Parameter 預設都是 requires_grad=True,因為它們本來就是要拿來做 gradient descent 最佳化的參數。

  • 當呼叫 training_loss.backward() 的時候,PyTorch 會從 loss 反向傳播,計算每個參數的梯度,並把梯度累積在參數的 .grad 屬性裡。
  • 接著再呼叫 optimizer.step(),SGD optimizer 會走訪每一個 Parameter,根據它的 .grad 內容更新參數值。

以上就是 optimizer 如何更新 nn.Linear 參數的流程,整理一下流程:

  1. model(t_u_train):計算 prediction。
  2. loss_fn(prediction, t_c_train):得到 loss。
  3. loss.backward():計算 weightbias 的梯度。
  4. optimizer.step():更新 weightbias

訓練迴圈的變化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def training_loop(n_epochs, optimizer, model, loss_fn,
t_u_train, t_u_val, t_c_train, t_c_val):
for epoch in range(1, n_epochs + 1):
t_p_train = model(t_u_train)
loss_train = loss_fn(t_p_train, t_c_train)

t_p_val = model(t_u_val)
loss_val = loss_fn(t_p_val, t_c_val)

optimizer.zero_grad()
loss_train.backward()
optimizer.step()

if epoch == 1 or epoch % 1000 == 0:
print(f"Epoch {epoch}, "
f"Training loss {loss_train.item():.4f}, "
f"Validation loss {loss_val.item():.4f}")

這段程式和前面手寫版本幾乎完全一樣,但差別在於以前呼叫模型時可能要傳入 params,像是 t_p_train = model(t_u_train, *params),而現在不用了,因為 linear_model 自己已經有 weightbias,所以直接寫 t_p_train = model(t_u_train) 即可。

nn.MSELoss() 取代手寫 loss function

最後,書中也把 loss function 改成 PyTorch 的內建版本,前面已經手寫過這樣的 loss function 了:

1
2
3
def loss_fn(t_p, t_c):
squared_diffs = (t_p - t_c) ** 2
return squared_diffs.mean()

這就是 mean square error,也就是 MSE。由於 PyTorch 提供了 nn.MSELoss() 這樣的函式讓你用,為何不用呢?對吧。

另外需注意 loss function 在 nn 裡也是 nn.Module 的子類別,所以要先建立 instance,再像函式一樣呼叫。

最後結合這個東西就可寫出一個相當簡潔的訓練程式碼了:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
linear_model = nn.Linear(1, 1)
optimizer = optim.SGD(linear_model.parameters(), lr=1e-2)

training_loop(
n_epochs=3000,
optimizer=optimizer,
model=linear_model,
loss_fn=nn.MSELoss(),
t_u_train=t_un_train,
t_u_val=t_un_val,
t_c_train=t_c_train,
t_c_val=t_c_val
)

print()
print(linear_model.weight)
print(linear_model.bias)

Output:

1
2
3
4
5
6
7
8
Epoch 1, Training loss 134.9599, Validation loss 183.1707
Epoch 1000, Training loss 4.8053, Validation loss 4.7307
Epoch 2000, Training loss 3.0285, Validation loss 3.0889
Epoch 3000, Training loss 2.8569, Validation loss 3.9105
Parameter containing:
tensor([[5.4319]], requires_grad=True)
Parameter containing:
tensor([-17.9693], requires_grad=True)

最終的參數模型是 tc5.4319,tu17.9693t_c \approx 5.4319, t_u \approx -17.9693,表示模型最終學到的參數數值是這樣。

總整理

nn.Module 基礎架構

在 PyTorch 中,建構網路的基礎區塊稱為模組(Modules),等同於其他框架的層(Layers),所有自訂的網路架構都必須繼承自 nn.Module

兩大核心屬性:

  1. Parameters(參數):如 weight 與 bias。本質上是特殊的 Tensor,其 requires_grad=True,會在訓練過程中被 Optimizer 不斷更新。
  2. Submodules(子模組):模組內可包裝其他 nn.Module,父模組會自動追蹤管理內部所有的參數。

須注意:絕對不要將子模組裝在標準 Python 的 list 或 dict 中,會導致 Optimizer 找不到參數,如有動態管理需求,則務必使用 PyTorch 專用的容器,nn.ModuleListnn.ModuleDict

模型執行:__call__ vs forward

  • forward(self, x):寫死模型運算邏輯的地方(前向傳播)。
  • model(x):實際使用模型的入口,PyTorch 實作了 __call__ 方法,呼叫 model(x) 會執行 forward(),也會在前後觸發各種 Hooks 機制。
  • 規範:永遠寫 y = model(x),嚴禁直接呼叫 y = model.forward(x),以免繞過系統內部的追蹤與狀態更新機制。

批次處理(Batch)的設計

PyTorch 的神經網路層預設支援並期望一次處理一批資料,而非單筆資料。

  • 維度規則:輸入 Tensor 的第 0 維(Dimension 0)一律代表 Batch Size(BB)。
    • 線性層 nn.Linear(Nin, Nout):輸入 Shape 為 [B, Nin],輸出 Shape 為 [B, Nout]
    • 影像資料:輸入 Shape 通常為 B×C×H×WB \times C \times H \times W(Batch Size, Channels, Height, Width)。
  • 為何要用 Batch?
    1. 硬體效率:極大化 GPU 的高度平行運算效能,避免硬體閒置。
    2. 收斂穩定:使用整個 Batch 的統計資訊來計算梯度,方向會比單筆資料更準確穩定。
  • 實作技巧:如果只有單一特徵的一維資料 [11],必須透過 .unsqueeze(1) 將其轉換成合法的矩陣形狀 [11, 1] 才能送入模型。

實戰應用:nn.Linear 與訓練迴圈

nn.Linear(in_features, out_features) 封裝了仿射變換(Affine Transformation,即 y=wx+by = wx + b),並自動生成對應維度的 weight 與 bias 參數。

  • parameters() 方法:呼叫 model.parameters() 會以遞迴方式搜集整個網路(包含所有子模組)中需要訓練的參數,直接將這個列表交給 Optimizer(如 optim.SGD)即可。
  • 內建 Loss Function:手寫的均方誤差可以完全用內建模組 nn.MSELoss() 取代。
  • 標準訓練 4 步驟(Training Loop):
    1. 前向傳播:prediction = model(x)
    2. 計算誤差:loss = loss_fn(prediction, target)
    3. 反向傳播:optimizer.zero_grad() 後執行 loss.backward()(計算梯度並存入 .grad
    4. 更新參數:optimizer.step()