深度學習模型訓練過程中，資料載入與 GPU 記憶體管理至關重要。本文將探討如何利用 PyTorch 最佳化資料載入流程，提升 GPU 記憶體使用效率，進而提升模型訓練速度與穩定性。首先，分析如何使用 torch.utils.data.DataLoader 進行批次資料載入，並藉由 py-spy 等效能分析工具找出程式碼中的效能瓶頸。接著，示範如何將運算密集的操作，例如隨機噪聲新增，從 CPU 轉移到 GPU，以充分發揮 GPU 的運算能力。最後，討論如何在 GPU 記憶體不足的情況下，使用 nvidia-smi 監控 GPU 狀態，並使用梯度檢查點（Gradient Checkpointing）等技術來減少 GPU 記憶體佔用，同時維持訓練效率。
深度解析 PyTorch 模型除錯：最佳化資料載入與 GPU 記憶體管理
在深度學習模型的訓練過程中，資料載入效率和 GPU 記憶體管理是兩個至關重要的因素。本文將探討如何使用 PyTorch 最佳化資料載入流程，並有效管理 GPU 記憶體，以提升模型訓練的效率和穩定性。
最佳化資料載入流程
在原始程式碼中，資料載入和預處理是在 CPU 上進行的，這可能導致訓練過程中的瓶頸。以下是一個典型的資料載入和訓練迴圈：
batch_size = 32
train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size)
optimizer = optim.Adam(model.parameters(), lr=2e-2)
criterion = nn.CrossEntropyLoss()
def train(model, optimizer, loss_fn, train_loader, val_loader, epochs=20, device='cuda:0'):
model.to(device)
for epoch in range(epochs):
print(f"epoch {epoch}")
model.train()
for batch in train_loader:
optimizer.zero_grad()
ww, target = batch
ww = ww.to(device)
target = target.to(device)
output = model(ww)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
model.eval()
num_correct = 0
num_examples = 0
for batch in val_loader:
input, target = batch
input = input.to(device)
target = target.to(device)
output = model(input)
correct = torch.eq(torch.max(output, dim=1)[1], target).view(-1)
num_correct += torch.sum(correct).item()
num_examples += correct.shape[0]
print("Epoch {}, accuracy = {:.2f}".format(epoch, num_correct / num_examples))
train(model, optimizer, criterion, train_data_loader, train_data_loader, epochs=10)
內容解密：
- 資料載入器（DataLoader）：使用
torch.utils.data.DataLoader來批次載入資料，batch_size設定為 32。 - 模型訓練：在每個 epoch 中，模型先進行訓練模式，然後遍歷訓練資料，載入批次資料到 GPU 上，進行前向傳播、計算損失、反向傳播和引數更新。
- 模型評估：在每個 epoch 結束後，模型切換到評估模式，並在驗證集上計算準確率。
使用 py-spy 分析效能瓶頸
透過使用 py-spy 工具，可以對程式進行效能分析，找出瓶頸所在：
py-spy -r 99 -d 120 --flame slowloader.svg -- python slowloader.py
火焰圖顯示，大部分時間花在影像載入和轉換為張量上，同時有 16.87% 的時間花在應用隨機噪聲上。
內容解密：
py-spy使用：透過py-spy對 Python 程式進行抽樣分析，生成火焰圖以視覺化效能瓶頸。- 效能瓶頸：發現應用隨機噪聲的操作佔用了相當比例的執行時間。
最佳化隨機噪聲應用
透過將隨機噪聲的應用從 CPU 轉移到 GPU，可以顯著提高效率。最佳化後的程式碼如下：
def add_noise_gpu(tensor, device):
random_noise = torch.randn_like(tensor).to(device)
return tensor.add_(random_noise)
在訓練迴圈中，於資料傳輸到 GPU 後新增噪聲：
input = input.to(device)
input = add_noise_gpu(input, device)
同時，移除資料載入器中的 BadRandom 轉換。
內容解密：
add_noise_gpu函式：在 GPU 上生成與輸入張量相同形狀的隨機噪聲，並將其新增到輸入張量中。- 訓練迴圈修改：在將輸入資料傳輸到 GPU 後，直接在 GPU 上應用隨機噪聲，避免了 CPU 和 GPU 之間的頻繁資料傳輸。
GPU 加速的隨機噪聲應用效能比較
透過將隨機噪聲的應用轉移到 GPU，效能得到了顯著提升。使用 Jupyter Notebook 進行簡單的效能測試：
cpu_t1 = torch.rand(64,3,224,224)
cpu_t2 = torch.rand(64,3,224,224)
%timeit cpu_t1 + cpu_t2
gpu_t1 = torch.rand(64,3,224,224).to("cuda")
gpu_t2 = torch.rand(64,3,224,224).to("cuda")
%timeit gpu_t1 + gpu_t2
結果顯示，GPU 上的張量運算比 CPU 快約 20 倍。
內容解密：
- CPU 與 GPU 效能比較：透過比較 CPU 和 GPU 上張量加法的執行時間，展示了 GPU 在矩陣運算上的顯著優勢。
- GPU 加速的意義：將耗時的運算轉移到 GPU 上，可以有效提升整體訓練效率。
除錯 GPU 記憶體問題
當模型訓練過程中遇到 GPU 記憶體不足的問題時，可以使用 nvidia-smi 工具來監控和除錯 GPU 記憶體使用情況。
nvidia-smi
內容解密：
nvidia-smi使用：透過nvidia-smi命令檢視當前 GPU 的記憶體使用情況，幫助定位記憶體佔用高的原因。- GPU 記憶體管理：瞭解模型訓練過程中 GPU 記憶體的使用情況，以便進行相應的最佳化，如減少批次大小、最佳化模型結構等。
最佳化流程示意圖
此圖示展示了從原始訓練迴圈到最終最佳化方案的流程，包括效能分析、瓶頸發現、最佳化措施實施以及 GPU 記憶體除錯等關鍵步驟。
偵錯 GPU 問題
當使用 GPU 進行深度學習運算時，記憶體管理變得至關重要。瞭解 GPU 的使用情況有助於避免記憶體不足的問題。
使用 nvidia-smi 監控 GPU 狀態
nvidia-smi 是一個強大的工具，可以用來檢查 GPU 的使用情況。圖 7-12 顯示了在終端機中執行 nvidia-smi 的輸出結果。在 Jupyter Notebook 中，可以使用 !nvidia-smi 來呼叫這個工具。
圖示：nvidia-smi 輸出結果
@startuml
skinparam backgroundColor #FEFEFE
skinparam componentStyle rectangle
title PyTorch 模型除錯與生產環境佈署
package "Kubernetes Cluster" {
package "Control Plane" {
component [API Server] as api
component [Controller Manager] as cm
component [Scheduler] as sched
database [etcd] as etcd
}
package "Worker Nodes" {
component [Kubelet] as kubelet
component [Kube-proxy] as proxy
package "Pods" {
component [Container 1] as c1
component [Container 2] as c2
}
}
}
api --> etcd : 儲存狀態
api --> cm : 控制迴圈
api --> sched : 調度決策
api --> kubelet : 指令下達
kubelet --> c1
kubelet --> c2
proxy --> c1 : 網路代理
proxy --> c2
note right of api
核心 API 入口
所有操作經由此處
end note
@enduml這個範例來自於一台搭載 1080 Ti 的個人電腦。可以看到，多個 Jupyter Notebook 正在執行，每個 notebook 都佔用了部分記憶體，其中一個甚至使用了 4GB！可以使用 os.getpid() 來取得目前 notebook 的 PID。結果發現，佔用最多記憶體的程式竟然是一個用於測試 GPU 轉換的實驗性 notebook。可以想像，當模型、批次資料以及前向和後向傳遞的資料都載入記憶體時，記憶體使用量會迅速增加。
管理 GPU 記憶體
除了深度學習程式外，還有一些其他程式正在執行，例如 X server 和 GNOME。除非你正在使用本地機器，否則你幾乎不會看到這些程式。
PyTorch 會為每個程式分配約 0.5GB 的記憶體。因此，最好一次只處理一個專案，並且不要讓 Jupyter Notebook 亂七八糟地執行（可以使用 Kernel 選單來關閉與 notebook 相關的 Python 程式）。
單獨執行 nvidia-smi 可以獲得目前 GPU 使用情況的快照，但可以使用 -l 引數來獲得持續的輸出。以下是一個範例指令，每 5 秒輸出時間戳、已使用記憶體、可用記憶體、總記憶體和 GPU 使用率：
nvidia-smi --query-gpu=timestamp,memory.used,memory.free,memory.total,utilization.gpu --format=csv -l 5
手動釋放 GPU 記憶體
如果認為 GPU 使用的記憶體過多，可以嘗試讓 Python 的垃圾回收器介入。如果有一個不再需要的張量 tensor_to_be_deleted，可以按照 fast.ai 函式庫的建議，使用 del 將其刪除：
import gc
del tensor_to_be_deleted
gc.collect()
如果在 Jupyter Notebook 中頻繁建立和重新建立模型，可以發現刪除一些參考並呼叫 gc.collect() 可以回收一些記憶體。
梯度檢查點（Gradient Checkpointing）
儘管使用了刪除和垃圾回收技巧，仍然可能會遇到記憶體不足的問題。對於大多數應用程式，下一步是減少訓練迴圈中透過模型的批次大小。這樣會增加每個 epoch 的訓練時間，而且模型的效果可能不如使用足夠記憶體來處理更大批次大小的模型。但是，可以透過使用梯度檢查點來在 PyTorch 中權衡計算和記憶體。
梯度檢查點的主要目標是透過分割模型來減少在任何時候都存在於 GPU 上的狀態數量。這種方法意味著，可以將批次大小增加 4 到 10 倍，但是訓練將會更加耗費計算資源。在前向傳遞期間，PyTorch 儲存輸入和引數到一個段，但不實際執行前向傳遞。在後向傳遞期間，PyTorch 取回這些值，並為該段計算前向傳遞。中間值被傳遞到下一個段，但這些必須按段進行。
程式碼範例：使用梯度檢查點的 AlexNet
from torch.utils.checkpoint import checkpoint_sequential
import torch.nn as nn
class CheckpointedAlexNet(nn.Module):
def __init__(self, num_classes=1000, chunks=2):
super(CheckpointedAlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
)
self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
self.classifier = nn.Sequential(
nn.Dropout(),
nn.Linear(256 * 6 * 6, 4096),
nn.ReLU(inplace=True),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(inplace=True),
nn.Linear(4096, num_classes),
)
def forward(self, x):
x = checkpoint_sequential(self.features, chunks, x)
x = self.avgpool(x)
x = x.view(x.size(0), 256 * 6 * 6)
x = self.classifier(x)
return x
程式碼解密：
checkpoint_sequential的使用：在forward方法中，使用checkpoint_sequential對self.features進行分割。這樣可以減少記憶體的使用。chunks引數：在__init__方法中，新增了一個chunks引數，用於控制分割的段數。預設值為 2。- 模型架構：
CheckpointedAlexNet的架構與原始的 AlexNet 相似，但是使用了梯度檢查點技術。
PyTorch 生產環境佈署
模型佈署
在前幾章中，我們已經學習瞭如何使用 PyTorch 進行影像、文字和聲音的分類別。現在，我們將探討如何將 PyTorch 應用佈署到生產環境中。在本章中，我們將建立透過 HTTP 和 gRPC 進行 PyTorch 模型推理的應用程式，並將其封裝成 Docker 容器，佈署到 Google Cloud 上的 Kubernetes 叢集。
模型服務
建立模型只是構建深度學習應用的一部分。如果模型永遠不會做出任何預測，那麼它是否有價值呢？我們需要一種簡單的方法來封裝我們的模型，以便它們能夠回應請求（無論是透過網路或其他方式），並在生產環境中以最小的努力執行。
幸運的是，Python 允許我們使用 Flask 框架快速建立網路服務。在本文中，我們將構建一個簡單的服務，載入根據 ResNet 的貓或魚模型，接受包含影像 URL 的請求，並傳回一個 JSON 回應，指示影像是否包含貓或魚。
建立 Flask 服務
首先，安裝 Flask 函式庫：
conda install -c anaconda flask
pip install flask
建立一個名為 catfish 的新目錄，並將模型定義複製到 model.py 中：
from torchvision import models
CatfishClasses = ["cat", "fish"]
CatfishModel = models.ResNet50()
CatfishModel.fc = nn.Sequential(nn.Linear(transfer_model.fc.in_features, 500),
nn.ReLU(),
nn.Dropout(),
nn.Linear(500, 2))
接下來，建立另一個 Python 指令碼 catfish_server.py，用於啟動網路服務：
from flask import Flask, jsonify
from . import CatfishModel
from torchvision import transforms
import torch
import os
def load_model():
return model
app = Flask(__name__)
@app.route("/")
def status():
return jsonify({"status": "ok"})
@app.route("/predict", methods=['GET', 'POST'])
def predict():
img_url = request.image_url
img_tensor = open_image(BytesIO(response.content))
prediction = model(img_tensor)
predicted_class = CatfishClasses[torch.argmax(prediction)]
return jsonify({"image": img_url, "prediction": predicted_class})
if __name__ == '__main__':
app.run(host=os.environ["CATFISH_HOST"], port=os.environ["CATFISH_PORT"])
透過設定 CATFISH_HOST 和 CATFISH_PORT 環境變數，可以在命令列上啟動網路伺服器：
CATFISH_HOST=127.0.0.1 CATFISH_PORT=8080 python catfish_server.py
在瀏覽器中存取 http://127.0.0.1:8080，應該會看到一個狀態為 “ok” 的 JSON 回應。
注意事項
- 不要直接將 Flask 服務佈署到生產環境，因為內建伺服器不適合生產使用。
- 當模型遇到未知的影像類別時，它將始終選擇一個已知的類別。可以在訓練過程中新增一個額外的類別 “Unknown”，或者檢視最終 softmax 輸出的機率。
程式碼解析：
from flask import Flask, jsonify
from . import CatfishModel
from torchvision import transforms
import torch
import os
內容解密：
這段程式碼匯入了必要的函式庫，包括 Flask、jsonify 和 CatfishModel。CatfishModel 是我們定義的模型，用於分類別貓和魚的影像。
app = Flask(__name__)
@app.route("/")
def status():
return jsonify({"status": "ok"})
內容解密：
這段程式碼定義了一個 Flask 應用程式，並建立了一個路由 /，用於檢查服務的狀態。當存取這個路由時，它將傳回一個 JSON 回應，指示服務狀態為 “ok”。
@app.route("/predict", methods=['GET', 'POST'])
def predict():
img_url = request.image_url
img_tensor = open_image(BytesIO(response.content))
prediction = model(img_tensor)
predicted_class = CatfishClasses[torch.argmax(prediction)]
return jsonify({"image": img_url, "prediction": predicted_class})
內容解密：
這段程式碼定義了一個路由 /predict，用於接收包含影像 URL 的請求。它將下載影像，將其轉換為張量，並使用模型進行預測。然後，它將傳回一個 JSON 回應，指示影像的預測類別。
