第一次用tensorflow模型进行推理。
初始状态下,显存的占用情况
模型加载后,电脑显存直接快拉满了,但是模型参数量并不大,这就很奇怪,究竟这是一股什么神秘的力量,竟要试图侵占我可怜的显卡宝贝
原因分析:安装了tensorflow-gpu后,运行程序默认是把GPU的内存全部占满的,但是正常人都不希望这样,都希望能自己掌握自己的资源,而不是任人摆布,可以这样操作。
解决代码:
import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
"""
好处:默认情况下,TensorFlow会尝试占满所有可用的GPU显存,
这可能会导致其他GPU任务无法运行。通过设置 allow_growth 为 True,
TensorFlow会根据实际需求动态分配显存,而不是一次性占满。
"""
"""
如果需要限制TensorFlow进程使用GPU显存的最大比例,
可以通过这个参数来设置。例如,设置为 0.8 表示最多使用80%的显存。
这在多任务环境中非常有用。
"""
sess = tf.compat.v1.Session(config=config)
适用场景
- 多任务环境:如果一台机器上有多个GPU任务同时运行,控制每个任务的显存使用是非常重要的,这样可以避免冲突。
- 显存有限的GPU:在显存较小的GPU上运行大型模型时,动态显存分配可以帮助防止因显存不足而导致的程序崩溃。