做深度学习的同学应该都有过这种经历启动第一个 TensorFlow 训练脚本后打开nvidia-smi一看明明模型不大显存却直接占到了接近满载。这不是模型吃了多少的问题而是 TensorFlow 在默认情况下会对当前可见的 GPU 做一整块预分配把整张卡的显存当作自己家的仓库。这个行为在 1.x 和 2.x 里都存在只是处理机制和配置入口很不一样。这篇文章就把 TensorFlow 2.x 和 1.x 限制显存的主流方法、底层参数含义、踩坑点一次讲透。看完之后你会明白为什么限制显存不等于限制模型大小为什么allow_growth有时不起作用以及多进程共享一块 GPU 时到底该怎么分配。1. 默认行为一张卡被单进程整套带走1.1 TensorFlow 为什么会把显存全部占满TensorFlow 的显存分配主要交给内部的 BFCBest-Fit with Coalescing分配器来做。BFC 分配器的工作方式和操作系统里的内存管理器很像它先把一大块显存接管下来然后在内部切成不同大小的块供算子按需申请。但关键在这一步如果不做任何配置TensorFlow 在初始化 GPU 设备时会直接把当前进程可见的整张卡显存全部申请到自己的池子里。对驱动层来说这张卡已经名义上被这个进程占完了哪怕你现在只跑了一个几 MB 的小张量nvidia-smi上显示的占用率也会是 90% 以上。很多刚入门的朋友会误以为这是模型太吃显存于是开始疯狂降低 batch size。其实模型本身占用的显存可能只有 12GB剩下的一大半都被预分配池子吞了。你限制显存限制的并不是算子使用的显存而是 BFC 分配器从显卡驱动那里申请的上限。1.2 真正需要限制显存的四个场景第一个场景是多人共用 GPU。实验室或公司里一张卡经常有几个人在跑任务如果每个人启动 TensorFlow 时都默认占满全卡后来的人只能干等。合理做法是每个人各自设一个上限比如一个人分 4GB另一个人分 4GB。第二个场景是单卡多任务。你可能要同时跑一个训练任务和一个推理任务或者一个脚本里加载多个模型。只要每个任务用到的显存都收敛就必须给它们分别画好边界否则第二个模型初始化时就会把第一个模型挤掉。第三个场景是 Jupyter Notebook 这类交互式环境。你反复执行代码单元格环境会保留历史计算图上的张量显存经常只涨不跌。加上 TensorFlow 默认把整卡占住写课程作业都会出问题。所以交互式环境里几乎必须限制显存。第四个场景是模型部署。线上服务进程需要长时间运行显存只能限制在一个稳定水位。如果让推理服务按需无限增长运行三五天后显存就可能被缓存顶到上限导致请求崩溃。1.3 一个容易混的概念进程显存不等于模型显存你限制显存之后nvidia-smi里看到的进程占用数字并不等于模型参数和中间特征图的大小。进程要正常运行还需要维持 CUDA context、cuDNN workspace、以及部分库的常驻内存。这部分开销通常有一两百 MB也就是你设置 2GB 上限时nvidia-smi里看到的进程占用可能会在 2.1GB 到 2.3GB 左右。所以我在实际配置时都会先按模型真实需求 上下文开销来算不会直接卡到模型需求值上。这一步虽然简单但很多同学第一次写限制配置时都会忽略导致明明给了足够空间却总是差一点。2. TensorFlow 2.x 限制显存的标准操作2.1 增长式分配set_memory_growthTensorFlow 2.x 官方的推荐入口是tf.config.experimental。最常用的方法是让显存按需增长也就是模型实际申请多少TensorFlow 就占多少不要一次性整卡占满。对应的代码是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)这里的关键点在于set_memory_growth必须在任何算子执行、任何模型构建之前调用。只要 TensorFlow 已经初始化了 GPU 上下文再设置就会触发RuntimeError因为分配策略已经定死了。另外需要注意这段代码只是开启用多少申请多少的模式并没有硬性规定上限。如果你的模型恰好需要 11GB而卡上只有 11GB那它还是会占满增长模式解决的是不用就别占着的问题不是最多只能用多少的问题。2.2 硬性上限用虚拟设备配置指定内存大小如果你希望一个进程严格限制在固定大小比如最多只能用 4GB剩下的 2GB 要留给另一个任务那就需要用到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。我的个人建议是训练脚本用set_memory_growth更省事因为只要显存不冲突它不会影响性能推理服务或需要稳定复现的脚本用memory_limit更靠谱因为固定水位好监控、好告警、好做容量规划。2.3 环境变量方式TF_FORCE_GPU_ALLOW_GROWTH除了在代码里设置TensorFlow 2.x 还支持通过环境变量TF_FORCE_GPU_ALLOW_GROWTH强制开启增长式分配。TF_FORCE_GPU_ALLOW_GROWTHtrue python train.py这种方式的优点是能不动训练代码直接通过启动命令控制进程的显存行为。在跑同学给的脚本、或者别人封装的模型仓库时我会先看对方代码里有没有写死显存配置如果没写就优先用这个环境变量临时接管。不过环境变量也有一些局限性。部分版本对它的支持不够完整或者会被代码里的设置覆盖。所以我只是把它当作应急手段代码里能设置的话还是老老实实在代码里设置。2.4 增长和硬上限到底怎么选set_memory_growth对应动态扩张memory_limit对应静态封顶。选择时主要看运行环境。如果你是在自己的电脑上训练显存怎么折腾都不会影响别人那直接用增长式最舒服既能充分利用显存又不会开机就占满。如果你是跑共享服务器且同一张卡上还有其他人的任务请务必使用memory_limit固定上限。否则就算设置了增长式只要你的模型稍微大一点它还是会迅速把卡占满和默认行为几乎没差别。如果是部署在线服务风控和容量规划很关键建议直接用memory_limit并把上限设置在峰值 20%的状态。比如某个模型峰值要 1.5GB我通常会给 1.8GB 左右留一点上下文开销和临时算子缓存。3. TensorFlow 1.x 的显存限制从 GPUOptions 到 Session 配置3.1 per_process_gpu_memory_fraction按比例设置显存上限TensorFlow 1.x 时代还没有tf.config.experimental这套新 API显存限制基本集中在GPUOptions和ConfigProto上。最常见的是通过比例限制比如把当前进程可用的显存限制在总显存的 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% 的显存。随着算子执行它会在这一份池子里继续分配。合理的地方在于参数直观比如一张 10GB 的卡0.4 大约就是 4GB。但注意这个比例是基于当前进程可见的显存总量来算的如果你先用CUDA_VISIBLE_DEVICES屏蔽了部分卡那比例对应的是可见卡的总量。3.2 allow_growth按需申请但小心碎片化1.x 里还有另一个经典参数allow_growth。设置它之后TensorFlow 不再一开始就把整张卡包下来而是根据每次算子的需求一点点向驱动申请显存。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)很像但老版本实现上有更多问题。最典型的是碎片化BFC 分配器按需申请内存时会因为各个张量的大小不一在显存池里留下很多不连续的小块。运行一段时间后可能整体显存还有几百 MB 空闲但找不到一块足够大的连续空间来放新张量结果照样 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)这是那个时代很常见的组合写法。现在回看backend.set_session这套机制早就被新 API 取代了但如果你还在维护老代码这仍然是必须认识的老朋友。3.3 在 2.x 中复用 1.x 的配置姿势有些老项目短时间内没法完全迁移到 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 的代码习惯。之后模型的构建和训练方式基本还是原来的写法。不过说实话如果项目还能改我不太建议长期停在 compat.v1 的兼容支线上。兼容层只是过渡方案新写的代码直接用 2.x 原生的tf.config.experimental更干净。老代码需要快速跑通的话用compat.v1能省不少事但不要把它当成长期架构。4. 实测不同限制参数下的显存占用和 OOM 边界4.1 用 nvidia-smi 观察真实占用在讨论各种配置之前先确认观察工具。nvidia-smi最常用的实时查看命令是watch -n 1 nvidia-smi如果想拿到格式化的占用数据方便记录到日志里可以用nvidia-smi --query-gpuindex,memory.used,memory.total,utilization.gpu --formatcsv -l 5-l 5表示每 5 秒刷新一次。实际调优时我更推荐看一下进程级显存而不是只看总量nvidia-smi --query-compute-appspid,used_memory,process_name --formatcsv这样能直接看到哪个 PID 占了多少不会被卡上其他进程干扰。4.2 一张 6GB 卡上的典型对比我拿一个常见的分割模型做参考输入大约512x512、batch size 为 8模型峰值需求大约在 3.4GB 左右。使用 TensorFlow 2.x分别用不同方式设置显存典型表现如下设置方式显存配置进程实际占用约训练速度表现不设置无占满整卡正常增长式set_memory_growth(True)约 3.5GB正常虚拟设备memory_limit4096约 4.2GB正常虚拟设备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 的缓冲训练速度基本不会受什么影响。真正的风险出现在限制值低于实际需求时。4.3 OOM 的具体表现不是等待而是立刻失败有相当一部分人以为显存不够时 TensorFlow 会像 CPU 内存那样先等一会儿或者和其他进程商量一下。实际完全不是这样。当 BFC 分配器向驱动申请一块连续显存失败时进程会立刻抛出ResourceExhaustedError并中断训练。报错信息大致长这样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还要关注特征图尺寸、通道数、以及可能临时创建的中间张量。4.4 为什么限制过小会引发连锁性的变慢如果把显存卡在模型峰值附近例如需求 3.4GB限制给 3.6GB你可能会发现训练速度没崩但比正常情况下慢了一截。原因主要有两类。第一类是显存不够时TensorFlow 会频繁进行块合并和碎片整理。BFC 里的块被拆开再合并本身就是额外开销每次算子执行都可能触发。第二类是部分算子发现显存紧张会退回到性能更差的实现或发起 host-device 拷贝。问题不会立刻以 OOM 形式出现而是让训练一步慢半拍非常隐蔽。所以在设计限制值时我一般留出 20% 的余量。比如模型峰值 3.4GB就限制在 4GB 或 4.5GB。这样既能保证稳定又不会给同卡其他任务带来太多负担。5. 多进程共享 GPU 时怎么规划显存才合理5.1 CUDA_VISIBLE_DEVICES 先做物理隔离多个进程分配同一张卡之前先想清楚每个进程到底应该看到哪张卡。CUDA_VISIBLE_DEVICES是在驱动层做可见性过滤的环境变量设置之后 TensorFlow 只能看到列表中的卡。import os os.environ[CUDA_VISIBLE_DEVICES] 1,2 import tensorflow as tf gpus tf.config.experimental.list_physical_devices(GPU) print(gpus)注意这个环境变量必须在import tensorflow之前设置。因为 TensorFlow 导入时会读取这个变量来确定设备可见性晚了就来不及了。设置1,2之后物理上的 1 号卡、2 号卡在进程内会变成逻辑0和1。也就是说list_physical_devices(GPU)返回索引是从 0 开始的逻辑索引不是物理索引。很多人踩过这个坑以为自己设置的是gpu:1实际是物理 0 号卡结果把别人的任务挤爆了。5.2 训练进程和推理进程的参数建议如果一张卡要同时跑一个训练任务和一个推理任务我比较推荐的安排是训练任务用set_memory_growth推理任务用memory_limit固定上限。原因在于训练过程的显存需求会随 batch size、输入尺寸和模型结构波动。用增长式至少能保证模型在峰值时不会因为提前锁死而 OOM。而推理任务通常输入波动不大显存需求相对稳定更适合用固定上限来保证服务质量。假如同一个模型推理峰值是 1GB我习惯给 1.3GB 或 1.5GB。这样即使输入偶尔变大一点也不至于挂掉同时给训练任务留出足够的剩余空间。5.3 多卡训练时逐卡设置上限用tf.distribute.MirroredStrategy做单机多卡训练时显存配置需要在策略创建之前完成。官方推荐是在代码最前面遍历所有物理 GPU并对每张卡分别设置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()如果每张卡型号一致直接循环设置即可。如果型号不一致比如一张 8GB、一张 16GB就要分别获取每张卡并逐个给不同的memory_limit否则策略初始化时会因为显存不一致出各种莫名其妙的报错。6. 避坑清单限制无效、显存只涨不降、速度异常6.1 明明设置了限制为什么还是占满这类问题我见过太多次几乎每次都是配置时机不对。常见情况有三种。第一种代码里先构建了模型或执行了算子再去调用set_memory_growth。此时 GPU 上下文已初始化TensorFlow 大概率会抛RuntimeError但也有老版本不报错、直接忽略配置的情况。所以配置要放在所有 TensorFlow 操作之前。第二种模块导入阶段就触发了 GPU 初始化。有些库在import时并不构建模型但会预加载 cuDNN 或创建上下文。如果你在导入之后才配置实际上已经晚了。第三种使用多进程库时子进程里的 TensorFlow 不会继承主进程的显存配置。如果主进程设置好了限制但训练是在spawn出来的子进程里跑的子进程仍会走默认的整卡占满逻辑。多进程代码需要考虑在子进程入口处重新执行一次显存配置。6.2 训练中显存只涨不降是泄漏吗很多人跑长时间训练发现nvidia-smi里进程占用从 3GB 慢慢涨到 4GB、5GB于是怀疑是显存泄漏。大多数情况下这不是真正意义上的泄漏而是 BFC 分配器的缓存机制在起作用。TensorFlow 在运行时会把用过的显存块留在池子里方便后续算子复用。这个池子的总占用会随着不同模块的负载浮动但不会在算子释放后立刻归还给驱动。真正判断泄漏的方法是看同一段循环里持续构造新张量、反复执行模型进程占用是否单调递增且永不回头。如果只是前期从低到高爬升之后就稳定在一个水位上那只是缓存池扩容不是泄漏。6.3 限制显存后训练莫名其妙变慢如果确认显存限制没有触发 OOM但训练速度明显下降优先检查你是不是把限制值压得太紧了。当 BFC 剩余空间很小每个算子申请显存时都面临刚刚好够、但不够整块的情况分配器就要频繁做块切分和合并。这个过程非常消耗时间。更麻烦的是cuDNN 在自动选择卷积实现时会搜索多个候选算法有些算法要求额外工作区空间。显存限制太紧会让 cuDNN 选不到最优算法只能退而求其次整个模型运行速度都会下降。正确的做法是给显存留足缓冲。不要追求刚好够用显存本质上也是性能资源。6.4 一条额外的注意点限制够用时不要过度限制最后说一句我自己在实践中的体会显存限制的价值是让多任务能共存不是为了省钱。如果整张卡只有你的一个任务就算模型只需 2GB也没必要强制限制到 4GB 以内。强行限制反而可能让 cuDNN 和 TensorFlow 的调度更保守影响训练吞吐。正确思路是先用默认方式跑通模型记录峰值显存然后再决定限制值。像那种上来就把per_process_gpu_memory_fraction设置成 0.3 的做法纯粹是给自己埋雷。显存管理的第一原则永远是先保证任务能跑下来再谈怎么分配。