【TensorFlow】GPU和显存资源的配置

最新推荐文章于 2025-02-08 20:07:55 发布

coderwait

最新推荐文章于 2025-02-08 20:07:55 发布

阅读量483

点赞数

分类专栏：深度学习框架

本文链接：https://blog.youkuaiyun.com/coderwait/article/details/109051564

版权

深度学习框架专栏收录该内容

1 篇文章

订阅专栏

本文详细介绍了如何在TensorFlow中指定和管理GPU资源。通过设置环境变量CUDA_VISIBLE_DEVICES，可以在命令行或Python代码中选择使用特定GPU。对于显存管理，可以设定固定比例或动态分配，以适应不同场景的需求，如个人使用、多用户共享服务器等。同时，建议在运行程序前检查GPU和显存使用情况，确保稳定运行。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

随时更新……

使用指定的GPU

GPU的编号在终端输入nvidia-smi查看。

（1）在终端设置

# 设置运行该程序可见的GPU
CUDA_VISIBLE_DEVICES = 0,1,2 python your_file.py  # 指定0#,1#,2#GPU
CUDA_VISIBLE_DEVICES = "3,4" python your_file.py  # 指定3#,4#GPU
CUDA_VISIBLE_DEVICES = ""  # 不使用GPU

（2）在Python代码设置

设置程序运行中可见的GPU设备。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

对于某个特定操作指定设备，则可以与tf.device一起使用来创建设备上下文，并且该上下文中的所有操作都将在同一指定设备上运行。

with tf.device('/CPU:0'):
  a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
 
c = tf.matmul(a, b)
print(c)

显存使用方式

（1）直接使用

会默认占用机器可用显卡剩余的所有GPU显存，即使根本用不到那么多。

with tf.Session() as sess:

由于直接占用全部显存，只要模型大小不超过显存大小就不会产生显存碎片。适合个人使用整个服务器，或部署应用。

（2）分配一定比例的显存

创建Session时设定最多占用的显存，避免其它用户没有显存可以用。

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用90%的显存
sess = tf.Session(config = config)

适合多用户共享服务器，但是可能存在显存浪费的情况。

（3）动态申请显存

在程序运行中动态申请显存，但是不会释放显存。若申请不到可用的显存就会报错。

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
with tf.Session(config=config) as sess:

适合多用户共享服务器，形式上避免了显存浪费，但很可能在运行程序时被其他用户抢占显存资源，导致程序崩溃。

记录操作分配给哪些设备

一些好的习惯

1. 运行程序前先查看GPU和显存使用情况。

参考链接

Tensorflow如何使用GPU训练（笔记）

tensorflow笔记：训练模型的时候如何指定GPU进行训练