如何通过Python代码限制TensorFlow使用的显存比例？-Python教程-PHP中文网

tensorflow 2.x 应在 import tensorflow 后、任何模型/计算图创建前，对每张 gpu 调用 tf.config.experimental.set_memory_growth(gpu, true)，实现按需分配显存，避免预占全部显存；该设置不可与 set_memory_limit 共存，且一旦 cuda 上下文初始化即失效。

如何通过python代码限制tensorflow使用的显存比例？

TensorFlow 2.x 怎么用 `set_memory_growth` 避免显存占满？

默认情况下，TensorFlow 2.x 会预分配几乎全部 GPU 显存，哪怕你只跑一个很小的模型，也会看到 nvidia-smi 显示显存已用 90%+。这不是真的被占用，而是预留（reserved），但会导致其他进程或同卡多实例无法启动。

最稳妥的解法是启用内存增长模式：让 TensorFlow 按需申请显存，而不是一次性占满。

必须在导入 tensorflow 后、创建任何模型或会话前调用
对每个可见 GPU 都要单独设置，不能只设 gpu[0] 就以为全局生效
代码示例（放在脚本最开头）：

import tensorflow as tf
gpus = tf.config.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)  # 初始化后就不能再修改，比如 Keras 已触发构建图就会报错

注意：一旦触发了 tf.function 编译、Keras 模型 .fit() 或 tf.data pipeline 构建，再调这个函数就无效，会抛 RuntimeError: Physical devices cannot be modified after being initialized。

想精确控制显存上限？用 `set_memory_limit` 而不是比例

TensorFlow 不提供“限制 40% 显存”这种比例接口，只支持按字节设硬上限（set_memory_limit）。你需要自己算出目标值，比如 16GB 卡限制为 8GB：

先查显卡总显存：nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits
换算成字节（如 8 * 1024 * 1024 * 1024），传给 set_memory_limit
同样必须在初始化前调用，且和 set_memory_growth 互斥——二者只能选其一

示例：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_limit(gpus[0], 8 * 1024 ** 3)  # 8GB
    except RuntimeError as e:
        print(e)

如果设得太小（比如低于模型权重+梯度所需空间），训练时会直接报 ResourceExhaustedError: OOM when allocating tensor，而不是自动降级——它不会“智能压缩”，只会硬失败。

Python 3.14.2

Python 3.14.2是Python编程语言在2025年12月5日发布的稳定版本，属于3.14系列的第二个维护更新。该版本包含了18项修复，重点解决了多进程、数据类及正则表达式等模块的回归问题，并修复了CVE-2025-12084等安全漏洞。此版本标志着自由线程模式（移除GIL）正式获得官方支持，是Python发展的重要里程碑。

下载

为什么 `os.environ['TF_FORCE_GPU_ALLOW_GROWTH']='true'` 有时不生效？

这个环境变量是旧版（TF 1.x）遗留方式，在 TF 2.4+ 中已被标记为 deprecated，且优先级低于代码中显式调用的 set_memory_growth。

如果代码里调了 set_memory_growth，环境变量会被忽略
如果没写代码，但设置了该变量，它只对第一个 GPU 生效，且无法覆盖已加载的 CUDA 上下文
常见误操作：在 import tensorflow 之后才设环境变量 → 完全无效

真正可靠的顺序只有：启动 Python 前设环境变量，或者在 import tensorflow 后立刻调 API —— 二者不要混用。

多卡场景下，每张卡的显存策略能不同吗？

可以，但得手动遍历 list_physical_devices('GPU') 并分别设置。比如只对第 0 卡限 6GB，第 1 卡开 growth：

gpus = tf.config.list_physical_devices('GPU')
if len(gpus) >= 2:
    tf.config.experimental.set_memory_limit(gpus[0], 6 * 1024**3)
    tf.config.experimental.set_memory_growth(gpus[1], True)

注意：Keras 多卡分布式策略（如 tf.distribute.MirroredStrategy）会自动管理显存，此时手动设限可能被绕过；若需精细控制，建议改用 tf.distribute.MultiWorkerMirroredStrategy 或单卡调试后再扩展。

显存控制的关键不在“设多少”，而在于“什么时候设”——晚于图构建一步，所有设置都变成废操作。