深度学习 | TensorFlow 2.x 和 1.x 限制显存超详细聊到 TensorFlow 跑模型显存这个话题几乎绕不开。我见过太多人在服务器上训练模型时一个脚本下去整块 GPU 直接被占满其他同事的进程活活被挤掉严重的时候连显卡的显示输出都跟着卡顿。还有一类场景是训完一个模型想接着做推理结果报错说显存不够明明模型很小可显存就是被前面的 session 占着不放。这些情况的根源都是 TensorFlow 默认会贪婪地把 GPU 上几乎所有的显存都申请走哪怕当前模型根本用不了那么多。所以今天这篇就把 TensorFlow 1.x 和 2.x 里限制显存这件事彻底讲透。先从它为什么默认全占说起再分别给出两代 API 的具体写法最后补充几个非常容易踩的坑。如果你正被显存不够、OOM 或者同事抱怨抢卡困扰这篇文章应该能直接解决问题。1. 先说原理为什么 TensorFlow 默认要把显存占满在写代码之前理解背后的机制很重要。TensorFlow 在检测到 GPU 后默认会调用一个叫“预先分配显存”的策略。具体来说框架启动时会向驱动查询当前 GPU 的总显存大小然后一次性申请其中绝大部分通常在 90% 以上作为进程的预分配池。后面训练中不管是张量计算、中间结果还是梯度缓存都从这个池子里拿。这个设计本身是为了性能。显存分配是个比较重的操作如果在每个 step 里反复申请和释放会因为内存碎片和驱动开销拖慢训练速度。一次性拿大头后面用的时候直接从池子划拨速度就快很多。所以官方给默认行为配了一个充分的理由稳定、高速、适合专用训练机。但问题在于这种“为性能优化”的默认策略在真实场景里很尴尬。开发调试阶段模型可能只在几 GB 显存里跑结果整块 24GB 的卡被占满别人没法用线上推理服务部署时模型需要的显存可能是动态波动的固定占大头既浪费资源又可能导致同一张卡上无法部署多个服务。另外像本地开发机还要兼顾显示输出显存被抢走后界面都可能卡顿。所以大部分情况下我们需要手动干预显存分配策略。TensorFlow 提供了两种思路一种是按需增长也就是用到多少申请多少用不到的时候不占另一种是设置硬上限无论模型需要多少最多给你指定的量。两种策略在 1.x 和 2.x 里的写法不同下面分别展开。2. TensorFlow 1.x 限制显存两种策略的完整写法TF 1.x 的 API 虽然官方早已停止维护但存量项目仍然非常多尤其是很多老论文复现、旧代码库和工业部署系统至今跑在 1.14 或者 1.15 上。这个版本的显存配置核心是一个叫ConfigProto的配置对象配合tf.Session传入。2.1 按需增长策略Grow by Request按需增长就是让 TensorFlow 不要一次性把显存全拿完而是用到多少拿多少。代码是这样写的import tensorflow as tf config tf.ConfigProto() config.gpu_options.allow_growth True session tf.Session(configconfig)就这么三行。设置allow_growth True后TensorFlow 在初始化时只申请一小块显存之后每检测到新的显存需求就在当前已申请的基础上继续追加直到达到显卡物理上限。这样做的好处很明显单进程只占实际使用的显存别的任务可以共享这块卡。但这里有个真实的坑。这个策略有个“只涨不跌”的特性。它不会在计算结束、显存空闲时自动把已经申请的内存还给驱动。比如你训练一个需要 6GB 显存的模型用按需增长跑了一段时间那么即使后来切换到只占 1GB 的小模型在同一进程内进程依然维持着 6GB 的已申请显存。原因在于 CUDA 缓存分配器通常不把内存归还给驱动除非进程终止。所以在做多任务切换的场景按需增长的效果没有想象中那么完美。2.2 固定上限策略Limit by Hard Cap如果你的需求很明确——只允许 TensorFlow 用某个数值的显存多一分都不给——那就用硬上限。做法是通过per_process_gpu_memory_fraction指定一个比例。例如显卡总显存 12GB你想限制到 4GB也就是三分之一多一点那么写import tensorflow as tf config tf.ConfigProto() config.gpu_options.per_process_gpu_memory_fraction 0.35 session tf.Session(configconfig)这里 0.35 表示总显存的比例不是具体字节数。这个比例是按整个进程来算的模型需要超过这个值就直接报错不给你动态追加的机会。当然也可以用gpu_options.allow_growth同时配合显存比例两者会结合作用但一般情况下只会选用其一。实际操作中我更推荐按需增长用于开发调试硬上限用于生产部署。前者保证你不会因为临时跑个小模型把整卡吃死后者则强制程序只能在预算内运行避免某次数据波动直接拖垮整机的其他服务。2.3 TF 1.x 中的多 GPU 限制写法很多时候你的机器不止一张卡这时候要分别限制每张卡。一种常用的写法是在CUDA_VISIBLE_DEVICES环境变量里先指定物理可见的设备序号再在配置里对每个设备设置允许增长的标记import os os.environ[CUDA_VISIBLE_DEVICES] 0,1 config tf.ConfigProto() config.allow_soft_placement True config.gpu_options.allow_growth True session tf.Session(configconfig)需要理解的是CUDA_VISIBLE_DEVICES并不是限制显存而是让 TensorFlow 只看到指定的卡。如果你希望物理第 0 号卡和第 1 号卡分别限制不同的大小那么需要分别对每个设备创建 session 或者手动指定分配的 scope这属于更进阶的用法下文会在进阶部分展开。注意TF 1.x 中tf.Session()如果不传 config它不会读取任何你预先设置的全局配置。很多新手在代码里写了config但忘了传进Session结果发现没有生效。这一步非常容易漏必须在构造 Session 时显式传入。3. TensorFlow 2.x 限制显存新版 API 的正确姿势TF 2.x 结构上做了大改原来的Session和ConfigProto被废弃了取而代之的是统一的tf.configAPI。限制显存的方式也从“配置会话”变成了“配置逻辑设备”。这个转变容易让人困惑但其实非常好理解。3.1 按需增长策略Growth by DemandTF 2.x 的按需增长一句话就能搞定import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_growth(physical_devices[0], True)这里tf.config.list_physical_devices(GPU)会返回当前机器上所有可见的物理 GPU 设备列表。拿到设备对象后调用set_memory_growth并传入True这台逻辑设备就会在运行时逐步申请显存。注意这里传入的是列表中的某个元素如果机器有多张卡你需要对列表中的每个设备都调用一次才能让所有卡都启用按需增长。这个策略和 1.x 时代一样开始只占用很小的显存模型实际使用多少TensorFlow 就在运行过程中逐步追加申请。比如训练一个 batch size 为 32 的模型显存占用可能从几百 MB 一路爬到数 GB期间不会一次性请求太多别人在这张卡上跑轻量任务不受影响。3.2 显存上限策略Cap on Memory如果你希望限制的是指定数值而不是按需增长可以用set_memory_limit指定字节数。比如给一张卡设上限为 4GB要先把字节数算清楚import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 4096 * 1024 * 1024)4096 乘以 1024 再乘以 1024就是 4GB 对应的字节数。这里有个容易出错的地方不同卡的实际可用显存可能略有差别比如部分显存被其他服务占用如果你设置的上限接近或超过实际可用的物理显存TensorFlow 可能会在初始化时报错。所以设置之前最好先通过nvidia-smi查看实际显存容量留出一些余量。更灵活的做法是写代码的时候先查询每张卡的显存大小再按百分比来设import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) for dev in physical_devices: tf.config.experimental.set_memory_limit(dev, 4096 * 1024 * 1024)这个循环会对所有卡统一设置 4GB 上限。如果你机器里有 3 张卡这样写最省事。3.3 在 2.x 中兼容 1.x 配置代码的坑TF 2.x 虽然推荐用新的 API但部分老代码的配置方式依然保留下了一个“兼容层”。具体来说tf.compat.v1.ConfigProto在 2.x 中仍然可用不过它只能配合tf.compat.v1.Session使用。很多人在 2.x 环境里复制了 1.x 的代码觉得只要包一层tf.compat.v1就能运行结果发现 new API 并不会读这些配置。所以建议是新项目直接使用tf.config系列接口旧项目迁移时别混用两种机制否则调了半天发现没效果其实是因为配置根本没有生效。另外还有一个小细节。tf.config的这些显存设置必须在任何设备或 TensorFlow 计算图被创建之前调用。如果你已经创建了某个模型或执行了某个 op再调用set_memory_growthTensorFlow 会直接抛出一个运行时错误。这是因为计算图初始化时已经确定了显存分配策略不允许中途变更。所以规则很简单把显存配置放在代码文件的最开头import 之后紧接着执行。4. 进阶技巧逻辑设备切分、多卡动态限制和容器场景基础的两类限制写完了但在实际项目中还会遇到一些更复杂的需求。比如一张大显存卡上要同时跑两个不同模型或者希望给多卡训练的任务设置某种调度策略再或者是在 Docker 容器里跑 TensorFlow宿主机的显存和容器内可见的不一样限制的策略也完全不同。这些都属于经验里比较值钱的部分单独讲一节。4.1 一张卡切成多个逻辑设备如果你的模型显存需求不大但卡很大例如有一块 24GB 的卡模型只要 6GB可以在一张卡上虚拟出 4 个逻辑设备分别跑 4 个不同的任务。TF 2.x 里用tf.config.physical_device_to_logical_device配合logical_device_count可以实现import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 12 * 1024 * 1024 * 1024) tf.config.experimental.set_logical_device_configuration( physical_devices[0], [tf.config.LogicalDeviceConfiguration(memory_limit6 * 1024 * 1024 * 1024), tf.config.LogicalDeviceConfiguration(memory_limit6 * 1024 * 1024 * 1024)] )手工切分逻辑设备时要注意所有逻辑设备的内存之和不能超过物理设备的上限并且设置必须放在第一次使用 GPU 之前。这种用法在生产环境的多服务并行部署中非常有用能最大化利用昂贵的 GPU 资源。不过它也有代价——逻辑设备之间无法高效共享显存。如果你同时跑的任务有显存波峰和波谷互相错开的场景逻辑切分反而不如按需增长来得自然。4.2 动态限制根据进程参数调整上限部署推理服务时经常需要根据模型的大小动态决定显存上限。比如算法同学训练了三个模型分别要 1GB、3GB、5GB 显存希望启动推理进程时通过参数指定显存大小。早期的做法是在代码里写死数字但更合理的是从环境变量读取import os import tensorflow as tf mem_limit os.getenv(PER_PROCESS_GPU_MEMORY, 4096) limit_bytes int(mem_limit) * 1024 * 1024 physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], limit_bytes)配合部署脚本PER_PROCESS_GPU_MEMORY2048 python server.py这样启动就能让同一个推理服务根据环境变量灵活调整显存。这种方式在 K8s 或者 Docker 环境里部署时非常顺手不用改代码就能适配不同规格的显存卡。4.3 容器场景下的显存限制注意点在 Docker 或 Kubernetes 中跑训练任务时容器通过--gpus参数映射宿主机的 GPU 设备。这里面隐藏着一个容易踩的坑容器内nvidia-smi显示的显存大小依然是整卡的物理显存而不是你通过--memory参数限制的内存大小。因为显存不同于系统内存Docker 本身默认不限制 GPU 显存--memory只约束 CPU 内存。所以容器里如果不在 TensorFlow 代码中设置显存上限一次 OOM 可能直接波及宿主机的整张卡。我见过生产环境的事故容器里跑了一个训练脚本没设置任何显存策略结果直接把宿主机整块 32GB 卡吃满其他容器内的推理服务全部报 failed to allocate memory。解决方式有两种一种是在容器入口处统一注入TF_GPU_ALLOCATORcuda_malloc_async环境变量配合动态增长另一种是像上面那样在代码里设置显存上限并把上限写进容器启动脚本的环境变量配置。4.4 多卡训练时对单卡做差异化限制多卡训练、特别是使用tf.distribute.MirroredStrategy时每个进程往往对应一张卡。如果希望不同卡上的进程显存上限不一样需要通过CUDA_VISIBLE_DEVICES提前指定当前进程可见哪张卡再只对列表的第一个设备做设置。这样做的好处是代码统一不管进程被分配到哪张卡它总是限制当前可见卡的显存。import os import tensorflow as tf # 假设外部传入了当前进程应该使用的物理卡号 gpu_id os.getenv(GPU_ID, 0) os.environ[CUDA_VISIBLE_DEVICES] gpu_id physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 8192 * 1024 * 1024)注意这时list_physical_devices只会返回一个设备因为其他卡对当前进程是不可见的。这是一个非常实用的部署技巧适合数据并行训练和分布式参数服务器的场景。5. 常见问题与排查技巧实录写了这么多年 TensorFlow 代码这两代 API 在显存上的坑我基本都踩过一遍。这里挑几个最典型的、也是最容易让新人卡住半天的整理成表格附带排查思路和解决办法。问题现象根本原因解决方式设置了 allow_growth 但显存还是被全占配置没有传进 Session或者运行环境存在多个 session 混用检查tf.Session(configconfig)是否显式传入TF2 中确认没有混用 compat.v1 与新 API程序一开始就报 failed to allocate memory物理显存本身不足或设置的上限大于实际可用显存用nvidia-smi查看实际显存调低上限值检查是否其他进程占用了大量显存设置 set_memory_growth 时抛运行时错误调用时机太晚计算图已经创建把配置移动到所有模型构建和训练代码之前最好紧跟 import 之后单卡上跑多个小模型总是互相挤爆没有启用允许增长或没有切分逻辑设备改用按需增长或者使用LogicalDeviceConfiguration手动切分容器里明明限制了内存显存还是暴涨Docker 的 --memory 不限制显存在入口脚本或代码中设置set_memory_limit或用TF_GPU_ALLOCATORcuda_malloc_async旧代码从 1.x 迁移到 2.x配置没效果新的 API 不读取旧的ConfigProto改为tf.config.experimental.set_memory_growth或set_memory_limit额外补充两个不常被提到但很重要的细节。第一显存增长策略和系统内存分配器的交互。在深度学习训练中Python 端频繁创建和销毁张量对象会带来一定的内存碎片。如果你启用了allow_growth显存碎片化可能比一次性拿全部显存更明显因为每次追加的粒度不等于实际需要的粒度。长时间跑任务时偶尔可以在训练 loop 里监控一下显存碎片率。如果碎片率过高导致性能下降可以考虑换回固定上限模式或者干脆增大 batch size减少频繁小规模分配。第二TF_GPU_ALLOCATOR环境变量的作用。TF 2.7 之后支持把底层显存分配器换成 CUDA 异步分配器TF_GPU_ALLOCATORcuda_malloc_async python train.py这个环境变量在某些新架构显卡上能明显提高显存利用率特别是在动态 shape 和大量小张量场景下。但它也不是万能药旧代码在部分卡上可能触发兼容性问题。如果你在开发时发现显存占用曲线忽高忽低或者分配速度明显变慢可以试着切换这个环境变量做对比。6. 讲点个人体会和推荐组合文章写到这里显存限制的核心策略基本都覆盖了。最后分享一些我在实际项目中沉淀的搭配规则。对于实验室共用服务器我推荐默认开启allow_growth策略。共用机器上每个人跑的任务复杂度和显存需求都不一样按需增长能最大限度做到互不干扰。但要注意提示团队里的同事用完大模型后如果继续跑小模型尽量重启进程避免因为缓存分配器不主动归还显存导致别人觉得你在偷占资源。对于线上推理服务推荐用固定上限。因为服务实例需要稳定可预测的资源占用扩容和缩容都依赖这个数字。在容器编排里按模型大小配置PER_PROCESS_GPU_MEMORY环境变量再让代码读取这个变量设置set_memory_limit管理起来最省心。对于大模型训练跑多卡则不建议限制单卡显存。分布式训练对显存需求通常只增不减限制反而可能导致 OOM。你需要做的是保证每张卡上只有这一个任务动态增长和固定上限都适合但更建议直接不限制让框架自行管理。在实际使用中我还发现一个容易被忽略的点混合精度训练往往会降低显存占用因为多数张量变成了 float16。但如果你启用了allow_growth显存的“爬坡”过程会更频繁运行第一个 step 和运行第十个 step 的显存占用可能差出一大截新手看着监控曲线容易误以为内存泄漏。其实这是正常现象框架在逐步启用更多内核时分配了额外暂存空间。多跑几个 step等显存曲线稳定下来再判断状态。最后一个实际建议不管用什么策略代码里最好加一段显存状态打印。TF 2.x 里可以这样写import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) for dev in physical_devices: print(fDevice: {dev}) tf.config.experimental.set_memory_growth(dev, True)跑起来后先用nvidia-smi确认下实际显存占用是否符合预期然后再开始训练。显存这个东西配置好了省心配置错了浪费的不只是算力还有你的调试时间。希望这篇能帮你少走点弯路。 
长河君 · 2026/10/20 15:14:27 · 阅读 1.2k TensorFlow显存限制实战：1.x与2.x防OOM策略详解
深度学习 | TensorFlow 2.x 和 1.x 限制显存超详细聊到 TensorFlow 跑模型显存这个话题几乎绕不开。我见过太多人在服务器上训练模型时一个脚本下去整块 GPU 直接被占满其他同事的进程活活被挤掉严重的时候连显卡的显示输出都跟着卡顿。还有一类场景是训完一个模型想接着做推理结果报错说显存不够明明模型很小可显存就是被前面的 session 占着不放。这些情况的根源都是 TensorFlow 默认会贪婪地把 GPU 上几乎所有的显存都申请走哪怕当前模型根本用不了那么多。所以今天这篇就把 TensorFlow 1.x 和 2.x 里限制显存这件事彻底讲透。先从它为什么默认全占说起再分别给出两代 API 的具体写法最后补充几个非常容易踩的坑。如果你正被显存不够、OOM 或者同事抱怨抢卡困扰这篇文章应该能直接解决问题。1. 先说原理为什么 TensorFlow 默认要把显存占满在写代码之前理解背后的机制很重要。TensorFlow 在检测到 GPU 后默认会调用一个叫“预先分配显存”的策略。具体来说框架启动时会向驱动查询当前 GPU 的总显存大小然后一次性申请其中绝大部分通常在 90% 以上作为进程的预分配池。后面训练中不管是张量计算、中间结果还是梯度缓存都从这个池子里拿。这个设计本身是为了性能。显存分配是个比较重的操作如果在每个 step 里反复申请和释放会因为内存碎片和驱动开销拖慢训练速度。一次性拿大头后面用的时候直接从池子划拨速度就快很多。所以官方给默认行为配了一个充分的理由稳定、高速、适合专用训练机。但问题在于这种“为性能优化”的默认策略在真实场景里很尴尬。开发调试阶段模型可能只在几 GB 显存里跑结果整块 24GB 的卡被占满别人没法用线上推理服务部署时模型需要的显存可能是动态波动的固定占大头既浪费资源又可能导致同一张卡上无法部署多个服务。另外像本地开发机还要兼顾显示输出显存被抢走后界面都可能卡顿。所以大部分情况下我们需要手动干预显存分配策略。TensorFlow 提供了两种思路一种是按需增长也就是用到多少申请多少用不到的时候不占另一种是设置硬上限无论模型需要多少最多给你指定的量。两种策略在 1.x 和 2.x 里的写法不同下面分别展开。2. TensorFlow 1.x 限制显存两种策略的完整写法TF 1.x 的 API 虽然官方早已停止维护但存量项目仍然非常多尤其是很多老论文复现、旧代码库和工业部署系统至今跑在 1.14 或者 1.15 上。这个版本的显存配置核心是一个叫ConfigProto的配置对象配合tf.Session传入。2.1 按需增长策略Grow by Request按需增长就是让 TensorFlow 不要一次性把显存全拿完而是用到多少拿多少。代码是这样写的import tensorflow as tf config tf.ConfigProto() config.gpu_options.allow_growth True session tf.Session(configconfig)就这么三行。设置allow_growth True后TensorFlow 在初始化时只申请一小块显存之后每检测到新的显存需求就在当前已申请的基础上继续追加直到达到显卡物理上限。这样做的好处很明显单进程只占实际使用的显存别的任务可以共享这块卡。但这里有个真实的坑。这个策略有个“只涨不跌”的特性。它不会在计算结束、显存空闲时自动把已经申请的内存还给驱动。比如你训练一个需要 6GB 显存的模型用按需增长跑了一段时间那么即使后来切换到只占 1GB 的小模型在同一进程内进程依然维持着 6GB 的已申请显存。原因在于 CUDA 缓存分配器通常不把内存归还给驱动除非进程终止。所以在做多任务切换的场景按需增长的效果没有想象中那么完美。2.2 固定上限策略Limit by Hard Cap如果你的需求很明确——只允许 TensorFlow 用某个数值的显存多一分都不给——那就用硬上限。做法是通过per_process_gpu_memory_fraction指定一个比例。例如显卡总显存 12GB你想限制到 4GB也就是三分之一多一点那么写import tensorflow as tf config tf.ConfigProto() config.gpu_options.per_process_gpu_memory_fraction 0.35 session tf.Session(configconfig)这里 0.35 表示总显存的比例不是具体字节数。这个比例是按整个进程来算的模型需要超过这个值就直接报错不给你动态追加的机会。当然也可以用gpu_options.allow_growth同时配合显存比例两者会结合作用但一般情况下只会选用其一。实际操作中我更推荐按需增长用于开发调试硬上限用于生产部署。前者保证你不会因为临时跑个小模型把整卡吃死后者则强制程序只能在预算内运行避免某次数据波动直接拖垮整机的其他服务。2.3 TF 1.x 中的多 GPU 限制写法很多时候你的机器不止一张卡这时候要分别限制每张卡。一种常用的写法是在CUDA_VISIBLE_DEVICES环境变量里先指定物理可见的设备序号再在配置里对每个设备设置允许增长的标记import os os.environ[CUDA_VISIBLE_DEVICES] 0,1 config tf.ConfigProto() config.allow_soft_placement True config.gpu_options.allow_growth True session tf.Session(configconfig)需要理解的是CUDA_VISIBLE_DEVICES并不是限制显存而是让 TensorFlow 只看到指定的卡。如果你希望物理第 0 号卡和第 1 号卡分别限制不同的大小那么需要分别对每个设备创建 session 或者手动指定分配的 scope这属于更进阶的用法下文会在进阶部分展开。注意TF 1.x 中tf.Session()如果不传 config它不会读取任何你预先设置的全局配置。很多新手在代码里写了config但忘了传进Session结果发现没有生效。这一步非常容易漏必须在构造 Session 时显式传入。3. TensorFlow 2.x 限制显存新版 API 的正确姿势TF 2.x 结构上做了大改原来的Session和ConfigProto被废弃了取而代之的是统一的tf.configAPI。限制显存的方式也从“配置会话”变成了“配置逻辑设备”。这个转变容易让人困惑但其实非常好理解。3.1 按需增长策略Growth by DemandTF 2.x 的按需增长一句话就能搞定import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_growth(physical_devices[0], True)这里tf.config.list_physical_devices(GPU)会返回当前机器上所有可见的物理 GPU 设备列表。拿到设备对象后调用set_memory_growth并传入True这台逻辑设备就会在运行时逐步申请显存。注意这里传入的是列表中的某个元素如果机器有多张卡你需要对列表中的每个设备都调用一次才能让所有卡都启用按需增长。这个策略和 1.x 时代一样开始只占用很小的显存模型实际使用多少TensorFlow 就在运行过程中逐步追加申请。比如训练一个 batch size 为 32 的模型显存占用可能从几百 MB 一路爬到数 GB期间不会一次性请求太多别人在这张卡上跑轻量任务不受影响。3.2 显存上限策略Cap on Memory如果你希望限制的是指定数值而不是按需增长可以用set_memory_limit指定字节数。比如给一张卡设上限为 4GB要先把字节数算清楚import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 4096 * 1024 * 1024)4096 乘以 1024 再乘以 1024就是 4GB 对应的字节数。这里有个容易出错的地方不同卡的实际可用显存可能略有差别比如部分显存被其他服务占用如果你设置的上限接近或超过实际可用的物理显存TensorFlow 可能会在初始化时报错。所以设置之前最好先通过nvidia-smi查看实际显存容量留出一些余量。更灵活的做法是写代码的时候先查询每张卡的显存大小再按百分比来设import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) for dev in physical_devices: tf.config.experimental.set_memory_limit(dev, 4096 * 1024 * 1024)这个循环会对所有卡统一设置 4GB 上限。如果你机器里有 3 张卡这样写最省事。3.3 在 2.x 中兼容 1.x 配置代码的坑TF 2.x 虽然推荐用新的 API但部分老代码的配置方式依然保留下了一个“兼容层”。具体来说tf.compat.v1.ConfigProto在 2.x 中仍然可用不过它只能配合tf.compat.v1.Session使用。很多人在 2.x 环境里复制了 1.x 的代码觉得只要包一层tf.compat.v1就能运行结果发现 new API 并不会读这些配置。所以建议是新项目直接使用tf.config系列接口旧项目迁移时别混用两种机制否则调了半天发现没效果其实是因为配置根本没有生效。另外还有一个小细节。tf.config的这些显存设置必须在任何设备或 TensorFlow 计算图被创建之前调用。如果你已经创建了某个模型或执行了某个 op再调用set_memory_growthTensorFlow 会直接抛出一个运行时错误。这是因为计算图初始化时已经确定了显存分配策略不允许中途变更。所以规则很简单把显存配置放在代码文件的最开头import 之后紧接着执行。4. 进阶技巧逻辑设备切分、多卡动态限制和容器场景基础的两类限制写完了但在实际项目中还会遇到一些更复杂的需求。比如一张大显存卡上要同时跑两个不同模型或者希望给多卡训练的任务设置某种调度策略再或者是在 Docker 容器里跑 TensorFlow宿主机的显存和容器内可见的不一样限制的策略也完全不同。这些都属于经验里比较值钱的部分单独讲一节。4.1 一张卡切成多个逻辑设备如果你的模型显存需求不大但卡很大例如有一块 24GB 的卡模型只要 6GB可以在一张卡上虚拟出 4 个逻辑设备分别跑 4 个不同的任务。TF 2.x 里用tf.config.physical_device_to_logical_device配合logical_device_count可以实现import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 12 * 1024 * 1024 * 1024) tf.config.experimental.set_logical_device_configuration( physical_devices[0], [tf.config.LogicalDeviceConfiguration(memory_limit6 * 1024 * 1024 * 1024), tf.config.LogicalDeviceConfiguration(memory_limit6 * 1024 * 1024 * 1024)] )手工切分逻辑设备时要注意所有逻辑设备的内存之和不能超过物理设备的上限并且设置必须放在第一次使用 GPU 之前。这种用法在生产环境的多服务并行部署中非常有用能最大化利用昂贵的 GPU 资源。不过它也有代价——逻辑设备之间无法高效共享显存。如果你同时跑的任务有显存波峰和波谷互相错开的场景逻辑切分反而不如按需增长来得自然。4.2 动态限制根据进程参数调整上限部署推理服务时经常需要根据模型的大小动态决定显存上限。比如算法同学训练了三个模型分别要 1GB、3GB、5GB 显存希望启动推理进程时通过参数指定显存大小。早期的做法是在代码里写死数字但更合理的是从环境变量读取import os import tensorflow as tf mem_limit os.getenv(PER_PROCESS_GPU_MEMORY, 4096) limit_bytes int(mem_limit) * 1024 * 1024 physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], limit_bytes)配合部署脚本PER_PROCESS_GPU_MEMORY2048 python server.py这样启动就能让同一个推理服务根据环境变量灵活调整显存。这种方式在 K8s 或者 Docker 环境里部署时非常顺手不用改代码就能适配不同规格的显存卡。4.3 容器场景下的显存限制注意点在 Docker 或 Kubernetes 中跑训练任务时容器通过--gpus参数映射宿主机的 GPU 设备。这里面隐藏着一个容易踩的坑容器内nvidia-smi显示的显存大小依然是整卡的物理显存而不是你通过--memory参数限制的内存大小。因为显存不同于系统内存Docker 本身默认不限制 GPU 显存--memory只约束 CPU 内存。所以容器里如果不在 TensorFlow 代码中设置显存上限一次 OOM 可能直接波及宿主机的整张卡。我见过生产环境的事故容器里跑了一个训练脚本没设置任何显存策略结果直接把宿主机整块 32GB 卡吃满其他容器内的推理服务全部报 failed to allocate memory。解决方式有两种一种是在容器入口处统一注入TF_GPU_ALLOCATORcuda_malloc_async环境变量配合动态增长另一种是像上面那样在代码里设置显存上限并把上限写进容器启动脚本的环境变量配置。4.4 多卡训练时对单卡做差异化限制多卡训练、特别是使用tf.distribute.MirroredStrategy时每个进程往往对应一张卡。如果希望不同卡上的进程显存上限不一样需要通过CUDA_VISIBLE_DEVICES提前指定当前进程可见哪张卡再只对列表的第一个设备做设置。这样做的好处是代码统一不管进程被分配到哪张卡它总是限制当前可见卡的显存。import os import tensorflow as tf # 假设外部传入了当前进程应该使用的物理卡号 gpu_id os.getenv(GPU_ID, 0) os.environ[CUDA_VISIBLE_DEVICES] gpu_id physical_devices tf.config.list_physical_devices(GPU) if physical_devices: tf.config.experimental.set_memory_limit(physical_devices[0], 8192 * 1024 * 1024)注意这时list_physical_devices只会返回一个设备因为其他卡对当前进程是不可见的。这是一个非常实用的部署技巧适合数据并行训练和分布式参数服务器的场景。5. 常见问题与排查技巧实录写了这么多年 TensorFlow 代码这两代 API 在显存上的坑我基本都踩过一遍。这里挑几个最典型的、也是最容易让新人卡住半天的整理成表格附带排查思路和解决办法。问题现象根本原因解决方式设置了 allow_growth 但显存还是被全占配置没有传进 Session或者运行环境存在多个 session 混用检查tf.Session(configconfig)是否显式传入TF2 中确认没有混用 compat.v1 与新 API程序一开始就报 failed to allocate memory物理显存本身不足或设置的上限大于实际可用显存用nvidia-smi查看实际显存调低上限值检查是否其他进程占用了大量显存设置 set_memory_growth 时抛运行时错误调用时机太晚计算图已经创建把配置移动到所有模型构建和训练代码之前最好紧跟 import 之后单卡上跑多个小模型总是互相挤爆没有启用允许增长或没有切分逻辑设备改用按需增长或者使用LogicalDeviceConfiguration手动切分容器里明明限制了内存显存还是暴涨Docker 的 --memory 不限制显存在入口脚本或代码中设置set_memory_limit或用TF_GPU_ALLOCATORcuda_malloc_async旧代码从 1.x 迁移到 2.x配置没效果新的 API 不读取旧的ConfigProto改为tf.config.experimental.set_memory_growth或set_memory_limit额外补充两个不常被提到但很重要的细节。第一显存增长策略和系统内存分配器的交互。在深度学习训练中Python 端频繁创建和销毁张量对象会带来一定的内存碎片。如果你启用了allow_growth显存碎片化可能比一次性拿全部显存更明显因为每次追加的粒度不等于实际需要的粒度。长时间跑任务时偶尔可以在训练 loop 里监控一下显存碎片率。如果碎片率过高导致性能下降可以考虑换回固定上限模式或者干脆增大 batch size减少频繁小规模分配。第二TF_GPU_ALLOCATOR环境变量的作用。TF 2.7 之后支持把底层显存分配器换成 CUDA 异步分配器TF_GPU_ALLOCATORcuda_malloc_async python train.py这个环境变量在某些新架构显卡上能明显提高显存利用率特别是在动态 shape 和大量小张量场景下。但它也不是万能药旧代码在部分卡上可能触发兼容性问题。如果你在开发时发现显存占用曲线忽高忽低或者分配速度明显变慢可以试着切换这个环境变量做对比。6. 讲点个人体会和推荐组合文章写到这里显存限制的核心策略基本都覆盖了。最后分享一些我在实际项目中沉淀的搭配规则。对于实验室共用服务器我推荐默认开启allow_growth策略。共用机器上每个人跑的任务复杂度和显存需求都不一样按需增长能最大限度做到互不干扰。但要注意提示团队里的同事用完大模型后如果继续跑小模型尽量重启进程避免因为缓存分配器不主动归还显存导致别人觉得你在偷占资源。对于线上推理服务推荐用固定上限。因为服务实例需要稳定可预测的资源占用扩容和缩容都依赖这个数字。在容器编排里按模型大小配置PER_PROCESS_GPU_MEMORY环境变量再让代码读取这个变量设置set_memory_limit管理起来最省心。对于大模型训练跑多卡则不建议限制单卡显存。分布式训练对显存需求通常只增不减限制反而可能导致 OOM。你需要做的是保证每张卡上只有这一个任务动态增长和固定上限都适合但更建议直接不限制让框架自行管理。在实际使用中我还发现一个容易被忽略的点混合精度训练往往会降低显存占用因为多数张量变成了 float16。但如果你启用了allow_growth显存的“爬坡”过程会更频繁运行第一个 step 和运行第十个 step 的显存占用可能差出一大截新手看着监控曲线容易误以为内存泄漏。其实这是正常现象框架在逐步启用更多内核时分配了额外暂存空间。多跑几个 step等显存曲线稳定下来再判断状态。最后一个实际建议不管用什么策略代码里最好加一段显存状态打印。TF 2.x 里可以这样写import tensorflow as tf physical_devices tf.config.list_physical_devices(GPU) for dev in physical_devices: print(fDevice: {dev}) tf.config.experimental.set_memory_growth(dev, True)跑起来后先用nvidia-smi确认下实际显存占用是否符合预期然后再开始训练。显存这个东西配置好了省心配置错了浪费的不只是算力还有你的调试时间。希望这篇能帮你少走点弯路。