邊緣設備上部署模型的限制之一——顯存占用：模型的參數量只是冰山一角

在邊緣設備上部署深度學習模型已成為趨勢，但資源限制是其核心挑戰之一。其中，顯存（或更廣義的內存）占用是開發者們必須仔細考量的重要因素。許多人認為顯存占用主要取決于模型的參數量，這種看法雖然沒錯，但并不全面。實際上，顯存的占用遠不止模型參數量那么簡單。

關于邊緣設備（Edge Device）的介紹，可以參見我的這一篇文章：Edge Device（邊緣設備）：連接物理世界與數字世界的橋梁

模型參數量：顯存占用的基石

首先，我們不得不承認，模型參數量確實是顯存占用的一個主要且直接的決定因素。

• 模型權重存儲：神經網絡中的每一個權重和偏置都需要存儲在顯存中。模型的參數越多，所需存儲的空間就越大。
• 數據類型影響：參數的數據類型也至關重要。例如，使用單精度浮點數（FP32）存儲每個參數需要4字節；而如果采用半精度（FP16）或8位整型量化（INT8），則分別只需要2字節或1字節。通過數據類型優化，即使參數量不變，也能顯著降低顯存需求。

---

顯存占用：遠不止參數量那么簡單

然而，如果我們只關注參數量，就會錯過顯存占用的其他幾個關鍵組成部分，這些部分有時甚至比參數量本身占用更大的顯存。

1. 中間激活（Intermediate Activations）

在模型推理過程中，數據流經每一層網絡時會產生中間結果，我們稱之為激活（Activations）或特征圖（Feature Maps）。這些中間激活需要臨時存儲在顯存中，以便后續層進行計算。

• 批處理大小（Batch Size）：批處理大小直接影響中間激活的顯存占用。批次越大，每次推理需要處理的樣本越多，每層生成的特征圖尺寸就越大，從而占據更多的顯存。
• 網絡深度與寬度：模型的層數（深度）和每層的神經元數量或特征圖通道數（寬度）都會直接影響中間激活的存儲需求。更深或更寬的網絡往往會產生海量的中間激活。

2. 優化器狀態（Optimizer State）

這一點主要與模型的訓練而非推理有關。在使用Adam、RMSprop等優化器進行模型訓練時，這些優化器會為每個模型參數維護額外的狀態變量（例如，動量和方差的估計值）。這些狀態變量通常與參數數量相同，并且同樣需要存儲在顯存中。在某些情況下，優化器狀態的顯存占用甚至可能使總顯存需求翻倍或更多。

3. 輸入和輸出數據

模型在進行推理之前，輸入數據（如圖像、音頻、傳感器數據等）需要加載到顯存中。輸入數據越大（例如高分辨率圖像），占用的顯存就越多。同樣，模型的最終輸出結果也需要一定的顯存來存儲，盡管這部分占用通常相對較小。

4. 輔助數據結構與操作開銷

除了上述幾點，還有一些其他因素也會貢獻顯存占用：

• 模型圖結構：深度學習框架（如TensorFlow Lite、PyTorch Mobile）在加載模型時，需要存儲模型的計算圖結構信息，這也會占用少量顯存。
• 緩沖區/緩存：某些特定的計算操作可能需要額外的緩沖區來進行數據轉換或中間計算。
• 量化和剪枝的潛在開銷：盡管量化和剪枝旨在減少模型大小，但在某些實現中，它們可能在處理過程中引入臨時的額外存儲開銷。

---

總結與思考

綜上所述，顯存占用是一個多維度的問題。一個擁有高參數量的模型，如果經過了高效的量化（Quantization）和剪枝（Pruning）等優化技術，并且在推理時采用較小的批處理大小，其最終的顯存占用可能反而低于一個參數量較少但未經優化、且使用大批處理大小的模型。

在邊緣設備上部署深度學習模型時，我們必須跳出“只看參數量”的誤區，全面評估和優化顯存占用。這包括對模型架構的選擇、數據類型的量化、批處理大小的設定以及中間激活的管理等方面進行深入考量，以確保模型能在有限的資源下穩定、高效地運行。

希望這篇文章能幫助你更好地理解邊緣設備上顯存占用的構成！😊