欢迎来到尧图网

客户服务 关于我们

您的位置:首页 > 房产 > 家装 > 解决tensorflow模型加载时把GPU显存占满的问题(亲测效果显著)

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

2024/11/30 15:33:28 来源:https://blog.csdn.net/weixin_44648900/article/details/141585958  浏览:    关键词:解决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上运行大型模型时,动态显存分配可以帮助防止因显存不足而导致的程序崩溃。

版权声明:

本网仅为发布的内容提供存储空间,不对发表、转载的内容提供任何形式的保证。凡本网注明“来源:XXX网络”的作品,均转载自其它媒体,著作权归作者所有,商业转载请联系作者获得授权,非商业转载请注明出处。

我们尊重并感谢每一位作者,均已注明文章来源和作者。如因作品内容、版权或其它问题,请及时与我们联系,联系邮箱:809451989@qq.com,投稿邮箱:809451989@qq.com