做深度學(xué)習(xí)的同學(xué)應(yīng)該都有過這種經(jīng)歷啟動第一個 TensorFlow 訓(xùn)練腳本后打開nvidia-smi一看明明模型不大顯存卻直接占到了接近滿載。這不是模型吃了多少的問題而是 TensorFlow 在默認情況下會對當(dāng)前可見的 GPU 做一整塊預(yù)分配把整張卡的顯存當(dāng)作自己家的倉庫。這個行為在 1.x 和 2.x 里都存在只是處理機制和配置入口很不一樣。這篇文章就把 TensorFlow 2.x 和 1.x 限制顯存的主流方法、底層參數(shù)含義、踩坑點一次講透?？赐曛竽銜靼诪槭裁聪拗骑@存不等于限制模型大小為什么allow_growth有時不起作用以及多進程共享一塊 GPU 時到底該怎么分配。1. 默認行為一張卡被單進程整套帶走1.1 TensorFlow 為什么會把顯存全部占滿TensorFlow 的顯存分配主要交給內(nèi)部的 BFCBest-Fit with Coalescing分配器來做。BFC 分配器的工作方式和操作系統(tǒng)里的內(nèi)存管理器很像它先把一大塊顯存接管下來然后在內(nèi)部切成不同大小的塊供算子按需申請。但關(guān)鍵在這一步如果不做任何配置TensorFlow 在初始化 GPU 設(shè)備時會直接把當(dāng)前進程可見的整張卡顯存全部申請到自己的池子里。對驅(qū)動層來說這張卡已經(jīng)名義上被這個進程占完了哪怕你現(xiàn)在只跑了一個幾 MB 的小張量nvidia-smi上顯示的占用率也會是 90% 以上。很多剛?cè)腴T的朋友會誤以為這是模型太吃顯存于是開始瘋狂降低 batch size。其實模型本身占用的顯存可能只有 12GB剩下的一大半都被預(yù)分配池子吞了。你限制顯存限制的并不是算子使用的顯存而是 BFC 分配器從顯卡驅(qū)動那里申請的上限。1.2 真正需要限制顯存的四個場景第一個場景是多人共用 GPU。實驗室或公司里一張卡經(jīng)常有幾個人在跑任務(wù)如果每個人啟動 TensorFlow 時都默認占滿全卡后來的人只能干等。合理做法是每個人各自設(shè)一個上限比如一個人分 4GB另一個人分 4GB。第二個場景是單卡多任務(wù)。你可能要同時跑一個訓(xùn)練任務(wù)和一個推理任務(wù)或者一個腳本里加載多個模型。只要每個任務(wù)用到的顯存都收斂就必須給它們分別畫好邊界否則第二個模型初始化時就會把第一個模型擠掉。第三個場景是 Jupyter Notebook 這類交互式環(huán)境。你反復(fù)執(zhí)行代碼單元格環(huán)境會保留歷史計算圖上的張量顯存經(jīng)常只漲不跌。加上 TensorFlow 默認把整卡占住寫課程作業(yè)都會出問題。所以交互式環(huán)境里幾乎必須限制顯存。第四個場景是模型部署。線上服務(wù)進程需要長時間運行顯存只能限制在一個穩(wěn)定水位。如果讓推理服務(wù)按需無限增長運行三五天后顯存就可能被緩存頂?shù)缴舷迣?dǎo)致請求崩潰。1.3 一個容易混的概念進程顯存不等于模型顯存你限制顯存之后nvidia-smi里看到的進程占用數(shù)字并不等于模型參數(shù)和中間特征圖的大小。進程要正常運行還需要維持 CUDA context、cuDNN workspace、以及部分庫的常駐內(nèi)存。這部分開銷通常有一兩百 MB也就是你設(shè)置 2GB 上限時nvidia-smi里看到的進程占用可能會在 2.1GB 到 2.3GB 左右。所以我在實際配置時都會先按模型真實需求 上下文開銷來算不會直接卡到模型需求值上。這一步雖然簡單但很多同學(xué)第一次寫限制配置時都會忽略導(dǎo)致明明給了足夠空間卻總是差一點。2. TensorFlow 2.x 限制顯存的標(biāo)準(zhǔn)操作2.1 增長式分配set_memory_growthTensorFlow 2.x 官方的推薦入口是tf.config.experimental。最常用的方法是讓顯存按需增長也就是模型實際申請多少TensorFlow 就占多少不要一次性整卡占滿。對應(yīng)的代碼是import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) if gpus: try: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) except RuntimeError as e: print(e)這里的關(guān)鍵點在于set_memory_growth必須在任何算子執(zhí)行、任何模型構(gòu)建之前調(diào)用。只要 TensorFlow 已經(jīng)初始化了 GPU 上下文再設(shè)置就會觸發(fā)RuntimeError因為分配策略已經(jīng)定死了。另外需要注意這段代碼只是開啟用多少申請多少的模式并沒有硬性規(guī)定上限。如果你的模型恰好需要 11GB而卡上只有 11GB那它還是會占滿增長模式解決的是不用就別占著的問題不是最多只能用多少的問題。2.2 硬性上限用虛擬設(shè)備配置指定內(nèi)存大小如果你希望一個進程嚴(yán)格限制在固定大小比如最多只能用 4GB剩下的 2GB 要留給另一個任務(wù)那就需要用到set_virtual_device_configuration。import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) if gpus: tf.config.experimental.set_virtual_device_configuration( gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit4096)] )memory_limit的單位是 MB所以4096就是 4GB。配置完之后TensorFlow 的 BFC 分配器只會在這個 4GB 的池子里做分配超過就會直接報 OOM。我的個人建議是訓(xùn)練腳本用set_memory_growth更省事因為只要顯存不沖突它不會影響性能推理服務(wù)或需要穩(wěn)定復(fù)現(xiàn)的腳本用memory_limit更靠譜因為固定水位好監(jiān)控、好告警、好做容量規(guī)劃。2.3 環(huán)境變量方式TF_FORCE_GPU_ALLOW_GROWTH除了在代碼里設(shè)置TensorFlow 2.x 還支持通過環(huán)境變量TF_FORCE_GPU_ALLOW_GROWTH強制開啟增長式分配。TF_FORCE_GPU_ALLOW_GROWTHtrue python train.py這種方式的優(yōu)點是能不動訓(xùn)練代碼直接通過啟動命令控制進程的顯存行為。在跑同學(xué)給的腳本、或者別人封裝的模型倉庫時我會先看對方代碼里有沒有寫死顯存配置如果沒寫就優(yōu)先用這個環(huán)境變量臨時接管。不過環(huán)境變量也有一些局限性。部分版本對它的支持不夠完整或者會被代碼里的設(shè)置覆蓋。所以我只是把它當(dāng)作應(yīng)急手段代碼里能設(shè)置的話還是老老實實在代碼里設(shè)置。2.4 增長和硬上限到底怎么選set_memory_growth對應(yīng)動態(tài)擴張memory_limit對應(yīng)靜態(tài)封頂。選擇時主要看運行環(huán)境。如果你是在自己的電腦上訓(xùn)練顯存怎么折騰都不會影響別人那直接用增長式最舒服既能充分利用顯存又不會開機就占滿。如果你是跑共享服務(wù)器且同一張卡上還有其他人的任務(wù)請務(wù)必使用memory_limit固定上限。否則就算設(shè)置了增長式只要你的模型稍微大一點它還是會迅速把卡占滿和默認行為幾乎沒差別。如果是部署在線服務(wù)風(fēng)控和容量規(guī)劃很關(guān)鍵建議直接用memory_limit并把上限設(shè)置在峰值 20%的狀態(tài)。比如某個模型峰值要 1.5GB我通常會給 1.8GB 左右留一點上下文開銷和臨時算子緩存。3. TensorFlow 1.x 的顯存限制從 GPUOptions 到 Session 配置3.1 per_process_gpu_memory_fraction按比例設(shè)置顯存上限TensorFlow 1.x 時代還沒有tf.config.experimental這套新 API顯存限制基本集中在GPUOptions和ConfigProto上。最常見的是通過比例限制比如把當(dāng)前進程可用的顯存限制在總顯存的 40%import tensorflow as tf gpu_options tf.GPUOptions(per_process_gpu_memory_fraction0.4) config tf.ConfigProto(gpu_optionsgpu_options) sess tf.Session(configconfig)這么做之后TensorFlow 啟動時就不會一次性占滿整張卡而是先申請大約 40% 的顯存。隨著算子執(zhí)行它會在這一份池子里繼續(xù)分配。合理的地方在于參數(shù)直觀比如一張 10GB 的卡0.4 大約就是 4GB。但注意這個比例是基于當(dāng)前進程可見的顯存總量來算的如果你先用CUDA_VISIBLE_DEVICES屏蔽了部分卡那比例對應(yīng)的是可見卡的總量。3.2 allow_growth按需申請但小心碎片化1.x 里還有另一個經(jīng)典參數(shù)allow_growth。設(shè)置它之后TensorFlow 不再一開始就把整張卡包下來而是根據(jù)每次算子的需求一點點向驅(qū)動申請顯存。import tensorflow as tf gpu_options tf.GPUOptions(allow_growthTrue) config tf.ConfigProto(gpu_optionsgpu_options) sess tf.Session(configconfig)它的效果和 2.x 里的set_memory_growth(True)很像但老版本實現(xiàn)上有更多問題。最典型的是碎片化BFC 分配器按需申請內(nèi)存時會因為各個張量的大小不一在顯存池里留下很多不連續(xù)的小塊。運行一段時間后可能整體顯存還有幾百 MB 空閑但找不到一塊足夠大的連續(xù)空間來放新張量結(jié)果照樣 OOM。在 1.x 里如果我還需要配合 Keras 使用會把 Session 塞給 Kerasimport tensorflow as tf config tf.ConfigProto() config.gpu_options.allow_growth True sess tf.Session(configconfig) tf.keras.backend.set_session(sess)這是那個時代很常見的組合寫法?，F(xiàn)在回看backend.set_session這套機制早就被新 API 取代了但如果你還在維護老代碼這仍然是必須認識的老朋友。3.3 在 2.x 中復(fù)用 1.x 的配置姿勢有些老項目短時間內(nèi)沒法完全遷移到 2.x 的新 API這就需要用到compat.v1兼容層。比如下面這種方式import tensorflow as tf tf.compat.v1.disable_eager_execution() gpu_options tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction0.3) config tf.compat.v1.ConfigProto(gpu_optionsgpu_options) sess tf.compat.v1.Session(configconfig)這里先禁用 Eager 模式是為了讓 Session 機制兼容 1.x 的代碼習(xí)慣。之后模型的構(gòu)建和訓(xùn)練方式基本還是原來的寫法。不過說實話如果項目還能改我不太建議長期停在 compat.v1 的兼容支線上。兼容層只是過渡方案新寫的代碼直接用 2.x 原生的tf.config.experimental更干凈。老代碼需要快速跑通的話用compat.v1能省不少事但不要把它當(dāng)成長期架構(gòu)。4. 實測不同限制參數(shù)下的顯存占用和 OOM 邊界4.1 用 nvidia-smi 觀察真實占用在討論各種配置之前先確認觀察工具。nvidia-smi最常用的實時查看命令是watch -n 1 nvidia-smi如果想拿到格式化的占用數(shù)據(jù)方便記錄到日志里可以用nvidia-smi --query-gpuindex,memory.used,memory.total,utilization.gpu --formatcsv -l 5-l 5表示每 5 秒刷新一次。實際調(diào)優(yōu)時我更推薦看一下進程級顯存而不是只看總量nvidia-smi --query-compute-appspid,used_memory,process_name --formatcsv這樣能直接看到哪個 PID 占了多少不會被卡上其他進程干擾。4.2 一張 6GB 卡上的典型對比我拿一個常見的分割模型做參考輸入大約512x512、batch size 為 8模型峰值需求大約在 3.4GB 左右。使用 TensorFlow 2.x分別用不同方式設(shè)置顯存典型表現(xiàn)如下設(shè)置方式顯存配置進程實際占用約訓(xùn)練速度表現(xiàn)不設(shè)置無占滿整卡正常增長式set_memory_growth(True)約 3.5GB正常虛擬設(shè)備memory_limit4096約 4.2GB正常虛擬設(shè)備memory_limit2048約 2.1GB大概率 OOM1.x 比例per_process_gpu_memory_fraction0.5約 3.0GB正常1.x 比例per_process_gpu_memory_fraction0.2約 1.2GB直接 OOM可以看到只要限制值高于模型峰值加上約 0.5GB 的緩沖訓(xùn)練速度基本不會受什么影響。真正的風(fēng)險出現(xiàn)在限制值低于實際需求時。4.3 OOM 的具體表現(xiàn)不是等待而是立刻失敗有相當(dāng)一部分人以為顯存不夠時 TensorFlow 會像 CPU 內(nèi)存那樣先等一會兒或者和其他進程商量一下。實際完全不是這樣。當(dāng) BFC 分配器向驅(qū)動申請一塊連續(xù)顯存失敗時進程會立刻拋出ResourceExhaustedError并中斷訓(xùn)練。報錯信息大致長這樣tensorflow.python.framework.errors_impl.ResourceExhaustedError: OOM when allocating tensor with shape [8,512,512,3] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_bfc看到GPU_bfc就知道是 BFC 分配器沒有足夠空間。這個報錯本身很有價值它能暴露你模型某個算子的具體張量形狀直接定位到瓶頸。排查時不要只盯著 batch size還要關(guān)注特征圖尺寸、通道數(shù)、以及可能臨時創(chuàng)建的中間張量。4.4 為什么限制過小會引發(fā)連鎖性的變慢如果把顯存卡在模型峰值附近例如需求 3.4GB限制給 3.6GB你可能會發(fā)現(xiàn)訓(xùn)練速度沒崩但比正常情況下慢了一截。原因主要有兩類。第一類是顯存不夠時TensorFlow 會頻繁進行塊合并和碎片整理。BFC 里的塊被拆開再合并本身就是額外開銷每次算子執(zhí)行都可能觸發(fā)。第二類是部分算子發(fā)現(xiàn)顯存緊張會退回到性能更差的實現(xiàn)或發(fā)起 host-device 拷貝。問題不會立刻以 OOM 形式出現(xiàn)而是讓訓(xùn)練一步慢半拍非常隱蔽。所以在設(shè)計限制值時我一般留出 20% 的余量。比如模型峰值 3.4GB就限制在 4GB 或 4.5GB。這樣既能保證穩(wěn)定又不會給同卡其他任務(wù)帶來太多負擔(dān)。5. 多進程共享 GPU 時怎么規(guī)劃顯存才合理5.1 CUDA_VISIBLE_DEVICES 先做物理隔離多個進程分配同一張卡之前先想清楚每個進程到底應(yīng)該看到哪張卡。CUDA_VISIBLE_DEVICES是在驅(qū)動層做可見性過濾的環(huán)境變量設(shè)置之后 TensorFlow 只能看到列表中的卡。import os os.environ[CUDA_VISIBLE_DEVICES] 1,2 import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) print(gpus)注意這個環(huán)境變量必須在import tensorflow之前設(shè)置。因為 TensorFlow 導(dǎo)入時會讀取這個變量來確定設(shè)備可見性晚了就來不及了。設(shè)置1,2之后物理上的 1 號卡、2 號卡在進程內(nèi)會變成邏輯0和1。也就是說list_physical_devices(GPU)返回索引是從 0 開始的邏輯索引不是物理索引。很多人踩過這個坑以為自己設(shè)置的是gpu:1實際是物理 0 號卡結(jié)果把別人的任務(wù)擠爆了。5.2 訓(xùn)練進程和推理進程的參數(shù)建議如果一張卡要同時跑一個訓(xùn)練任務(wù)和一個推理任務(wù)我比較推薦的安排是訓(xùn)練任務(wù)用set_memory_growth推理任務(wù)用memory_limit固定上限。原因在于訓(xùn)練過程的顯存需求會隨 batch size、輸入尺寸和模型結(jié)構(gòu)波動。用增長式至少能保證模型在峰值時不會因為提前鎖死而 OOM。而推理任務(wù)通常輸入波動不大顯存需求相對穩(wěn)定更適合用固定上限來保證服務(wù)質(zhì)量。假如同一個模型推理峰值是 1GB我習(xí)慣給 1.3GB 或 1.5GB。這樣即使輸入偶爾變大一點也不至于掛掉同時給訓(xùn)練任務(wù)留出足夠的剩余空間。5.3 多卡訓(xùn)練時逐卡設(shè)置上限用tf.distribute.MirroredStrategy做單機多卡訓(xùn)練時顯存配置需要在策略創(chuàng)建之前完成。官方推薦是在代碼最前面遍歷所有物理 GPU并對每張卡分別設(shè)置import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) for gpu in gpus: tf.config.experimental.set_virtual_device_configuration( gpu, [tf.config.experimental.VirtualDeviceConfiguration(memory_limit4096)] ) strategy tf.distribute.MirroredStrategy()如果每張卡型號一致直接循環(huán)設(shè)置即可。如果型號不一致比如一張 8GB、一張 16GB就要分別獲取每張卡并逐個給不同的memory_limit否則策略初始化時會因為顯存不一致出各種莫名其妙的報錯。6. 避坑清單限制無效、顯存只漲不降、速度異常6.1 明明設(shè)置了限制為什么還是占滿這類問題我見過太多次幾乎每次都是配置時機不對。常見情況有三種。第一種代碼里先構(gòu)建了模型或執(zhí)行了算子再去調(diào)用set_memory_growth。此時 GPU 上下文已初始化TensorFlow 大概率會拋RuntimeError但也有老版本不報錯、直接忽略配置的情況。所以配置要放在所有 TensorFlow 操作之前。第二種模塊導(dǎo)入階段就觸發(fā)了 GPU 初始化。有些庫在import時并不構(gòu)建模型但會預(yù)加載 cuDNN 或創(chuàng)建上下文。如果你在導(dǎo)入之后才配置實際上已經(jīng)晚了。第三種使用多進程庫時子進程里的 TensorFlow 不會繼承主進程的顯存配置。如果主進程設(shè)置好了限制但訓(xùn)練是在spawn出來的子進程里跑的子進程仍會走默認的整卡占滿邏輯。多進程代碼需要考慮在子進程入口處重新執(zhí)行一次顯存配置。6.2 訓(xùn)練中顯存只漲不降是泄漏嗎很多人跑長時間訓(xùn)練發(fā)現(xiàn)nvidia-smi里進程占用從 3GB 慢慢漲到 4GB、5GB于是懷疑是顯存泄漏。大多數(shù)情況下這不是真正意義上的泄漏而是 BFC 分配器的緩存機制在起作用。TensorFlow 在運行時會把用過的顯存塊留在池子里方便后續(xù)算子復(fù)用。這個池子的總占用會隨著不同模塊的負載浮動但不會在算子釋放后立刻歸還給驅(qū)動。真正判斷泄漏的方法是看同一段循環(huán)里持續(xù)構(gòu)造新張量、反復(fù)執(zhí)行模型進程占用是否單調(diào)遞增且永不回頭。如果只是前期從低到高爬升之后就穩(wěn)定在一個水位上那只是緩存池擴容不是泄漏。6.3 限制顯存后訓(xùn)練莫名其妙變慢如果確認顯存限制沒有觸發(fā) OOM但訓(xùn)練速度明顯下降優(yōu)先檢查你是不是把限制值壓得太緊了。當(dāng) BFC 剩余空間很小每個算子申請顯存時都面臨剛剛好夠、但不夠整塊的情況分配器就要頻繁做塊切分和合并。這個過程非常消耗時間。更麻煩的是cuDNN 在自動選擇卷積實現(xiàn)時會搜索多個候選算法有些算法要求額外工作區(qū)空間。顯存限制太緊會讓 cuDNN 選不到最優(yōu)算法只能退而求其次整個模型運行速度都會下降。正確的做法是給顯存留足緩沖。不要追求剛好夠用顯存本質(zhì)上也是性能資源。6.4 一條額外的注意點限制夠用時不要過度限制最后說一句我自己在實踐中的體會顯存限制的價值是讓多任務(wù)能共存不是為了省錢。如果整張卡只有你的一個任務(wù)就算模型只需 2GB也沒必要強制限制到 4GB 以內(nèi)。強行限制反而可能讓 cuDNN 和 TensorFlow 的調(diào)度更保守影響訓(xùn)練吞吐。正確思路是先用默認方式跑通模型記錄峰值顯存然后再決定限制值。像那種上來就把per_process_gpu_memory_fraction設(shè)置成 0.3 的做法純粹是給自己埋雷。顯存管理的第一原則永遠是先保證任務(wù)能跑下來再談怎么分配。