解决tensorflow模型加载时把GPU显存占满的问题(亲测效果显著)

第一次用tensorflow模型进行推理。

初始状态下，显存的占用情况

模型加载后，电脑显存直接快拉满了，但是模型参数量并不大，这就很奇怪,究竟这是一股什么神秘的力量，竟要试图侵占我可怜的显卡宝贝

原因分析:安装了tensorflow-gpu后，运行程序默认是把GPU的内存全部占满的，但是正常人都不希望这样，都希望能自己掌握自己的资源，而不是任人摆布，可以这样操作。

解决代码：

import tensorflow as tf
import os

# 指定使用的GPU
os.environ['CUDA_VISIBLE_DEVICES'] = "0"  # 指定使用的GPU编号，"0"表示使用第一块GPU。如果有多个GPU，可以更改为其他编号。

# 配置TensorFlow使用的GPU资源
config = tf.compat.v1.ConfigProto()

# 设置GPU内存分配策略
config.gpu_options.allow_growth = True  # 动态分配GPU内存。程序刚开始会少量申请显存，然后随着需求增长，逐渐增加显存的分配，避免一次性占满整个显存。
"""
好处：默认情况下，TensorFlow会尝试占满所有可用的GPU显存，
这可能会导致其他GPU任务无法运行。通过设置 allow_growth 为 True，
TensorFlow会根据实际需求动态分配显存，而不是一次性占满。
"""

# 如果想设置最大GPU显存占用比例，可以使用下面的代码（可选）
# config.gpu_options.per_process_gpu_memory_fraction = 0.8  # 限制TensorFlow进程最多使用80%的GPU显存
"""
如果需要限制TensorFlow进程使用GPU显存的最大比例，
可以通过这个参数来设置。例如，设置为 0.8 表示最多使用80%的显存。
这在多任务环境中非常有用。
"""

# 创建Session，并应用配置
sess = tf.compat.v1.Session(config=config)

适用场景

• 多任务环境：如果一台机器上有多个GPU任务同时运行，控制每个任务的显存使用是非常重要的，这样可以避免冲突。
• 显存有限的GPU：在显存较小的GPU上运行大型模型时，动态显存分配可以帮助防止因显存不足而导致的程序崩溃。