1. 一個讓很多人困惑的現(xiàn)象模型文件明明只有幾十兆加載到內(nèi)存里卻占了幾百兆甚至幾個G這事兒我遇到過太多次了。最早做嵌入式端側(cè)推理的時候一個量化后不到20MB的模型跑在512MB內(nèi)存的開發(fā)板上程序一啟動系統(tǒng)就開始瘋狂swap當(dāng)時百思不得其解。后來把賬一筆一筆算清楚才發(fā)現(xiàn)問題根本不在模型文件本身。這篇文章就是要把卷積神經(jīng)網(wǎng)絡(luò)運(yùn)行時的內(nèi)存賬本徹底攤開。核心圍繞三個關(guān)鍵詞參數(shù)量、MACs乘加運(yùn)算次數(shù)、FP32。如果你正在做模型部署、端側(cè)推理優(yōu)化或者單純想搞清楚“為什么我的模型這么吃內(nèi)存”這篇內(nèi)容應(yīng)該能幫你把思路理順。我會從內(nèi)存到底被誰吃了講起把卷積層的三筆賬——權(quán)重賬、激活賬、工作區(qū)賬——一筆一筆算給你看最后給出可以直接抄的優(yōu)化方案。需要提前說明的是文中涉及的具體數(shù)值和配置是基于常見工程實踐給出的參考值不同框架、不同硬件平臺會有差異但計算邏輯和優(yōu)化思路是通用的。2. 先搞清楚內(nèi)存到底被誰吃了2.1 模型文件大小不等于運(yùn)行時內(nèi)存很多人有一個直覺模型文件多大運(yùn)行時就應(yīng)該占多大內(nèi)存。這個直覺在推理階段大致成立但有幾個關(guān)鍵前提——你得把權(quán)重完整加載進(jìn)內(nèi)存而且推理過程中產(chǎn)生的中間張量不能忽略。模型文件比如.pt、.onnx、.tflite存儲的是權(quán)重參數(shù)通常以FP32格式保存。一個包含100萬個參數(shù)的卷積網(wǎng)絡(luò)FP32格式下文件大小就是 100萬 × 4字節(jié) 4MB。但運(yùn)行時除了這4MB權(quán)重還需要激活值A(chǔ)ctivations每一層卷積的輸出特征圖前向傳播過程中必須保留至少當(dāng)前層和下一層需要工作區(qū)Workspace某些算子實現(xiàn)需要臨時緩沖區(qū)比如im2col展開、Winograd變換框架運(yùn)行時開銷內(nèi)存分配器、線程池、算子注冊表等所以一個4MB的模型文件運(yùn)行時占40MB甚至400MB完全有可能。2.2 三筆賬的框架我把卷積層的內(nèi)存消耗拆成三筆賬賬目來源典型占比是否可優(yōu)化權(quán)重賬卷積核參數(shù)10%-30%可量化、可剪枝激活賬特征圖輸出40%-70%可復(fù)用、可重計算工作區(qū)賬算子臨時緩沖10%-40%可調(diào)算法、可限制這三筆賬加起來才是卷積層真正的內(nèi)存開銷。下面逐筆拆解。2.3 為什么FP32是默認(rèn)選項FP32單精度浮點用4個字節(jié)存儲一個數(shù)1位符號位、8位指數(shù)位、23位尾數(shù)位。深度學(xué)習(xí)訓(xùn)練默認(rèn)用FP32因為梯度更新對數(shù)值精度敏感FP16容易溢出或下溢。但推理階段FP32往往不是最優(yōu)選擇。FP16只需要2字節(jié)INT8只需要1字節(jié)。一個FP32下占100MB的模型轉(zhuǎn)成INT8后權(quán)重部分直接降到25MB。這也是為什么端側(cè)部署幾乎都會做量化。注意量化不是無損的。INT8量化后精度通常掉0.5%-2%具體取決于校準(zhǔn)數(shù)據(jù)集的質(zhì)量和量化方案。對精度敏感的任務(wù)如醫(yī)學(xué)圖像分割要謹(jǐn)慎。3. 第一筆賬權(quán)重到底占多少3.1 參數(shù)量怎么算卷積層的參數(shù)量公式很簡單參數(shù)量 卷積核數(shù)量 × 輸入通道數(shù) × 卷積核高 × 卷積核寬 卷積核數(shù)量偏置舉個例子一個標(biāo)準(zhǔn)的3×3卷積輸入64通道輸出128通道參數(shù)量 128 × 64 × 3 × 3 128 73728 128 73856FP32下占內(nèi)存73856 × 4字節(jié) ≈ 288KB?？雌饋聿淮蟮粋€ResNet-50有53個卷積層加起來就是2500萬參數(shù)左右FP32下約100MB。3.2 參數(shù)量大不等于計算量大這里有一個常見的認(rèn)知誤區(qū)參數(shù)量大的層計算量不一定大。深度可分離卷積就是典型例子。標(biāo)準(zhǔn)卷積3×3輸入64輸出128參數(shù)量73856MACs128 × 64 × 3 × 3 × H × WH、W是輸出特征圖尺寸深度可分離卷積拆成兩步逐通道卷積64 × 3 × 3 576參數(shù)逐點卷積128 × 64 × 1 × 1 8192參數(shù)總參數(shù)量8768只有標(biāo)準(zhǔn)卷積的12%但MACs的下降比例取決于特征圖尺寸。當(dāng)特征圖較大時逐點卷積的MACs占比會上升。所以參數(shù)量和MACs是兩個獨立的賬本優(yōu)化時要分開看。3.3 權(quán)重內(nèi)存的優(yōu)化手段權(quán)重賬的優(yōu)化最直接量化FP32轉(zhuǎn)FP16內(nèi)存減半轉(zhuǎn)INT8內(nèi)存降到四分之一。這是性價比最高的手段。剪枝去掉不重要的連接稀疏化存儲。但稀疏矩陣的實際加速比取決于硬件支持。權(quán)重共享某些網(wǎng)絡(luò)結(jié)構(gòu)如MobileNet本身參數(shù)量就小不需要額外優(yōu)化。我實測下來FP16量化對大多數(shù)視覺模型精度影響在0.1%以內(nèi)幾乎可以無腦上。INT8需要仔細(xì)做校準(zhǔn)但收益也最大。4. 第二筆賬激活值才是內(nèi)存大戶4.1 激活值為什么占內(nèi)存前向傳播時每一層卷積的輸出特征圖必須保留因為下一層要用。假設(shè)一個卷積層輸出128通道特征圖尺寸是64×64FP32下占內(nèi)存128 × 64 × 64 × 4字節(jié) 2MB一個ResNet-50有幾十個這樣的層如果全部保留激活值內(nèi)存輕松超過100MB。但實際推理時并不是所有層的激活值都需要同時保留——只有當(dāng)前層和下一層需要。所以理論上激活值內(nèi)存可以控制在兩層的大小。但問題在于很多框架為了調(diào)試方便或者圖優(yōu)化不徹底會保留更多中間結(jié)果。這就是為什么同樣的模型不同框架跑出來的內(nèi)存占用差異巨大。4.2 激活值的內(nèi)存復(fù)用策略激活值內(nèi)存優(yōu)化的核心思路是復(fù)用。具體來說原地操作In-placeReLU、BN等逐元素操作可以直接覆蓋輸入內(nèi)存不需要額外分配。內(nèi)存池預(yù)分配一塊大內(nèi)存不同層的激活值輪流使用同一塊區(qū)域。計算換內(nèi)存某些層的激活值不保存反向傳播時重新計算。推理階段不需要反向傳播所以這條主要針對訓(xùn)練。以PyTorch為例torch.no_grad()下推理框架會自動做一定程度的內(nèi)存復(fù)用。但如果你用ONNX Runtime或者TensorRT它們有更激進(jìn)的內(nèi)存池策略激活值內(nèi)存可以壓得更低。4.3 一個真實的激活值計算案例假設(shè)你要部署一個輸入224×224×3的圖像分類模型第一層卷積輸出64通道特征圖112×112激活值 64 × 112 × 112 × 4字節(jié) 3.2MB第二層輸出128通道特征圖56×56激活值 128 × 56 × 56 × 4字節(jié) 1.6MB看起來每層都不大但如果有50層且框架沒有做內(nèi)存復(fù)用總激活值內(nèi)存就是幾十MB到上百M(fèi)B。這就是為什么模型文件只有幾十MB運(yùn)行時卻占幾百M(fèi)B。實操心得用torch.cuda.memory_summary()或者ONNX Runtime的profiling工具可以精確看到每一層的內(nèi)存分配情況。我一般會先跑一遍profiling找出內(nèi)存峰值出現(xiàn)在哪一層再針對性優(yōu)化。5. 第三筆賬工作區(qū)內(nèi)存容易被忽略5.1 工作區(qū)是什么工作區(qū)Workspace是算子實現(xiàn)過程中需要的臨時內(nèi)存。最典型的是im2col把卷積運(yùn)算轉(zhuǎn)換成矩陣乘法時需要把輸入特征圖展開成一個矩陣。這個展開后的矩陣就是工作區(qū)。假設(shè)輸入64通道3×3卷積輸出特征圖56×56im2col矩陣大小 64 × 3 × 3 × 56 × 56 × 4字節(jié) ≈ 7.2MB這還只是一層的工作區(qū)。如果框架沒有復(fù)用工作區(qū)內(nèi)存多層疊加起來就很可觀。5.2 不同算法的工作區(qū)開銷卷積的實現(xiàn)算法有很多種工作區(qū)開銷差異很大算法工作區(qū)開銷計算效率適用場景直接卷積低低小卷積核im2col GEMM高高大特征圖Winograd中很高3×3卷積FFT高中大卷積核Winograd是3×3卷積的常用優(yōu)化它通過變換減少乘法次數(shù)但需要額外的變換矩陣存儲。FFT卷積在大卷積核如7×7以上時效率高但工作區(qū)開銷也大。5.3 工作區(qū)內(nèi)存的限制與調(diào)優(yōu)很多推理框架允許你設(shè)置工作區(qū)內(nèi)存上限。比如TensorRT的workspace_size參數(shù)ONNX Runtime的arena_extend_strategy。設(shè)置得太小框架可能回退到低效算法設(shè)置得太大內(nèi)存占用高。我的經(jīng)驗是先設(shè)一個較大的值跑通然后用profiling工具看實際用了多少再逐步調(diào)小。一般端側(cè)部署會把工作區(qū)限制在幾十MB以內(nèi)服務(wù)器端可以放寬到幾百M(fèi)B。6. 把三筆賬合起來算一個完整例子6.1 模型設(shè)定假設(shè)一個簡化的卷積網(wǎng)絡(luò)輸入1×3×224×224Conv13→643×3stride1padding1Conv264→1283×3stride2padding1Conv3128→2563×3stride2padding1全局平均池化 全連接輸出6.2 權(quán)重賬Conv164×3×3×3 64 1792參數(shù) Conv2128×64×3×3 128 73856參數(shù) Conv3256×128×3×3 256 295168參數(shù) 全連接假設(shè)輸出1000類256×1000 1000 257000參數(shù)總參數(shù)量約62.8萬FP32下約2.5MB。6.3 激活賬Conv1輸出64×224×224×4 12.8MB Conv2輸出128×112×112×4 6.4MB Conv3輸出256×56×56×4 3.2MB如果全部保留22.4MB。如果只保留當(dāng)前層和下一層最大約12.8MB。6.4 工作區(qū)賬Conv1的im2col矩陣3×3×3×224×224×4 ≈ 1.8MB Conv2的im2col矩陣64×3×3×112×112×4 ≈ 28.9MB Conv3的im2col矩陣128×3×3×56×56×4 ≈ 14.5MB工作區(qū)峰值約28.9MB。6.5 總賬賬目內(nèi)存占用權(quán)重2.5MB激活值12.8MB優(yōu)化后工作區(qū)28.9MB框架開銷10-20MB合計約55-65MB模型文件2.5MB運(yùn)行時55MB以上差了20多倍。這就是“模型文件很小運(yùn)行為什么還吃內(nèi)存”的答案。7. 優(yōu)化實戰(zhàn)把內(nèi)存降下來7.1 量化權(quán)重和激活值FP32轉(zhuǎn)FP16權(quán)重和激活值內(nèi)存直接減半。轉(zhuǎn)INT8再減半。但I(xiàn)NT8需要校準(zhǔn)且不是所有算子都支持。我的建議服務(wù)器端FP16足夠精度損失可忽略端側(cè)INT8配合量化感知訓(xùn)練效果更好極端場景混合精度敏感層用FP16其他用INT87.2 限制工作區(qū)內(nèi)存以O(shè)NNX Runtime為例import onnxruntime as ort options ort.SessionOptions() options.enable_cpu_mem_arena True options.arena_extend_strategy kSameAsRequested options.add_session_config_entry(session.intra_op.allow_spinning, 0) session ort.InferenceSession(model.onnx, options)kSameAsRequested策略會讓內(nèi)存池按需增長而不是一次性分配一大塊。實測下來內(nèi)存峰值可以降低20%-30%。7.3 使用內(nèi)存復(fù)用和原地操作PyTorch推理時import torch model.eval() with torch.no_grad(): output model(input_tensor)torch.no_grad()會關(guān)閉梯度計算減少中間變量。model.eval()會切換BN和Dropout到推理模式避免額外的統(tǒng)計量存儲。7.4 選擇合適的內(nèi)存分配器不同框架的內(nèi)存分配器策略不同。PyTorch默認(rèn)用caching allocator會緩存已分配的內(nèi)存塊減少頻繁分配釋放的開銷但會導(dǎo)致內(nèi)存占用偏高。如果內(nèi)存緊張可以設(shè)置環(huán)境變量export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128這會讓分配器更積極地釋放大塊內(nèi)存。8. 常見問題與排查技巧8.1 為什么模型加載后內(nèi)存比文件大很多這是最常見的問題。原因通常是權(quán)重從FP32轉(zhuǎn)成了其他格式但加載時又轉(zhuǎn)回FP32框架在加載時做了圖優(yōu)化生成了額外的中間表示內(nèi)存分配器預(yù)分配了比實際需要更多的內(nèi)存排查方法用psutil或者框架自帶的內(nèi)存profiling工具看內(nèi)存是在哪個階段漲上去的。8.2 推理過程中內(nèi)存持續(xù)增長這通常是內(nèi)存泄漏。常見原因每次推理都創(chuàng)建新的Session或Context沒有釋放輸入輸出張量沒有復(fù)用每次都在新分配某些算子的工作區(qū)沒有正確釋放排查方法跑多次推理觀察內(nèi)存是否線性增長。如果是用tracemalloc或者valgrind定位泄漏點。8.3 量化后內(nèi)存沒降多少可能原因只量化了權(quán)重激活值還是FP32框架在運(yùn)行時做了反量化又變回FP32工作區(qū)內(nèi)存沒有跟著量化排查方法看框架的量化文檔確認(rèn)是否支持全量化權(quán)重激活值都量化。8.4 常見問題速查表問題可能原因解決方法模型加載后內(nèi)存暴漲格式轉(zhuǎn)換、圖優(yōu)化用profiling定位階段推理時內(nèi)存持續(xù)增長內(nèi)存泄漏復(fù)用Session和Tensor量化后內(nèi)存沒降只量化了權(quán)重開啟全量化工作區(qū)內(nèi)存過大算法選擇不當(dāng)限制workspace_size激活值內(nèi)存過高沒有內(nèi)存復(fù)用開啟內(nèi)存池、原地操作避坑技巧我一般會在模型部署前做一個“內(nèi)存預(yù)算表”把權(quán)重、激活值、工作區(qū)的預(yù)期內(nèi)存都列出來然后跟實際測量值對比。如果偏差超過20%就說明有隱藏的內(nèi)存開銷需要進(jìn)一步排查。9. 一些個人經(jīng)驗踩過幾次坑之后我養(yǎng)成了一個習(xí)慣任何模型部署前先算三筆賬。權(quán)重賬用參數(shù)量公式算激活賬用特征圖尺寸算工作區(qū)賬用im2col矩陣大小估算。三筆賬加起來再乘以1.5的安全系數(shù)就是內(nèi)存預(yù)算。這個習(xí)慣幫我避免了很多次“上線后OOM”的尷尬。有一次一個模型文件只有8MB我算下來運(yùn)行時需要120MB實際部署時果然在128MB內(nèi)存的設(shè)備上跑得很勉強(qiáng)。后來把工作區(qū)限制到32MB激活值用FP16總內(nèi)存降到70MB才穩(wěn)定下來。另外不同框架的內(nèi)存行為差異很大。同樣的模型PyTorch可能占200MBTensorRT可能只占80MB。所以選框架時內(nèi)存占用也是一個重要考量因素。端側(cè)部署我一般優(yōu)先考慮TensorRT、NCNN、MNN這些專門優(yōu)化過的框架。最后再分享一個小技巧如果你用的是ONNX模型可以用onnxsim做圖簡化去掉冗余算子有時候能減少10%-20%的內(nèi)存占用。這個工具用起來很簡單pip install onnxsim onnxsim input.onnx output.onnx模型文件小了運(yùn)行時內(nèi)存也可能跟著降。雖然不保證每次都有效但值得一試。