本文章是本人通過讀《Pytorch實用教程》第二版做的學習筆記,深度學習的核心部分:數據準備 ?? 模型構建 ?? 模型訓練 ?? 模型評估與應用。根據上面的思路,我們分為幾個部分:
-
第一部分:PyTorch 基礎 - 涵蓋了從基本概念(如 Tensor)到數據加載、模型搭建和訓練的所有核心知識。
-
第二部分:PyTorch 實戰 - 提供了一些非常棒的案例,讓你把學到的知識應用到真實世界的問題中,比如圖像分類和自然語言處理。
-
第三部分:PyTorch 模型部署 - 教你如何將訓練好的模型應用到生產環境中。
一 PyTorch 基礎
1.1 Tensor
[!question] 先提一個問題,Tensor和Numpy中的array相比,最關鍵的區別是什么?
Tensor 與 NumPy 數組最根本的區別是Tensor能夠給深度學習提供GPU計算和自動求梯度的能力:
- GPU 加速計算: 利用 GPU 大規模并行計算的能力,極大縮短訓練時間,是深度學習能夠發展到今天的關鍵。
- 自動求梯度 (Automatic Differentiation): 通過記錄計算圖。我們只需要搭建好模型(前向傳播),PyTorch 就能自動計算出所有參數的梯度,然后我們就可以用這些梯度來更新模型(反向傳播),讓模型變得越來越好。
創建Tensor
我們了解了Tensor能做什么后,來看下如何創建Tensor,創建Tensor的方式很多,我們只說兩種,其他方式可移步pytorch官方文檔
torch.tensor(numpy_array)
: 會創建一個 Tensor,并復制一份 NumPy 數組的數據。之后修改原來的 NumPy 數組,不會影響到這個新的 Tensor。是最常用、最直接的方法,尤其適合將 Python 的列表 (list) 或元組 (tuple) 轉換成 Tensor。torch.from_numpy(numpy_array)
: 創建的 Tensor 和原來的 NumPy 數組會共享內存。這意味著,如果你修改其中一個,另一個也會跟著改變。
通過下面的例子了解下:
import torch
import numpy as np# 使用 torch.from_numpy() (共享內存)
numpy_arr = np.array([1, 2, 3, 4])
tensor_shared = torch.from_numpy(numpy_arr)print(f"修改前,NumPy 數組是: {numpy_arr}")
print(f"修改前,Tensor 是: {tensor_shared}")# 我們只修改 NumPy 數組
numpy_arr[0] = 99 print("---------------------------")
print(f"修改后,NumPy 數組是: {numpy_arr}")
print(f"修改后,Tensor 也跟著變了: {tensor_shared}") # 注意這里的變化
在什么場景下使用 torch.from_numpy()
會比 torch.tensor()
更有優勢?
在進行前向傳播 (forward pass) 的時候,我們經常需要把來自不同來源(比如用 OpenCV 或其他庫處理過的圖像數據)的 NumPy 數組送入 PyTorch 模型。
這時候,如果數據量非常大(比如成千上萬張高清圖片),使用 torch.tensor()
就會遇到一個問題:內存占用會翻倍。因為它需要額外申請一塊內存來存放復制過來的數據。
而使用 torch.from_numpy()
就優雅地解決了這個問題。它非常高效,因為它避免了不必要的數據復制,直接利用已經存在于內存中的 NumPy 數據。這在處理大規模數據集時,能節省大量的內存空間和數據復制的時間。
torch.tensor()
:更安全,因為數據是獨立的。適合一般情況。torch.from_numpy()
:更高效,因為共享內存。適合處理大型數據集,特別是作為模型輸入時。
“全 1”或“全 0” Tensor
“全 1”或“全 0” Tensor的用途
- 參數初始化,比如,神經網絡中,經常把“偏置”參數全部初始化為0
- 作為掩碼,假設你有一個 Tensor,但你只想保留其中一部分的數據,把另一部分數據“屏蔽”掉(比如變成 0)。這時你就可以創建一個由 0 和 1 組成的“掩碼” Tensor,然后把它和你的原始 Tensor 相乘。任何數字乘以 0 都會變成 0,乘以 1 則保持不變。這樣就實現了精確的“屏蔽”效果,這在很多高級應用(比如 NLP 里的注意力機制)中非常有用。
主要的函數是torch.ones()和tensor.zeros()
Tensor的屬性
主要用于檢查Tensor的屬性,比如其尺寸(形狀)、包含的數字類型(dtype)以及存儲位置
.shape
:張量的大小和尺寸。📏.dtype
:存儲在里面的數字的數據類型(如小數的float32
或整數的int64
)。.device
:Tensor 的物理存儲位置—— 在 CPU 或 GPU (cuda
) 上。
import torch# Our sample tensor
my_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)# Checking its attributes
print(f"Shape: {my_tensor.shape}")
print(f"Data Type: {my_tensor.dtype}")
print(f"Device: {my_tensor.device}")# 輸出結果
Shape: torch.Size([2, 3])
Data Type: torch.float32
Device: cpu
Tensor的計算
Tensor計算,最常用的,也是最重要的就是矩陣乘法,主要函數是torch.matmul()
,作為一個方便快捷方式,可以使用@
符號
import torch# Create a 2x3 tensor
tensor_A = torch.randn(2, 3)# Create a 3x2 tensor
tensor_B = torch.randn(3, 2)# Perform matrix multiplication using the @ operator
result = tensor_A @ tensor_Bprint(f"Shape of A: {tensor_A.shape}")
print(f"Shape of B: {tensor_B.shape}")
print(f"Shape of the result: {result.shape}")# 輸出
Shape of A: torch.Size([2, 3])
Shape of B: torch.Size([3, 2])
Shape of the result: torch.Size([2, 2])
以上是非常簡單的Tensor操作,如果想了解更深,可以去官網查看更多API
1.2 數據加載
數據加載主要介紹Dataset和DateLoader
Dataset
📖: 這就像食譜 。它知道所有單獨的數據點(“食譜”)是什么以及在哪里可以找到它們。它的主要工作是定義數據點的總數以及如何在詢問時獲取單個數據點。DataLoader
👨?🍳:這就像廚師 。廚師獲取食譜(數據集
)并有效地準備模型的數據。廚師負責打亂食譜,將它們分組(例如,一次提供 8 道菜),甚至可以使用多個廚房助手(多處理)來加快速度。
Dataset
在PyTorch中,任何我們想要自定義的數據集,都必須遵循一個固定的“格式”,要求我們實現最核心的兩個“功能”:
__len__(self)
: 這個函數用來告訴 PyTorch,我們的數據集中一共有多少個樣本。就像是食譜的目錄,告訴廚師一共有多少道菜。__getitem__(self, idx)
: 這個函數用來獲取單個樣本。當我們給它一個索引號idx
(比如 5),它就要能準確地把第 5 個樣本(比如第 5 張圖片和它對應的標簽)拿出來。__getitem__
函數負責返回一對:(image_tensor, label)。
只要我們定義好了這兩個函數,PyTorch的DataLoader就知道該如何與數據進行交互了。
以下是一個自定義數據集CatDogDataset
的骨架。
from torch.utils.data import Dataset
# We'll probably need libraries to handle file paths and open images
import os
from PIL import Image class CatDogDataset(Dataset):def __init__(self, image_dir, transform=None):# 1. Get a list of all our image file names from the directory.self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir)]self.transform = transformdef __len__(self):# 2. The total number of samples is just the number of image files.return len(self.image_paths)def __getitem__(self, idx):# 3. Get the full file path for the requested index `idx`.image_path = self.image_paths[idx]# 4. Load the image using its path.image = Image.open(image_path)# 5. Determine the label from the file name.if 'cat' in image_path:label = 0else: # 'dog'label = 1# 6. (Optional but important) Apply any transformations (like resizing).if self.transform:image = self.transform(image)# 7. Return the image and its label.return image, label
DataLoader
DataLoader采用我們的Dataset(食譜)并處理準備訓練數據的工作,其主要工作是:
- 批處理:將單個數據點分組為批次,我們可以一次向模型提供一批32或64張圖像,而不是一次向模型提供一張圖像。這對GPU來說效率要高得多。
- Shuffle:在每個訓練周期(epoch)開始隨機洗牌數據
- 并行處理:使用多個后臺工作線程同時加載數據,這樣模型不必等待
以下是我們為剛剛定義的 CatDogDataset
創建 DataLoader
的方法:
from torch.utils.data import DataLoader# 1. First, we create an instance of our "cookbook".
# (Assuming we have a folder named 'data/cats_and_dogs')
cat_dog_dataset = CatDogDataset(image_dir='data/cats_and_dogs')# 2. Now, we give the cookbook to our "chef".
# We'll ask for batches of size 32 and to shuffle the data.
data_loader = DataLoader(dataset=cat_dog_dataset, batch_size=32, shuffle=True)# 3. Now we can iterate over the data_loader to get our batches.
for image_batch, label_batch in data_loader:# In each loop, we get a batch of 32 images and their 32 labels.# We can then feed these batches directly to our model.pass
可以看到DataLoader中shuffle是True,通過打亂數據,確保模型看到的每一批都是整個數據集的隨機混合樣本。這迫使模型學習真正的區別特征(如尖耳朵或胡須),使最終模型更加穩健 ,并且能夠更好地推廣到新的、看不見的數據
[!summary] 小結
我們現在已經構建了一個完整的數據管道。我們有:
- 基本數據塊 (Tensor)。
- 知道在哪里可以找到數據( 數據集 )的“食譜”。
- 準備訓練的“廚師”(DataLoader)。
數據準備好,下一步我們就要構建模型
1.3 模型搭建和訓練
模型搭建
在 PyTorch 中,所有的神經網絡模型都應該繼承一個叫做 nn.Module
的“基類”。
可以把 nn.Module
想象成一個樂高積木的“底板” ?。它本身提供了一些非常重要的基礎功能(比如追蹤模型的所有參數),而我們則需要往這塊底板上添加我們自己的“積木塊”(也就是神經網絡的各種層,比如線性層、卷積層等)。
搭建一個模型,通常也需要實現兩個核心部分:
-
__init__(self)
: 這是模型的“構造函數”。我們在這里定義模型需要用到的所有“積木塊”(神經網絡層)。 -
forward(self, x)
: 這里是“前向傳播”的核心。我們在這里連接__init__
中定義的積木塊,明確規定數據x
應該如何一步步流過這些層,最終得到輸出。
簡單來說:__init__
負責“買零件”,forward
負責“組裝”。
下面我們來看一個簡單的例子,構建一個模型,可以查看手寫數字的圖片(來自著名的 MNIST 數據集)并對它是哪個數字進行分類(0 到 9)。
MNIST 圖像是 28x28 像素的灰度圖片。為了將其輸入到一個簡單的模型中,我們首先將其展平為一條 784 像素的單行 (28 * 28 = 784)。
以下是此任務的基本模型在 PyTorch 中的樣子:
import torch.nn as nnclass SimpleClassifier(nn.Module):def __init__(self):# First, we must call the __init__ of the parent class (nn.Module)super().__init__()# "Buying the parts": We need one 'linear' layer.# It will take the 784 pixels as input and must output 10 numbers,# one score for each possible digit (0-9).self.layer1 = nn.Linear(in_features=784, out_features=10)def forward(self, flattened_image):# "Assembling the parts": We define how the data flows.# The flattened image data simply goes through our one layer.output = self.layer1(flattened_image)return output
784是將28 * 28的圖片展平為一條784像素的單行,10是此目標分類個數是10
上面的模型非常簡單,只是一個線性層,。為了學習手寫之類的復雜模式,我們需要在層之間引入一個“秘密成分”: 非線性激活函數。最常見的激活函數之一是ReLU,讓我們加入激活函數。
import torch.nn as nn
import torch.nn.functional as Fclass SimpleClassifierWithReLU(nn.Module):def __init__(self):super().__init__()# "Buying the parts"self.layer1 = nn.Linear(784, 128) # Hidden layerself.layer2 = nn.Linear(128, 10) # Output layerdef forward(self, x):# "Assembling the parts"# Pass through the first layerx = self.layer1(x)# Apply the ReLU activation functionx = F.relu(x)# Pass through the output layeroutput = self.layer2(x)return output
模型訓練
訓練模型本質上是一個循環,在這個循環里,不斷地根據訓練的模型結果,進行“微調”,這個過程需要三個關鍵組件:
- 損失函數(Loss Function): 用來衡量模型預測結果和真實答案之間的差距。差距越大,損失值就越高。
- 優化器(Optimizer): 根據損失函數計算出的梯度,來更新模型的權重參數,目標是讓損失值越來越小。
- 訓練循環(Training Loop): 把所有步驟(獲取數據 -> 前向傳播 -> 計算損失 -> 反向傳播 -> 更新權重)串起來,并重復執行
損失函數
損失函數可以看做是記分員,他的工作室兩件事:
- 模型的預測結果
- 實際的正確答案,即標簽
來衡量預測結果和正確答案的差距,即損失。
- 如果模型的預測偏離很大,就給出高分
- 如果模型預測偏離小,就給出低分
訓練的最終目標是將這個分數最小化
有很多介紹損失函數的文章,詳細資料可網上去搜
優化器
損失函數是“記分員”, 優化器就像是“教練”。
教練會根據記分員給出的分數(損失值),以及每個參數對這個分數的影響程度(梯度),來制定一個“訓練計劃”,告訴模型的每一個參數(權重)應該如何微調——是該調高一點,還是該調低一點,以及調整的幅度應該多大。
最常用的優化器之一叫做 Adam。現在我們只需要知道,它的工作就是根據 backward()
計算出的梯度,來智能地更新模型的權重,從而讓損失值越來越小。
下面我們了解下優化器的發展,可以看做是分三步走:
- SGD(隨機梯度下降):這是最基礎、最經典的優化器,就像是優化器中的“老爺車”,理解它就能理解了所有優化器的出發點
- Momentum(動量):這是對SGD的一個重要改進,給“老爺車”加了一個“慣性系統”,讓它跑的更穩,更快
- Adam:這是目前最流行、最常用的優化器之一,非常智能,自適應不同的情況,就像一輛“現代跑車”
關于優化器的最常見比喻是,尋優過程就像是蒙眼下山。
-
SGD
可以把它想象成一個蒙著眼睛、想要走到山谷最低點的下山者 👨?🦯。
他看不見整個山谷的全貌,所以他只能采取一個最簡單的策略:- 在當前位置,伸腳向四周探一探,感受哪個方向是下坡最陡的。(這就是計算梯度)
- 然后,朝著這個最陡峭的方向,邁出一小步。(這就是更新權重)
- 重復這個過程,一步一步地往下走。
在簡單的蒙眼下山策略中,有個潛在問題,如果徒步者走入山坡上的小凹陷,從他們的位置來看,每個方向似乎要么是平坦的,要么是略微上坡的,導致進入“局部最優解”,無法到達真正的谷底(全局最優解),SGD的最大弱點之一:很容易被局部最優解或鞍點困住
-
Monentum
上面的問題,就是Monentum要解決的,Monentum是“動量”、“慣性”的意思,在優化器里Monentum算法會積累過去幾個步驟的梯度(就像是累計“質量”和“速度”),形成一個“動量”,當遇到梯度變小(比如平坦區域或局部小坑)的時候,下山者也能通過這個“動量”沖過去,不會像SGD那樣被卡主 -
Adam (Adaptive Moment Estimation)
Adam 是目前最受歡迎的優化器之一,我們可以把它看作一輛智能的現代跑車 🏎?。它不僅吸收了 Momentum 的“慣性”優點,還增加了一個更強大的功能:自適應學習率 (Adaptive Learning Rates)。
拿下山的例子,在下山的時候,步子大一點,能快速接近谷底;快到谷底的時候,步子變小,有利于找到最底的谷底
這里的“一小步”的大小,在深度學習里被稱為“學習率 (Learning Rate)”。
如果步子(學習率)太大,他可能會一步邁過頭,直接跨到對面山坡上,導致永遠在谷底附近來回震蕩,到不了最低點。
如果步子(學習率)太小,他下山的速度會非常非常慢
訓練循環
訓練循環 (Training Loop) 就是整個過程的核心,它把我們之前討論的所有部件——數據、模型、損失函數和優化器——全部串聯起來,協同工作。
標準的訓練循環就像一個固定的“儀式”,每一步都有明確的目的:
- 清零梯度 (
optimizer.zero_grad()
): 準備開始新一輪的計算。 - 前向傳播 (
model(inputs)
): 讓模型根據輸入數據進行預測。 - 計算損失 (
loss_fn(outputs, labels)
): 評估模型的預測有多糟糕。 - 反向傳播 (
loss.backward()
): 根據損失,計算出每個參數應該如何調整(即計算梯度)。 - 更新權重 (
optimizer.step()
): “教練”正式出手,根據梯度更新模型的參數。
這五個步驟會一遍又一遍地重復,每一次重復,模型都會變得比上一次更“聰明”一點。
[!tip] 為什么要先進行梯度清零
PyTorch 在設計backward()
函數時,就是讓它把新計算出的梯度累加到已有的.grad
屬性上,而不是覆蓋掉。
這么設計其實是有意為之的,因為它在一些高級應用(比如循環神經網絡 RNN 的某些變種)中非常有用。
但對于我們現在正在做的、最常見的訓練任務來說,每一輪的梯度計算都應該是一個全新的開始,完全獨立于上一輪。我們只關心當前批次數據所產生的梯度。
所以,如果我們不在每一輪循環開始時手動“清零”,那么舊的梯度就會像“幽靈”一樣一直影響著新的梯度,導致“教練”(優化器)拿到完全錯誤的信息,最終模型也就無法被正確地訓練了。
二 PyTorch實踐
下面我們來進行代碼實戰,我們的目標是編寫一個完整的 Python 腳本
- 加載 MNIST 手寫數字數據集。
- 構建
SimpleClassifierWithReLU
模型。 - 使用我們討論的損失函數和優化器訓練模型。
- 評估訓練模型的性能。
2.1 準備數據集
我們可以逐步構建腳本。任何 Python 文件的第一步始終是導入必要的庫 。
import torch
from torch import nn # nn 包含了模型層 (nn.Module, nn.Linear) 和損失函數 (nn.CrossEntropyLoss)
from torch.utils.data import DataLoader
from torchvision import datasets # 這是一個方便的庫,已經幫我們打包好了 MNIST 等常用數據集
from torchvision.transforms import ToTensor # 這是一個工具,可以把圖片轉換成 Tensor
import torch.optim as optim # 這里面有我們需要的各種優化器,比如 Adam
在 PyTorch 里,準備數據的代碼通常非常簡潔,因為 torchvision
庫已經為我們處理了大部分繁瑣的工作,比如下載數據集和進行基礎的轉換。
這是加載 MNIST 訓練數據的標準代碼:
# --- 1. 數據準備 ---
training_data = datasets.MNIST(root="data", # 指定數據下載后存放的目錄train=True, # 明確指出這是訓練集download=True, # 如果 'data' 目錄里沒有,就自動下載transform=ToTensor() # 把圖片數據轉換成 PyTorch Tensor
)# 我們可以用 DataLoader 來打包數據
train_dataloader = DataLoader(training_data, batch_size=64)
[!question] 為什么使用ToTensor
- 更改格式: 它獲取圖像(通常是來自 Pillow (PIL) 等庫的數據結構),并將其轉換為 PyTorch 張量。
- 縮放像素值: 這是關鍵部分。圖像像素通常是從 0(黑色)到 255(白色)的整數。
ToTensor()
將它們轉換為 0.0 到 1.0 之間的浮點數。
這種縮放是規范化的一種形式,對于幫助神經網絡高效訓練非常重要
2.2 構建 SimpleClassifierWithReLU
模型。
- 構建模型
# --- 2. Model Definition ---
class SimpleClassifierWithReLU(nn.Module):def __init__(self):super().__init__()# "Buying the parts"self.layer1 = nn.Linear(28*28, 128) # Input is 784, hidden layer is 128self.layer2 = nn.Linear(128, 10) # Output is 10 classesdef forward(self, x):# "Assembling the parts"# This is the flattening step we talked about!# It reshapes the (1, 28, 28) image into a (784) vector.x = x.view(x.size(0), -1)# Now the data flows through the layersx = self.layer1(x)x = nn.functional.relu(x) # Apply ReLU activationoutput = self.layer2(x)return output```
## 2.3 損失函數
```python
# --- 3. Loss Function and Optimizer ---
loss_fn = nn.CrossEntropyLoss()
2.4 優化器
# (在模型定義和損失函數之后)# 首先,我們需要創建模型的一個實例
model = SimpleClassifierWithReLU()# 然后,我們創建優化器,并把模型的參數告訴它
optimizer = optim.Adam(model.parameters(), lr=1e-3) # lr 是學習率
2.5 清零梯度
# --- 4. The Full Training Loop ---# First, create an instance of our model
model = SimpleClassifierWithReLU()# Create our loss function and optimizer
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)# We'll train for 5 "epochs" (5 full passes over the dataset)
epochs = 5for epoch in range(epochs):print(f"Epoch {epoch+1}\n-------------------------------")# Loop over the data loaderfor batch, (X, y) in enumerate(train_dataloader):# 1. ZERO GRADIENTSoptimizer.zero_grad()# 2. FORWARD PASSpred = model(X)# 3. COMPUTE LOSSloss = loss_fn(pred, y)# 4. BACKWARD PASSloss.backward()# 5. UPDATE WEIGHTSoptimizer.step()# Optional: Print progressif batch % 100 == 0:loss, current = loss.item(), batch * len(X)print(f"loss: {loss:>7f} [{current:>5d}/{len(train_dataloader.dataset):>5d}]")print("Done Training!")
2.6 模型評估
# (Assuming we have a test_dataloader prepared just like our train_dataloader)
size = len(test_dataloader.dataset)
num_batches = len(test_dataloader)
model.eval() # Set the model to evaluation mode
test_loss, correct = 0, 0with torch.no_grad(): # We don't need gradients for testingfor X, y in test_dataloader:pred = model(X)test_loss += loss_fn(pred, y).item()# Find the index of the highest score, which is the model's predictioncorrect += (pred.argmax(1) == y).type(torch.float).sum().item()test_loss /= num_batches
correct /= size
print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
三 部署應用
部署的目標是采用我們訓練好的PyTorch模型并將其轉換為以下格式:
- Fast:針對純預測速度(推理)進行優化,而不是訓練
- Portable:能夠在不同的環境中運行,例如 Web 服務器、手機或邊緣設備,這些設備甚至可能沒有安裝 Python 或 PyTorch
主要側重將兩項關鍵技術:ONNX和TensorRT
讓我們打個比方:
-
我們的 PyTorch 模型 :它就像一個復雜的定制引擎??,在我們的車間(Python 環境)中設計。它功能強大,但需要我們所有的專用工具才能運行。
-
ONNX:這就像為我們的引擎創建通用技術藍圖 📜。這是一種標準格式,準確描述了引擎的工作原理。任何能閱讀此標準藍圖的人都可以構建我們引擎的副本,即使他們沒有我們原創的創意工坊工具。這解決了便攜性問題。
-
TensorRT: 這是英偉達打造的高性能工廠🏭。它采用藍圖(ONNX 文件)并使用先進的優化技術來構建我們引擎的極快生產級版本,專門針對在 NVIDIA GPU 上運行進行了調整。這解決了速度問題。
因此,典型的工作流程是:PyTorch 模型 -> ONNX 藍圖 -> 優化的 TensorRT 引擎 。
3.1 ONNX
ONNX(開放神經網絡交換)是通用藍圖 。從我們的 PyTorch 模型創建此藍圖的過程稱為導出 。
這個函數以一種巧妙的方式工作。它不僅保存模型的架構,還跟蹤模型。這意味著我們需要:
- 創建一段虛假的輸入數據(“虛擬輸入”)。
- 通過我們的模型傳遞此虛擬輸入。
- 當數據流過時,
torch.onnx.export()
會記錄發生的每一個作。 - 然后,它將此記錄的作序列保存到
.onnx
文件中。
以下是我們一直在使用的 SimpleClassifierWithReLU
模型的執行此作:
import torch# (Assume the SimpleClassifierWithReLU class is defined above)# 1. Create an instance of our model
model = SimpleClassifierWithReLU()
# IMPORTANT: Before exporting, you would typically load your saved trained weights
# and put the model in evaluation mode.
# model.load_state_dict(torch.load("path_to_weights.pth"))
model.eval()# 2. Create a dummy input tensor with the correct shape.
# Our model expects a flattened image, but the original input is (batch_size, channels, height, width)
# So let's create a dummy 28x28 image.
# We'll use a batch size of 1.
dummy_input = torch.randn(1, 1, 28, 28) # 3. Export the model
torch.onnx.export(model, # The model to exportdummy_input, # A sample input to trace the model"mnist_classifier.onnx", # The name of the output fileinput_names=["input_image"], # A name for the input nodeoutput_names=["output_scores"] # A name for the output node
)print("Model successfully exported to mnist_classifier.onnx")
[!question] 這個過程最獨特的部分是需要
dummy_input
。為什么我們需要向導出功能提供這些虛假數據嗎?
答案是為了追蹤,舉個例子,寫下你的朋友從椅子這里走到前門那,這個過程中你的朋友走的每一步你都記錄了下來。這個過程中,你的朋友就是dummy_input
; 穿過房間的每一步就是前向傳遞;你記錄的文件就是.onnx
文件
如何使用這個文件_在_未安裝 PyTorch 的環境中進行預測(例如,在簡單的 Web 服務器上)?
我們有.onnx
文件,但現在我們需要一個“閱讀器”或“引擎”來理解并執行它。該引擎稱為 ONNX Runtime
可以這樣想:
- 要打開
.pdf
文件,您需要一個像 Adobe Reader 這樣的程序。 - 若要運行
.onnx
模型,需要像 ONNX Runtime 這樣的庫。
runtime是安裝在部署環境中的單獨輕量級庫。它根本不需要 PyTorch。它唯一的工作是加載 .onnx
文件并非常非常快地執行預測(推理)。
下面介紹如何使用 Python 中的 onnxruntime
庫來運行導出的模型:
import onnxruntime as ort
import numpy as np# 1. Create an "inference session" by loading the .onnx file
session = ort.InferenceSession("mnist_classifier.onnx")# 2. Get the name of the input layer (we named it "input_image" during export)
input_name = session.get_inputs()[0].name# 3. Prepare a sample input. ONNX Runtime works well with NumPy arrays.
# The shape must match what the model expects: (1, 1, 28, 28)
sample_input = np.random.rand(1, 1, 28, 28).astype(np.float32)# 4. Run the prediction
# The result is a list containing the output arrays
results = session.run(None, {input_name: sample_input})# 5. Interpret the result
output_scores = results[0]
predicted_digit = np.argmax(output_scores)print(f"The ONNX Runtime produced output scores: \n{output_scores}")
print(f"The predicted digit is: {predicted_digit}")
3.2 TensorRT
將其視為一個高度專業化的工廠🏭,它獲取模型的藍圖(.onnx
文件)并重新構建它,使其在特定的 NVIDIA GPU 上盡可能快地進行物理處理。
這不僅僅是運行模型;它正在積極優化它。TensorRT 使用幾種巧妙的技術來做到這一點,但我們可以專注于三個主要想法:
- 精密校準(量化): 它巧妙地使用不太精確的數字來更快地進行數學運算,例如使用
1.5
而不是1.5000001
。 - 圖層融合 :它將模型的多個步驟組合成一個超級高效的步驟。
- 內核自動調整 :它為您的特定 GPU 硬件上的每個作找到絕對最快的代碼。
精準校準(量化)
一種更廣泛地稱為量化的技術。核心思想出奇地簡單。想象一下你正在測量一些東西。
- 您可以使用一個非常精確的數字,例如
3.14159265
- 或者,對于大多數實際目的,您可以只使用
3.14
。
第二個數字不太精確,但它更短且更容易使用。電腦也有同樣的感覺!
默認情況下,神經網絡使用高精度 32 位數字(稱為FP32
)進行訓練。TensorRT 分析模型并找出它可以在哪些方面安全地使用較低精度的數字,例如 16 位 (FP16
) 甚至 8 位整數 (INT8
),而不會對最終結果造成太大損害。
由于這些數字較小,GPU 可以更快地處理它們并將更多數字放入內存中。
不過,這不僅僅是盲舍入。“校準”部分意味著 TensorRT 使用實際數據的一小部分樣本來智能地找出將數字轉換為較低精度的最佳方法,同時將信息丟失降至最低。這里的主要好處是速度的巨大提升。
但如果操作太“粗糙”,累積的誤差就會越來越大,最終導致模型的準確率下降。 TensorRT 的“校準”(Calibration)步驟就顯得至關重要。它會非常智能地分析模型,只在那些對最終結果影響不大的地方使用低精度計算,而在關鍵部分仍然保持高精度。這是一個在速度和精度之間的權衡。
圖層融合
圖層融合,概念也很直觀,這個概念也很直觀。想象一下在工廠的流水線上組裝一個玩具:
- 步驟1: 工人A拿起玩具的身體。
- 步驟2: 工人A把身體遞給工人B。
- 步驟3: 工人B給玩具裝上頭部。
- 步驟4: 工人B把玩具遞給工人C。
- 步驟5: 工人C給頭部畫上眼睛。
現在,如果我們把這三個步驟融合成一個,讓工人A一個人完成“拿起身體 -> 裝上頭部 -> 畫上眼睛”這整套動作,會發生什么?整個流水線的效率會有什么變化?
答案是效率直線上升,什么原因導致的呢?
有兩個主要原因:
- 減少開銷: 消除了工作線程之間的“切換”時間。GPU 不必啟動新任務,從內存中讀取數據,然后將其寫回,只是為了發生下一個小步驟。
- 節省內存 :在融合發生時,數據可以保留在 GPU 內核的本地內存(緩存)中,而不是在每一步之間發送回主 GPU 內存。
因此,TensorRT 會查看模型的藍圖,并自動找到可以合法有效地融合到單個自定義作中的層序列。
內核自動調整
想象一下,你是一位頂級賽車工程師 👨?🔧,正在為一場特定的比賽調校賽車。對于賽車上的每一個螺絲,你手上都有一整套工具箱,里面有幾十種看似相同但實際有細微差別的扳手。
為了追求極致的速度,你會不厭其煩地用每一種扳手去試著擰緊那個螺絲,直到找到那一把能讓你用最快、最完美的方式完成工作的扳手。
在 GPU 的世界里:
- “核心” (Kernel) 就是一個為特定任務(比如一次卷積計算)編寫的高度優化的底層代碼,就像是那把“扳手”。
- 對于同一個任務,NVIDIA 的工程師們已經準備好了很多種不同的“核心”實現。
- “自動調整” (Auto-Tuning) 就是 TensorRT 在構建模型時,會像那位工程師一樣,為模型中的每一個操作,在你的特定 GPU 上實際運行和測試多種不同的“核心”,然后選擇那個表現最快的。