在試圖理解為什么我的 PyTorch 模型的最大批量大小受到限制時，我注意到不是模型本身，也不是將張量加載到使用最多記憶體的 GPU 上。大多數記憶體在第一次生成預測時用完，例如在訓練回圈中使用以下行：
output = model(images)
其中 images 是一些輸入張量，而 model 是我的 PyTorch 模型。在運行這條線之前，我有 9GB 的可用 GPU 記憶體，之后我下降到 2.5GB（然后在運行loss = criterion(outputs, labels).
兩個問題：
- 這是正常的嗎？
- 為什么會這樣？這么多記憶體是干什么用的？據我了解，模型已經加載，實際輸入張量在進行呼叫之前已經在 GPU 上。輸出張量本身不可能那么大。它與存盤計算圖有關嗎？
uj5u.com熱心網友回復：
這是正常的：這里的關鍵是，如果您想通過后向模式微分計算梯度，則必須存盤所有中間張量（整個計算圖）。您可以通過使用.no_grad背景關系管理器來避免這種情況：
with torch.no_grad():
output = model(images)
您會觀察到使用的記憶體要少得多，因為不會存盤計算圖。但這也意味著你不能再計算導數了。但是，如果您只想評估模型而不需要任何優化，這是標準方法。
如果您仍想優化，有一種方法可以減少記憶體消耗，稱為checkpointing。每當您在后向傳遞中需要一個中間張量時，它將從輸入（或實際上從最后一個“檢查點”）再次計算，而不會將中間張量存盤到該張量。但這只是在計算上更昂貴。您正在用記憶體與計算時間進行交易。
轉載請註明出處，本文鏈接：https://www.uj5u.com/shujuku/482347.html