PyTorch 的 torch.nn 模塊，是構建神經網絡最核心的模塊。它提供了模型基類、常用網絡層、激活函數、損失函數、容器結構、參數管理和訓練狀態控制等能力。

簡單地說，torch.nn 模塊回答的是：神經網絡在 PyTorch 中如何表示、如何組織、如何完成前向計算，以及模型參數如何被自動管理。

如果說 torch 模塊負責張量與基礎計算，那么 torch.nn 模塊就負責把這些張量計算組織成“模型”。在實際深度學習項目中，我們很少只寫零散的矩陣乘法，而是會把線性層、卷積層、激活函數、歸一化層、Dropout 層和損失函數組合成一個完整的神經網絡。

一、認識 torch.nn 模塊

torch.nn 是 PyTorch 中專門用于神經網絡建模的模塊。它并不是只提供幾個網絡層，而是圍繞模型構建提供了一套完整機制。

圖 1：torch.nn 模塊在 PyTorch 訓練流程中的位置

在 PyTorch 中，一個神經網絡通常由以下部分組成：

? 模型基類：nn.Module

? 網絡層：如 nn.Linear、nn.Conv2d、nn.Embedding

? 激活函數：如 nn.ReLU、nn.Sigmoid、nn.GELU

? 損失函數：如 nn.CrossEntropyLoss、nn.MSELoss

? 容器結構：如 nn.Sequential、nn.ModuleList

? 參數管理：如 parameters()、state_dict()

? 訓練與推理狀態：如 train()、eval()

一個最簡單的使用方式如下：

輸出形狀為：

torch.Size([2, 3])

這里可以看出，nn.Linear 不只是一個普通函數，而是一個帶有可學習參數的模塊。它內部保存了權重和偏置，并會在訓練過程中被更新。

二、nn.Module：所有神經網絡模塊的基礎

在 PyTorch 中，大多數神經網絡模型都會繼承 nn.Module。它是模型組織的基礎。

圖 2：nn.Module 的基本結構

一個典型模型通常包含兩個部分：

? __init__()：定義模型中需要用到的層

? forward()：定義數據如何從輸入流向輸出

例如：

這里最重要的不是代碼長度，而是理解模型的組織方式：

? __init__() 負責“有哪些層”

? forward() 負責“數據怎么流動”

? model(x) 會自動調用 forward(x)

? 模型中的層會被自動注冊，參數也會被自動管理

這就是 PyTorch 模型定義的基本范式。

三、forward()：定義前向傳播邏輯

forward() 方法決定輸入張量如何一步步變成輸出張量。它體現的是模型的計算路徑。

例如，一個兩層神經網絡可以理解為：

→ 輸出 logits

對應代碼是：

在分類任務中，模型最后輸出的通常不是直接類別，而是一組分數，常稱為 。

例如：

如果 logits.shape 是 (8, 3)，可以理解為：

? 有 8 個樣本

? 每個樣本對應 3 個類別分數

? 分數越大，模型越傾向于該類別

需要注意的是，很多分類模型的最后一層不需要手動加 。如果后面使用 nn.CrossEntropyLoss()，它會在內部處理與分類概率相關的計算。因此，模型通常直接輸出原始分數即可。

四、nn.Linear：全連接層

nn.Linear 是最基礎、最常用的神經網絡層之一。它常用于多層感知機、分類頭、回歸頭和 Transformer 中的前饋網絡部分。

（也稱為“線性層”）的核心計算可以理解為：

其中：

? X 表示輸入張量

? W 表示權重矩陣

? b 表示偏置

? Y 表示輸出張量

代碼示例：

需要注意的是，nn.Linear(4, 2) 中的 4 和 2 分別表示：

? in_features=4：每個樣本輸入特征數

? out_features=2：每個樣本輸出特征數

它并不表示有 4 個樣本或 2 個樣本。樣本數量通常由輸入張量的第 0 維決定。

五、激活函數：引入非線性能力

如果神經網絡只有線性層，那么多層線性變換疊加后，本質上仍然可以合并成一個線性變換。為了讓模型表達更復雜的非線性關系，需要引入。

常見激活函數包括：

? nn.ReLU

? nn.Sigmoid

? nn.Tanh

? nn.LeakyReLU

? nn.GELU

? nn.Softmax

示例：

在現代深度學習中， 和它的變體非常常見；在 Transformer、BERT、GPT 等模型中， 也經常出現。

需要注意的是，激活函數并不是隨便添加的。它通常位于線性層或卷積層之后，用于增加模型的非線性表達能力。

六、nn.Sequential：按順序組織模型

如果模型結構是簡單的“從前到后依次執行”，可以使用 nn.Sequential 簡化代碼。

例如，下面這個模型與前面手寫 forward() 的兩層網絡等價：

nn.Sequential 的優點是簡潔，適合線性堆疊結構。例如：

? 多層感知機

? 簡單卷積網絡

? 分類頭

? 回歸頭

? 特征映射模塊

但如果模型中包含復雜分支、跳躍連接、多輸入多輸出或條件邏輯，nn.Sequential 就不夠靈活。這時更適合繼承 nn.Module 并手寫 forward()。

七、ModuleList 與 ModuleDict：保存多個子模塊

有些模型中，層的數量可能比較多，或者需要在循環中組織多個子模塊。這時可以使用 nn.ModuleList 或 nn.ModuleDict。

1、ModuleList：列表形式保存子模塊

ModuleList 的關鍵作用是：讓 PyTorch 能夠正確識別并管理其中的子模塊參數。

不要簡單地用普通 Python 列表保存網絡層。普通列表中的層可能不會被模型正確注冊，從而導致參數無法被優化器更新。

2、ModuleDict：字典形式保存子模塊

ModuleDict 適合需要按名稱選擇不同子模塊的場景，例如多任務學習、多個輸出頭、多種特征處理分支等。

八、損失函數：衡量預測與目標的差異

神經網絡訓練不僅需要模型輸出，還需要。損失函數用于衡量模型預測結果與真實目標之間的差異。

常見損失函數包括：

? nn.CrossEntropyLoss：多分類任務常用

? nn.BCELoss：二分類概率輸出場景

? nn.BCEWithLogitsLoss：二分類或多標簽任務常用

? nn.MSELoss：回歸任務常用

? nn.L1Loss：平均絕對誤差損失

? nn.NLLLoss：負對數似然損失

1、分類任務中的 CrossEntropyLoss

nn.CrossEntropyLoss 常用于多分類任務。它接收的是模型輸出的原始分數 logits，以及真實類別編號 target。

這里需要特別注意：

? logits 的形狀通常是 (batch_size, num_classes)

? target 的形狀通常是 (batch_size,)

? target 中保存的是類別編號，而不是 one-hot 向量；target 應是整數類型張量，常見類型是 torch.long

需要注意的是，CrossEntropyLoss 內部已經包含了與 Softmax 相關的計算，因此模型最后一層通常直接輸出 logits，不需要手動添加 Softmax。

下面這種寫法更符合 PyTorch 多分類任務的常見形式：

這是初學者非常容易出錯的地方。如果提前對 logits 做 Softmax，或者把 target 寫成 one-hot 向量，就可能導致損失計算語義錯誤，進而影響訓練效果。

2、回歸任務中的 MSELoss

回歸任務中，預測值和真實值通常都使用浮點張量，并且形狀應盡量保持一致。

例如，如果 pred 的形狀是 (8, 1)，而 target 的形狀是 (8,)，有時可能會觸發廣播，導致計算含義不符合預期。因此，回歸任務中要特別注意預測值和目標值的形狀。

九、卷積層、池化層與圖像特征提取

在圖像任務中，常用的神經網絡層包括、、和。

典型組合可以寫成：

這里的結構可以理解為：

? 卷積層負責提取局部特征

? 激活函數引入非線性

? 池化層降低空間尺寸

? 展平操作把特征圖變成向量

? 全連接層輸出分類結果

在卷積網絡中，輸入通常采用 (N, C, H, W) 格式。如果圖像數據是 (N, H, W, C)，需要先調整維度順序。

十、Dropout 與歸一化層

在神經網絡訓練中，除了線性層和卷積層，還經常使用 與來改善訓練效果。

1、Dropout：隨機失活

Dropout 常用于緩解過擬合。它在訓練階段隨機丟棄部分神經元輸出，使模型不要過度依賴某些局部特征。

需要注意的是，Dropout 在訓練模式和推理模式下行為不同。因此，訓練時要調用 model.train()，推理或驗證時要調用 model.eval()。

2、BatchNorm 與 LayerNorm

常見歸一化層包括：

? nn.BatchNorm1d

? nn.BatchNorm2d

? nn.LayerNorm

示例：

歸一化層的作用不是改變任務目標，而是讓中間特征的數值分布更穩定，從而幫助模型更順利地訓練。

十一、Embedding、RNN 與 Transformer 相關層

torch.nn 不只用于普通前饋網絡和卷積網絡，也提供了文本、序列和注意力模型相關的常用層。

1、Embedding：把離散編號映射為向量

nn.Embedding 常用于自然語言處理和推薦系統。它可以把詞編號、用戶編號、物品編號等離散 ID 映射為連續向量。

示例：

這里可以理解為：

? 輸入是整數編號

? 輸出是連續向量

? 每個編號都有一個可學習的向量表示

2、RNN、LSTM 與 GRU：處理序列數據的循環網絡結構

RNN、LSTM 與 GRU 都屬于處理序列數據的循環神經網絡結構，常用于文本、語音、時間序列等任務。

RNN（循環神經網絡）是基礎循環結構，（長短期記憶網絡）和 （門控循環單元）是對普通 RNN 的改進版本，目的都是讓模型更好地處理序列中的上下文信息。

示例：

當 batch_first=True 時，輸入形狀通常是：

(batch_size, sequence_length, input_size)

序列模型中最容易混淆的是維度順序。寫代碼時要特別確認：

? 哪一維是批次

? 哪一維是時間步

? 哪一維是特征

3、MultiheadAttention 與 Transformer

nn.MultiheadAttention 和 nn.Transformer 可用于構建注意力機制與 Transformer 結構。

這里重點是理解：

? 輸入通常是序列張量

? 每個位置會與其他位置建立關系

? 輸出形狀通常與輸入主結構保持一致

Transformer 結構較復雜，初學階段不必急于展開所有細節。先理解 Embedding、序列張量形狀和注意力輸入輸出關系，會更穩妥。

十二、參數管理：parameters() 與 state_dict()

神經網絡訓練的本質，是不斷更新模型中的可學習參數。nn.Module 會自動管理模型中的參數。

1、查看模型參數

輸出結果類似：

2.bias torch.Size([3])

這里可以看到，ReLU 沒有可學習參數，而 Linear 有權重和偏置。

2、parameters()：交給優化器更新

這也是為什么模型層必須正確注冊到 nn.Module 中。如果某些層沒有被注冊，優化器可能無法找到它們的參數。

3、state_dict()：保存模型參數

加載參數通常寫成：

state_dict() 是 PyTorch 中非常常用的模型保存方式。它保存的是參數，而不是完整訓練代碼。因此，加載時通常需要先重新定義相同結構的模型。

十三、train() 與 eval()：訓練模式和推理模式

某些層在訓練和推理時行為不同，例如：

? Dropout：訓練時隨機丟棄，推理時不丟棄

? BatchNorm：訓練時使用當前批次統計量，推理時使用累計統計量

因此，在訓練和驗證時要明確切換模式。

需要注意：

? model.train() 不會自動訓練模型，只是切換為訓練模式

? model.eval() 不會停止梯度計算，只是切換為推理模式

? 推理時通常還要使用 torch.no_grad() 關閉梯度記錄

十四、一個完整示例：用 nn 構建并訓練簡單分類模型

圖 3：torch.nn 模型訓練閉環

下面用一個完整示例，把 torch.nn、損失函數和優化器串起來。

這個示例體現了神經網絡訓練的基本閉環：

輸入數據 → 模型前向計算 → 計算損失 → 清空梯度 → 反向傳播 → 參數更新

其中，torch.nn 主要負責模型結構和損失函數；優化器負責根據梯度更新參數。

十五、使用 torch.nn 時應注意的問題

1、不要忘記繼承 nn.Module 并調用 super().__init__()

自定義模型時，通常要寫：

如果忘記調用 super().__init__()，子模塊和參數可能無法被正常注冊，后續訓練會出現問題。

2、網絡層應定義在 init() 中

推薦寫法：

不推薦在 forward() 中反復新建帶參數的層：

原因是：訓練需要持續更新同一組參數。如果每次前向傳播都創建新層，模型參數就無法穩定學習。

3、注意輸入輸出形狀

神經網絡層對輸入形狀有明確要求。例如：

? nn.Linear 通常關注輸入最后一維

? nn.Conv2d 通常要求輸入為 (N, C, H, W)

? nn.LSTM 在 batch_first=True 時常用 (N, T, D)

? nn.CrossEntropyLoss 常用 (N, C) 的預測值和 (N,) 的類別標簽

很多錯誤都可以通過打印形狀定位：

print(target.shape)

4、區分 logits、概率和類別

分類模型中，常見三個概念：

? logits：模型輸出的原始分數

? 概率：經過 Softmax 后得到的歸一化結果

? 類別：分數最大或概率最大的類別編號

例如：

訓練時使用 CrossEntropyLoss，通常直接輸入 logits，不需要先手動計算 Softmax。

5、訓練與推理模式要切換

涉及 Dropout、BatchNorm 等層時，訓練和推理行為不同。

訓練階段：

model.train()

驗證或推理階段：

    output = model(x)

這兩個操作不應省略。否則，驗證結果可能不穩定，推理結果也可能受到訓練行為影響。

6、用 ModuleList 保存多個層，而不是普通列表

如果多個層需要放在列表里，推薦使用：

不要簡單寫成：

普通列表中的層可能不會被正確注冊，優化器也可能無法更新其中的參數。

7、損失函數要匹配任務類型

不同任務應該使用不同損失函數：

? 多分類任務：常用 CrossEntropyLoss

? 二分類任務：常用 BCEWithLogitsLoss

? 回歸任務：常用 MSELoss 或 L1Loss

? 多標簽任務：常用 BCEWithLogitsLoss

損失函數選錯，模型即使能運行，也可能學不到正確目標。

小結

torch.nn 模塊負責神經網絡模型的構建與組織，是從張量計算走向深度學習建模的關鍵模塊。學習這一模塊，重點不只是記住層名稱，而是理解 nn.Module 如何管理層和參數，forward() 如何定義數據流，損失函數如何連接預測與目標，以及訓練模式和推理模式為何必須區分。

“點贊有美意，贊賞是鼓勵”