tensorflow gpu使用说明

最新推荐文章于 2024-08-13 13:48:39 发布

转载最新推荐文章于 2024-08-13 13:48:39 发布 · 219 阅读

文章标签：

本文介绍了如何高效利用GPU资源运行TensorFlow程序。首先通过nvidia-smi命令检查GPU使用情况，然后通过设置环境变量CUDA_VISIBLE_DEVICES指定使用哪个GPU。为了避免占用全部GPU显存，可以通过配置tf.ConfigProto限制显存使用。

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

原文地址为： tensorflow gpu使用说明

由于tensorflow默认抢占服务器所有GPU显存，只允许一个小内存的程序也会占用所有GPU资源。下面提出使用GPU运行tensorflow的几点建议：

1.在运行之前先查看GPU的使用情况：

$ nvidia-smi # 查看GPU此时的使用情况

或者

$ nvidia-smi -l # 实时返回GPU使用情况

2.目前实验室服务器有0，1，2，3四个GPU，找到空闲的GPU号，可以使用环境变量CUDA_VISIBLE_DEVICES：

环境变量的定义格式：
CUDA_VISIBLE_DEVICES=1 Only device 1 will be seen
CUDA_VISIBLE_DEVICES=0,1 Devices 0 and 1 will be visible
CUDA_VISIBLE_DEVICES=”0,1” Same as above, quotation marks are optional
CUDA_VISIBLE_DEVICES=0,2,3 Devices 0, 2, 3 will be visible; device 1 is masked

输入以下命令运行程序：

$ export CUDA_VISIBLE_DEVICES=0 # 假设此时 GPU 0 空闲

为了防止新开终端忘记export，比较保险的做法是每次运行tensorflow之前定义使用的GPU：

$ CUDA_VISIBLE_DEVICES=0 python mnist.py # 假设此时 GPU 0 空闲, mnist.py为你想运行的程序。

3.这样tensorflow此时只会在指定的GPU上运行，但是仍然会占用整个GPU的显存，不过不和其他人公用GPU时也不会有影响，下面介绍两种限定GPU占用的方法：

(1)在tensorflow中定义session时作如下设置，该设置会启用最少的GPU显存来运行程序。
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
(2)在tensorflow中定义session时作如下设置，该设置会强制程序只占用指定比例的GPU显存。
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 占用GPU40%的显存
session = tf.Session(config=config)

注：
- 在tensorflow代码中with tf.device(‘/gpu:0’):只会指定在GPU 0上计算，但仍然会默认占用所有GPU资源。

转载请注明本文地址： tensorflow gpu使用说明