PyTorch 是 Python 生態(tài)中最常用的深度學(xué)習(xí)框架之一。它并不是只提供若干神經(jīng)網(wǎng)絡(luò)層，也不是簡(jiǎn)單的 GPU 計(jì)算工具，而是圍繞深度學(xué)習(xí)的完整流程，組織出一組相互配合的模塊：張量計(jì)算、自動(dòng)求導(dǎo)、神經(jīng)網(wǎng)絡(luò)建模、數(shù)據(jù)加載、優(yōu)化器、設(shè)備加速、模型保存、分布式訓(xùn)練與部署擴(kuò)展等。

在學(xué)習(xí) PyTorch 的過(guò)程中，如果一時(shí)記不清有哪些頂層名稱，可以直接在 Python 中查看：

print(dir(torch))

輸出中通常可以看到許多頂層名稱，例如：Tensor、nn、autograd、optim、utils、cuda、mps、amp、linalg、fft、sparse、distributed、jit、export、profiler 等。需要注意的是，dir(torch) 顯示的是頂層名稱集合，并不等于學(xué)習(xí)時(shí)必須逐一掌握的模塊清單。

從整體上看，可以把 PyTorch 的主要模塊理解為三類分工：

第一類是核心訓(xùn)練工作流模塊，負(fù)責(zé)張量表示、自動(dòng)求導(dǎo)、模型定義、數(shù)據(jù)加載、參數(shù)優(yōu)化和設(shè)備加速。

第二類是模型構(gòu)建與計(jì)算支持模塊，負(fù)責(zé)神經(jīng)網(wǎng)絡(luò)層、函數(shù)式接口、線性代數(shù)、概率分布、稀疏張量、傅里葉變換等計(jì)算能力。

第三類是工程擴(kuò)展與生態(tài)模塊，負(fù)責(zé)編譯加速、模型導(dǎo)出、分布式訓(xùn)練、性能分析、日志可視化和領(lǐng)域擴(kuò)展。

這種模塊化設(shè)計(jì)使得 PyTorch 不只是“神經(jīng)網(wǎng)絡(luò)層的集合”，而是一個(gè)圍繞深度學(xué)習(xí)訓(xùn)練與部署組織起來(lái)的完整計(jì)算框架。

一、第一層：核心訓(xùn)練工作流模塊

這一層模塊直接對(duì)應(yīng)深度學(xué)習(xí)最常見(jiàn)的訓(xùn)練流程：準(zhǔn)備數(shù)據(jù)、定義模型、前向計(jì)算、計(jì)算損失、反向傳播、更新參數(shù)、保存模型。在初學(xué)階段，這一層最值得優(yōu)先掌握。

1、torch：張量與基礎(chǔ)計(jì)算模塊

torch 是 PyTorch 的核心模塊。它提供張量對(duì)象、數(shù)學(xué)運(yùn)算、數(shù)組變形、索引切片、隨機(jī)數(shù)生成、設(shè)備遷移等基礎(chǔ)能力。可以把它理解為 PyTorch 中最底層、最常用的計(jì)算入口。

張量（Tensor）是 PyTorch 中最基本的數(shù)據(jù)結(jié)構(gòu)。輸入數(shù)據(jù)、模型參數(shù)、中間激活值、梯度、損失值，最終都以張量形式參與計(jì)算。

常用工具：

? torch.tensor()：從 Python 列表、數(shù)值等對(duì)象創(chuàng)建張量

? torch.zeros()：創(chuàng)建全 0 張量

? torch.ones()：創(chuàng)建全 1 張量

? torch.randn()：創(chuàng)建服從標(biāo)準(zhǔn)正態(tài)分布的隨機(jī)張量

? torch.arange()：創(chuàng)建等差序列張量

? torch.cat()：沿指定維度拼接張量

? torch.stack()：把多個(gè)張量堆疊成新維度

? torch.reshape()：改變張量形狀

? torch.matmul()：矩陣乘法

? torch.save()：保存張量、模型參數(shù)或其他對(duì)象

? torch.load()：加載已保存的對(duì)象

示例：

這個(gè)示例說(shuō)明：在 PyTorch 中，很多深度學(xué)習(xí)計(jì)算本質(zhì)上都是張量之間的數(shù)學(xué)運(yùn)算。

2、torch.autograd：自動(dòng)求導(dǎo)模塊

torch.autograd 是 PyTorch 的自動(dòng)求導(dǎo)系統(tǒng)。它負(fù)責(zé)記錄張量運(yùn)算過(guò)程，并根據(jù)計(jì)算圖自動(dòng)計(jì)算梯度。官方文檔將 torch.autograd 描述為用于實(shí)現(xiàn)任意標(biāo)量值函數(shù)自動(dòng)微分的類與函數(shù)集合。

在深度學(xué)習(xí)中，模型訓(xùn)練依賴反向傳播。手動(dòng)推導(dǎo)和實(shí)現(xiàn)每個(gè)參數(shù)的梯度既復(fù)雜又容易出錯(cuò)，而 autograd 可以根據(jù)前向計(jì)算過(guò)程自動(dòng)完成梯度計(jì)算。

常用工具：

? requires_grad=True：表示需要跟蹤該張量的梯度

? loss.backward()：從損失值出發(fā)執(zhí)行反向傳播

? tensor.grad：查看張量對(duì)應(yīng)的梯度

? torch.no_grad()：關(guān)閉梯度記錄，常用于推理階段

? tensor.detach()：從當(dāng)前計(jì)算圖中分離張量

示例：

這個(gè)示例中，PyTorch 會(huì)自動(dòng)計(jì)算 y 關(guān)于 x 的導(dǎo)數(shù)。對(duì)于神經(jīng)網(wǎng)絡(luò)而言，模型參數(shù)通常都會(huì)參與自動(dòng)求導(dǎo)，訓(xùn)練時(shí)不需要手動(dòng)寫(xiě)出每個(gè)參數(shù)的梯度公式。

3、torch.nn：神經(jīng)網(wǎng)絡(luò)模塊

torch.nn 是 PyTorch 中構(gòu)建神經(jīng)網(wǎng)絡(luò)最核心的模塊。它提供模型基類、常用網(wǎng)絡(luò)層、激活函數(shù)、損失函數(shù)、容器結(jié)構(gòu)等內(nèi)容。官方文檔中，torch.nn.Module 是所有神經(jīng)網(wǎng)絡(luò)模塊的基類；PyTorch 也使用模塊來(lái)表示神經(jīng)網(wǎng)絡(luò)。

在 PyTorch 中，通常會(huì)通過(guò)繼承 nn.Module 來(lái)定義模型。一個(gè)模型一般包含兩部分：

? 在 __init__() 中定義網(wǎng)絡(luò)層

? 在 forward() 中定義前向傳播邏輯

常用工具：

? nn.Module：所有神經(jīng)網(wǎng)絡(luò)模型的基類

? nn.Linear：全連接層

? nn.Conv2d：二維卷積層

? nn.MaxPool2d：二維最大池化層

? nn.BatchNorm2d：二維批量歸一化層

? nn.LayerNorm：層歸一化

? nn.Dropout：隨機(jī)失活層

? nn.Embedding：詞嵌入層

? nn.RNN：循環(huán)神經(jīng)網(wǎng)絡(luò)層

? nn.LSTM：長(zhǎng)短期記憶網(wǎng)絡(luò)層

? nn.GRU：門控循環(huán)單元

? nn.Transformer：Transformer 結(jié)構(gòu)模塊

? nn.CrossEntropyLoss：交叉熵?fù)p失

? nn.MSELoss：均方誤差損失

示例：

這個(gè)示例體現(xiàn)了 PyTorch 建模的基本思想：模型是一個(gè) nn.Module，輸入張量經(jīng)過(guò) forward() 得到輸出張量。

4、torch.nn.functional：函數(shù)式神經(jīng)網(wǎng)絡(luò)接口

torch.nn.functional 通常簡(jiǎn)寫(xiě)為 F。它提供許多函數(shù)式操作，例如激活函數(shù)、池化、歸一化、損失函數(shù)等。

torch.nn 與 torch.nn.functional 的區(qū)別可以簡(jiǎn)單理解為：

? torch.nn 更偏“層對(duì)象”，適合定義帶參數(shù)、需要注冊(cè)到模型中的模塊

? torch.nn.functional 更偏“函數(shù)調(diào)用”，適合沒(méi)有參數(shù)或臨時(shí)調(diào)用的計(jì)算操作

常用工具：

? F.relu()：ReLU 激活函數(shù)

? F.softmax()：Softmax 函數(shù)

? F.log_softmax()：對(duì)數(shù) Softmax

? F.cross_entropy()：交叉熵?fù)p失

? F.mse_loss()：均方誤差損失

? F.dropout()：函數(shù)式隨機(jī)失活

? F.conv2d()：函數(shù)式二維卷積

? F.max_pool2d()：函數(shù)式二維最大池化

示例：

這個(gè)示例中，nn.Linear 用于定義帶參數(shù)的線性層，而 F.relu() 用于執(zhí)行無(wú)參數(shù)的激活操作。

5、torch.optim：優(yōu)化器模塊

torch.optim 用于實(shí)現(xiàn)參數(shù)優(yōu)化算法。神經(jīng)網(wǎng)絡(luò)訓(xùn)練的目標(biāo)是讓損失函數(shù)逐步減小，而優(yōu)化器負(fù)責(zé)根據(jù)梯度更新模型參數(shù)。官方文檔將 torch.optim 描述為實(shí)現(xiàn)各種優(yōu)化算法的包，并支持常見(jiàn)優(yōu)化方法。

常用工具：

? optim.SGD：隨機(jī)梯度下降

? optim.Adam：自適應(yīng)矩估計(jì)優(yōu)化器，深度學(xué)習(xí)中非常常用

? optim.AdamW：解耦權(quán)重衰減版本的 Adam，訓(xùn)練 Transformer 類模型時(shí)很常見(jiàn)

? optim.RMSprop：基于平方梯度移動(dòng)平均的優(yōu)化器

? optim.Adadelta：自適應(yīng)學(xué)習(xí)率優(yōu)化器

? optim.lr_scheduler：學(xué)習(xí)率調(diào)度器子模塊

訓(xùn)練時(shí)常見(jiàn)的優(yōu)化步驟是：

optimizer.step()

含義分別是：

? 清空上一輪梯度

? 根據(jù)損失執(zhí)行反向傳播

? 按照優(yōu)化算法更新參數(shù)

示例：

這個(gè)示例展示了 PyTorch 最小訓(xùn)練閉環(huán)：

前向計(jì)算 → 損失計(jì)算 → 反向傳播 → 參數(shù)更新

6、torch.utils.data：數(shù)據(jù)加載模塊

torch.utils.data 主要用于組織數(shù)據(jù)集和批量加載數(shù)據(jù)。深度學(xué)習(xí)訓(xùn)練通常不是一次性把全部數(shù)據(jù)送入模型，而是按批次讀取、打亂、并行加載。官方文檔指出，DataLoader 是 PyTorch 數(shù)據(jù)加載工具的核心，它表示一個(gè)基于數(shù)據(jù)集的 Python 可迭代對(duì)象。

常用工具：

? Dataset：自定義數(shù)據(jù)集的基類

? TensorDataset：把張量包裝成數(shù)據(jù)集

? DataLoader：按批次加載數(shù)據(jù)

? random_split()：隨機(jī)劃分?jǐn)?shù)據(jù)集

? Sampler：控制樣本抽樣方式

? WeightedRandomSampler：按權(quán)重抽樣，常用于類別不平衡任務(wù)

示例：

這個(gè)示例說(shuō)明：Dataset 負(fù)責(zé)表示“有哪些樣本”，DataLoader 負(fù)責(zé)表示“如何按批次讀取樣本”。

7、torch.cuda、torch.mps 與 torch.amp：設(shè)備加速模塊

深度學(xué)習(xí)訓(xùn)練通常依賴硬件加速。PyTorch 可以把張量和模型移動(dòng)到不同設(shè)備上，例如 CPU、NVIDIA GPU、Apple Silicon GPU 等。

常用工具：

? torch.cuda.is_available()：檢查 CUDA 是否可用

? torch.device()：構(gòu)建設(shè)備對(duì)象

? tensor.to(device)：把張量移動(dòng)到指定設(shè)備

? model.to(device)：把模型參數(shù)移動(dòng)到指定設(shè)備

? torch.cuda.empty_cache()：釋放 CUDA 緩存

? torch.mps：用于訪問(wèn) Apple Metal Performance Shaders 后端，適合部分 macOS 設(shè)備

? torch.amp：自動(dòng)混合精度模塊，可讓部分操作使用低精度類型，以提升訓(xùn)練速度并降低顯存占用

示例：

這個(gè)示例體現(xiàn)了 PyTorch 設(shè)備管理的基本方式：模型和數(shù)據(jù)必須放在同一設(shè)備上，才能正常參與計(jì)算。

綜合示例 1：從數(shù)據(jù)到模型訓(xùn)練的最小閉環(huán)

下面這個(gè)示例把第一層中的幾個(gè)核心模塊串起來(lái)：先構(gòu)造數(shù)據(jù)，再加載批次，然后定義模型、損失函數(shù)和優(yōu)化器，最后完成簡(jiǎn)單訓(xùn)練。

這個(gè)示例體現(xiàn)了 PyTorch 工作流的典型順序：

數(shù)據(jù)準(zhǔn)備 → 批量加載 → 模型定義 → 前向計(jì)算 → 損失計(jì)算 → 反向傳播 → 參數(shù)更新

二、第二層：模型構(gòu)建與計(jì)算支持模塊

這一層模塊主要負(fù)責(zé)“模型內(nèi)部如何計(jì)算”。如果說(shuō)第一層解決的是“怎樣完成訓(xùn)練流程”，那么第二層解決的是“模型需要哪些計(jì)算能力”。

1、torch.nn 中的層、容器與損失函數(shù)

torch.nn 不只是提供 nn.Module 基類，還提供了大量常用網(wǎng)絡(luò)層、容器和損失函數(shù)。

常用層：

? nn.Linear：全連接層，常用于 MLP、分類頭和回歸頭

? nn.Conv1d：一維卷積，常用于序列信號(hào)

? nn.Conv2d：二維卷積，常用于圖像任務(wù)

? nn.Conv3d：三維卷積，常用于視頻或體數(shù)據(jù)

? nn.BatchNorm1d、nn.BatchNorm2d：批量歸一化

? nn.LayerNorm：層歸一化，Transformer 中常見(jiàn)

? nn.Dropout：隨機(jī)失活，常用于緩解過(guò)擬合

? nn.Embedding：嵌入層，常用于文本、推薦系統(tǒng)等任務(wù)

? nn.MultiheadAttention：多頭注意力層

常用容器：

? nn.Sequential：按順序組織多個(gè)層

? nn.ModuleList：保存多個(gè)子模塊

? nn.ModuleDict：用字典形式保存多個(gè)子模塊

常用損失函數(shù)：

? nn.CrossEntropyLoss：多分類交叉熵?fù)p失

? nn.BCELoss：二分類交叉熵?fù)p失

? nn.BCEWithLogitsLoss：帶 Logits 的二分類交叉熵?fù)p失

? nn.MSELoss：均方誤差損失

? nn.L1Loss：平均絕對(duì)誤差損失

? nn.NLLLoss：負(fù)對(duì)數(shù)似然損失

2、torch.linalg：線性代數(shù)模塊

torch.linalg 用于線性代數(shù)計(jì)算。它在機(jī)器學(xué)習(xí)、深度學(xué)習(xí)、科學(xué)計(jì)算中都非常重要，尤其適合矩陣分解、范數(shù)、矩陣秩、逆矩陣、特征值等計(jì)算。官方文檔中，torch.linalg 包含矩陣秩、條件數(shù)等典型線性代數(shù)函數(shù)。

常用工具：

? torch.linalg.norm()：計(jì)算向量或矩陣范數(shù)

? torch.linalg.inv()：計(jì)算矩陣逆

? torch.linalg.det()：計(jì)算行列式

? torch.linalg.matrix_rank()：計(jì)算矩陣秩

? torch.linalg.eig()：計(jì)算特征值和特征向量

? torch.linalg.svd()：奇異值分解

? torch.linalg.qr()：QR 分解

示例：

這一模塊更偏數(shù)學(xué)計(jì)算基礎(chǔ)。在實(shí)現(xiàn) PCA、矩陣分解、低秩近似等算法時(shí)，它非常有用。

3、torch.fft：傅里葉變換模塊

torch.fft 用于離散傅里葉變換及相關(guān)計(jì)算。官方文檔將 torch.fft 定位為離散傅里葉變換和相關(guān)函數(shù)集合。

常用工具：

? torch.fft.fft()：一維快速傅里葉變換

? torch.fft.ifft()：一維逆傅里葉變換

? torch.fft.fft2()：二維快速傅里葉變換

? torch.fft.ifft2()：二維逆傅里葉變換

? torch.fft.rfft()：實(shí)值輸入的一維傅里葉變換

? torch.fft.fftfreq()：生成傅里葉頻率

應(yīng)用場(chǎng)景：

? 信號(hào)處理

? 圖像頻域分析

? 科學(xué)計(jì)算

? 某些頻域神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)

4、torch.sparse：稀疏張量模塊

torch.sparse 用于處理稀疏張量。稀疏張量適合表示大多數(shù)元素為 0 的數(shù)據(jù)，例如圖結(jié)構(gòu)、推薦系統(tǒng)中的用戶—物品矩陣、高維文本特征等。官方文檔中，torch.sparse 提供從稠密張量構(gòu)造稀疏張量以及多種稀疏布局支持。

常用工具：

? torch.sparse_coo_tensor()：創(chuàng)建 COO 格式稀疏張量

? torch.sparse_csr_tensor()：創(chuàng)建 CSR 格式稀疏張量

? tensor.to_sparse()：把稠密張量轉(zhuǎn)換為稀疏張量

? tensor.to_dense()：把稀疏張量轉(zhuǎn)換為稠密張量

? torch.sparse.mm()：稀疏矩陣乘法

? torch.sparse.softmax()：稀疏版本 Softmax

示例：

稀疏張量并不是初學(xué)階段的重點(diǎn)，但在圖神經(jīng)網(wǎng)絡(luò)、推薦系統(tǒng)和大規(guī)模特征處理中很重要。

5、torch.distributions：概率分布模塊

torch.distributions 用于表示概率分布、采樣和計(jì)算概率值。它常用于生成模型、強(qiáng)化學(xué)習(xí)、貝葉斯建模和變分推斷等任務(wù)。

常用工具：

? Normal：正態(tài)分布

? Bernoulli：伯努利分布

? Categorical：類別分布

? MultivariateNormal：多元正態(tài)分布

? Uniform：均勻分布

? Distribution.sample()：從分布中采樣

? Distribution.log_prob()：計(jì)算對(duì)數(shù)概率

示例：

在普通監(jiān)督學(xué)習(xí)中，可能很少直接使用這個(gè)模塊；但在 VAE、策略梯度、概率建模中，它非常常見(jiàn)。

6、torch.random：隨機(jī)數(shù)模塊

torch.random 與 PyTorch 的隨機(jī)數(shù)狀態(tài)有關(guān)。深度學(xué)習(xí)實(shí)驗(yàn)通常需要控制隨機(jī)性，例如初始化參數(shù)、打亂數(shù)據(jù)、隨機(jī)增強(qiáng)、Dropout 等。

常用工具：

? torch.manual_seed()：設(shè)置隨機(jī)種子

? torch.initial_seed()：查看初始隨機(jī)種子

? torch.get_rng_state()：獲取隨機(jī)數(shù)生成器狀態(tài)

? torch.set_rng_state()：設(shè)置隨機(jī)數(shù)生成器狀態(tài)

示例：

這個(gè)示例說(shuō)明：設(shè)置相同隨機(jī)種子后，相同隨機(jī)操作可以得到一致結(jié)果。需要注意的是，完整可復(fù)現(xiàn)實(shí)驗(yàn)還可能涉及 CUDA、DataLoader 多進(jìn)程和底層算子確定性等因素。

7、torch.func：函數(shù)變換模塊

torch.func 提供函數(shù)式變換能力，例如自動(dòng)向量化、函數(shù)式梯度、雅可比矩陣、海森矩陣等。它更適合進(jìn)階場(chǎng)景，尤其是元學(xué)習(xí)、高階梯度、批量函數(shù)變換和研究型代碼。

常用工具：

? torch.func.grad()：構(gòu)造梯度函數(shù)

? torch.func.vmap()：自動(dòng)向量化映射

? torch.func.jacrev()：反向模式計(jì)算雅可比矩陣

? torch.func.jacfwd()：前向模式計(jì)算雅可比矩陣

? torch.func.hessian()：計(jì)算海森矩陣

初學(xué)階段只需要知道：torch.func 不是普通訓(xùn)練循環(huán)的必需模塊，而是面向更高級(jí)自動(dòng)微分與函數(shù)變換的工具。

綜合示例 2：定義一個(gè)更標(biāo)準(zhǔn)的訓(xùn)練函數(shù)

下面的示例展示 PyTorch 中較常見(jiàn)的訓(xùn)練組織方式：把單輪訓(xùn)練封裝為函數(shù)，使代碼更接近真實(shí)項(xiàng)目結(jié)構(gòu)。

這段代碼體現(xiàn)了 PyTorch 訓(xùn)練代碼的基本組織方式：模型類、訓(xùn)練函數(shù)、數(shù)據(jù)加載器、損失函數(shù)和優(yōu)化器彼此分工明確。

三、第三層：工程擴(kuò)展與生態(tài)模塊

這一層模塊同樣重要，但專題性更強(qiáng)。初學(xué)階段不必全部展開(kāi)，只需要知道它們分別解決什么問(wèn)題。等到模型規(guī)模變大、訓(xùn)練變慢、需要部署或需要多卡訓(xùn)練時(shí)，再逐步深入。

1、torch.compile 與 torch.compiler：編譯加速模塊

torch.compile() 是 PyTorch 2.x 之后非常重要的加速入口。它可以嘗試把模型計(jì)算圖編譯成更高效的執(zhí)行形式，從而提升訓(xùn)練或推理速度。torch.compiler 則提供與編譯相關(guān)的 API。

常用方式：

model = torch.compile(model)

適用場(chǎng)景：

? 模型訓(xùn)練速度優(yōu)化

? 推理性能優(yōu)化

? 在不大幅改變?cè)写a的情況下嘗試加速

需要注意的是，編譯加速并不保證所有模型都一定更快，也可能受到模型結(jié)構(gòu)、硬件設(shè)備、算子支持和輸入形狀變化的影響。

2、torch.export：模型導(dǎo)出模塊

torch.export 用于捕獲模型的張量計(jì)算圖，并生成可進(jìn)一步處理的導(dǎo)出表示。官方文檔中，torch.export.export() 會(huì)基于示例輸入追蹤模型執(zhí)行，并記錄 PyTorch 操作；相比傳統(tǒng)追蹤方式，它更強(qiáng)調(diào)對(duì)動(dòng)態(tài)形狀和程序語(yǔ)義的處理能力。

常用場(chǎng)景：

? 模型部署前的圖捕獲

? 編譯器后端處理

? 模型轉(zhuǎn)換和進(jìn)一步優(yōu)化

? 與 PyTorch 2.x 編譯生態(tài)配合使用

初學(xué)階段可以先知道：torch.export 更偏工程部署和編譯鏈路，不是普通訓(xùn)練入門的第一重點(diǎn)。

3、torch.jit：TorchScript 模塊

torch.jit 是較早期 PyTorch 中用于模型腳本化和圖模式表示的重要模塊。它可以把部分 PyTorch 模型轉(zhuǎn)換為 TorchScript 形式，以便脫離 Python 環(huán)境執(zhí)行或做進(jìn)一步優(yōu)化。

常用工具：

? torch.jit.trace()：通過(guò)示例輸入追蹤模型執(zhí)行路徑

? torch.jit.script()：把支持的 Python 模型代碼腳本化

? torch.jit.save()：保存 TorchScript 模型

? torch.jit.load()：加載 TorchScript 模型

在當(dāng)前 PyTorch 生態(tài)中，torch.compile、torch.export 等工具越來(lái)越重要，但理解 torch.jit 仍有助于理解 PyTorch 從動(dòng)態(tài)圖到圖表示的發(fā)展脈絡(luò)。

4、torch.distributed：分布式訓(xùn)練模塊

torch.distributed 用于多進(jìn)程、多 GPU、多節(jié)點(diǎn)訓(xùn)練。官方文檔指出，該包為跨多個(gè)計(jì)算節(jié)點(diǎn)運(yùn)行的多進(jìn)程并行提供 PyTorch 支持和通信原語(yǔ)。

常用工具：

? torch.distributed.init_process_group()：初始化分布式進(jìn)程組

? torch.distributed.get_rank()：獲取當(dāng)前進(jìn)程編號(hào)

? torch.distributed.get_world_size()：獲取總進(jìn)程數(shù)

? DistributedDataParallel：分布式數(shù)據(jù)并行訓(xùn)練，通常簡(jiǎn)寫(xiě)為 DDP

? DistributedSampler：分布式數(shù)據(jù)采樣器

典型應(yīng)用場(chǎng)景：

? 單機(jī)多卡訓(xùn)練

? 多機(jī)多卡訓(xùn)練

? 大模型訓(xùn)練

? 大規(guī)模圖像、語(yǔ)音、文本模型訓(xùn)練

對(duì)于初學(xué)者而言，先掌握單機(jī)單卡訓(xùn)練即可；當(dāng)模型和數(shù)據(jù)規(guī)模變大時(shí)，再學(xué)習(xí) DDP 會(huì)更自然。

5、torch.profiler：性能分析模塊

torch.profiler 用于分析模型訓(xùn)練或推理過(guò)程中的性能瓶頸。官方文檔中，torch.profiler 可導(dǎo)出設(shè)備上的內(nèi)存事件信息和時(shí)間線圖等內(nèi)容。

常用用途：

? 分析 CPU 和 GPU 時(shí)間消耗

? 觀察哪些算子最耗時(shí)

? 分析顯存占用

? 配合 TensorBoard 查看性能軌跡

常用工具：

? torch.profiler.profile()：?jiǎn)?dòng)性能分析

? torch.profiler.record_function()：為代碼塊添加分析標(biāo)簽

? ProfilerActivity.CPU：記錄 CPU 活動(dòng)

? ProfilerActivity.CUDA：記錄 CUDA 活動(dòng)

性能分析通常不是寫(xiě)第一個(gè)模型時(shí)需要掌握的內(nèi)容，但在模型訓(xùn)練變慢、顯存占用異常或推理延遲過(guò)高時(shí)非常重要。

6、torch.utils.tensorboard：日志與可視化模塊

torch.utils.tensorboard 用于把訓(xùn)練過(guò)程中的指標(biāo)寫(xiě)入 TensorBoard。例如損失曲線、準(zhǔn)確率曲線、學(xué)習(xí)率變化、圖像樣本、模型結(jié)構(gòu)等。

常用工具：

? SummaryWriter：寫(xiě)入 TensorBoard 日志

? add_scalar()：記錄標(biāo)量指標(biāo)

? add_image()：記錄圖像

? add_histogram()：記錄參數(shù)或梯度分布

? add_graph()：記錄模型圖

示例：

這個(gè)模塊適合在訓(xùn)練過(guò)程中觀察模型是否正常收斂，而不是只在訓(xùn)練結(jié)束后看最終結(jié)果。

7、torch.package 與 torch.hub：模型打包與模型獲取模塊

torch.package 用于創(chuàng)建包含 PyTorch 代碼和模型工件的包，便于保存、共享和加載。官方文檔將其描述為支持創(chuàng)建包含模型制品和任意 PyTorch 代碼的包。

torch.hub 則更偏向從遠(yuǎn)程倉(cāng)庫(kù)加載模型或代碼。早期很多預(yù)訓(xùn)練模型會(huì)通過(guò) torch.hub 提供加載入口。

常用場(chǎng)景：

? 保存和共享模型相關(guān)代碼

? 加載已有模型

? 快速?gòu)?fù)現(xiàn)實(shí)驗(yàn)

? 組織模型發(fā)布流程

需要注意的是，實(shí)際生產(chǎn)部署中，模型格式、依賴環(huán)境、推理框架和安全策略都需要進(jìn)一步考慮，不能只依賴單一保存函數(shù)。

8、torchvision、torchaudio 與其他生態(tài)庫(kù)

PyTorch 本體之外，還有一組常見(jiàn)生態(tài)庫(kù)，用于處理不同領(lǐng)域的數(shù)據(jù)和模型。PyTorch 官網(wǎng)也強(qiáng)調(diào)其生態(tài)工具和庫(kù)支持計(jì)算機(jī)視覺(jué)、自然語(yǔ)言處理等方向，并支持分布式訓(xùn)練和生產(chǎn)部署。

常見(jiàn)生態(tài)庫(kù)：

? torchvision：計(jì)算機(jī)視覺(jué)庫(kù)，提供圖像數(shù)據(jù)集、圖像變換、視覺(jué)模型等

? torchaudio：音頻處理庫(kù)，提供音頻數(shù)據(jù)加載、特征提取和音頻模型相關(guān)工具

? torchtext：文本處理相關(guān)工具，但當(dāng)前使用時(shí)需要注意版本兼容和維護(hù)狀態(tài)

? torchmetrics：常用評(píng)價(jià)指標(biāo)庫(kù)，常與 PyTorch Lightning 配合使用

? torchao：面向量化、低精度和推理優(yōu)化的生態(tài)工具

? TorchRL：強(qiáng)化學(xué)習(xí)相關(guān)工具庫(kù)

這些庫(kù)不一定都屬于 torch 頂層模塊，但在真實(shí)項(xiàng)目中經(jīng)常與 PyTorch 一起使用。

綜合示例 3：保存與加載模型參數(shù)

下面展示 PyTorch 中最常見(jiàn)的模型保存方式：保存 state_dict，即模型參數(shù)字典。

在 PyTorch 中，推薦把“模型結(jié)構(gòu)定義代碼”和“模型參數(shù) state_dict”配合使用。這樣更清晰，也更符合常見(jiàn)工程實(shí)踐。

四、學(xué)習(xí) PyTorch 模塊時(shí)的優(yōu)先順序

PyTorch 模塊很多，但初學(xué)時(shí)不必平均用力。更合理的學(xué)習(xí)順序是先掌握訓(xùn)練閉環(huán)，再逐步擴(kuò)展到復(fù)雜模型和工程能力。

第一階段，應(yīng)優(yōu)先掌握：

? torch：張量創(chuàng)建、張量運(yùn)算、形狀變換、設(shè)備遷移。

? torch.autograd：自動(dòng)求導(dǎo)、backward()、梯度概念。

? torch.nn：模型定義、常用層、損失函數(shù)。

? torch.optim：優(yōu)化器、參數(shù)更新。

? torch.utils.data：數(shù)據(jù)集與數(shù)據(jù)加載器。

第二階段，可以繼續(xù)掌握：

? torch.nn.functional：函數(shù)式接口。

? torch.cuda、torch.mps、torch.amp：設(shè)備與混合精度訓(xùn)練。

? torch.linalg：矩陣計(jì)算。

? torch.distributions：概率建模。

? torch.utils.tensorboard：訓(xùn)練日志記錄。

第三階段，再根據(jù)實(shí)際需要學(xué)習(xí)：

? torch.compile：模型編譯加速。

? torch.export：模型導(dǎo)出。

? torch.distributed：分布式訓(xùn)練。

? torch.profiler：性能分析。

? torch.sparse：稀疏計(jì)算。

? torchvision、torchaudio 等領(lǐng)域生態(tài)庫(kù)。

這種順序符合深度學(xué)習(xí)項(xiàng)目的自然發(fā)展路徑：先能訓(xùn)練一個(gè)模型，再能組織復(fù)雜模型，最后再考慮加速、部署和規(guī)模化訓(xùn)練。

五、PyTorch 與 Scikit-learn 模塊組織方式的差異

Scikit-learn 更強(qiáng)調(diào)傳統(tǒng)機(jī)器學(xué)習(xí)中的統(tǒng)一接口。許多模型都遵循 fit()、predict()、transform() 等接口，因此可以方便地替換模型、組合流程和比較結(jié)果。

PyTorch 的重點(diǎn)不同。它更強(qiáng)調(diào)動(dòng)態(tài)計(jì)算圖、張量計(jì)算、自動(dòng)求導(dǎo)和靈活的模型定義方式。PyTorch 不會(huì)把完整訓(xùn)練流程封裝成一個(gè)統(tǒng)一的 fit() 方法，而是把更多控制權(quán)交給開(kāi)發(fā)者：用戶自己寫(xiě)前向傳播、損失計(jì)算、反向傳播、參數(shù)更新和評(píng)估邏輯。

因此，二者的模塊化思想有所不同：

? Scikit-learn 更像“規(guī)范化的機(jī)器學(xué)習(xí)工具箱”

? PyTorch 更像“可編程的深度學(xué)習(xí)計(jì)算框架”

? Scikit-learn 適合快速構(gòu)建傳統(tǒng)機(jī)器學(xué)習(xí)流程

? PyTorch 適合靈活構(gòu)建神經(jīng)網(wǎng)絡(luò)、復(fù)雜訓(xùn)練過(guò)程和深度學(xué)習(xí)模型

這并不意味著二者互相替代。在真實(shí)項(xiàng)目中，Scikit-learn 常用于特征工程、傳統(tǒng)機(jī)器學(xué)習(xí)基線和評(píng)估流程；PyTorch 則常用于深度神經(jīng)網(wǎng)絡(luò)、計(jì)算機(jī)視覺(jué)、自然語(yǔ)言處理、多模態(tài)模型和大規(guī)模訓(xùn)練。

小結(jié)

PyTorch 的主要模塊大致可分為三層：核心訓(xùn)練工作流模塊、模型構(gòu)建與計(jì)算支持模塊、工程擴(kuò)展與生態(tài)模塊。理解這些模塊的分工，有助于從整體上把握 PyTorch 的結(jié)構(gòu)，并逐步建立完整的深度學(xué)習(xí)訓(xùn)練與工程實(shí)踐意識(shí)。

“點(diǎn)贊有美意，贊賞是鼓勵(lì)”