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