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