AIHGF

TensorFlow - GPU ID 和显存占用设定
Keras - GPU ID 和显存占用设定Pytorch - 多 GPUs 时的指定使用特定 GPUTensor...
扫描右侧二维码阅读全文
17
2018/11

TensorFlow - GPU ID 和显存占用设定

Keras - GPU ID 和显存占用设定

Pytorch - 多 GPUs 时的指定使用特定 GPU

Tensorflow 在利用显卡进行计算时,默认是将所有可用显卡资源一次性占满的,而不是在特定 GPU 上的自动显存分配. 对此,需要另外设定.

1. 查看 GPU 信息

import tensorflow as tf 


# 查看是否存在 GPU 
gpu_device_name = tf.test.gpu_device_name()
print(gpu_device_name)
# 输出: /device:GPU:0

# 查看 GPU 是否可用
tf.test.is_gpu_available()
print(tf.test.is_gpu_available())
# 输出: True 或 False 
# 获取本机所有可用设备
from tensorflow.python.client import device_lib

local_device_info = device_lib.list_local_devices()
print(local_device_info)

输出如:

[name: "/device:CPU:0"
 device_type: "CPU"
 memory_limit: 268435456
 locality {
 }
 incarnation: 5780123020024041651, name: "/device:XLA_GPU:0"
 device_type: "XLA_GPU"
 memory_limit: 17179869184
 locality {
 }
 incarnation: 11627108740584874474
 physical_device_desc: "device: XLA_GPU device", name: "/device:XLA_CPU:0"
 device_type: "XLA_CPU"
 memory_limit: 17179869184
 locality {
 }
 incarnation: 11108046644215655995
 physical_device_desc: "device: XLA_CPU device", name: "/device:GPU:0"
 device_type: "GPU"
 memory_limit: 7901999924
 locality {
   bus_id: 1
   links {
   }
 }
 incarnation: 5058955362673155543
 physical_device_desc: "device: 1, name: GeForce GTX 980 Ti, pci bus id: 0000:01:00.0, compute capability: 5.2"]

2. 设定 GPU ID

设定在特定 GPU ID 上进行计算,主要有两种方式 :

[1] - 命令行运行时添加指定 CUDA_VISIBLE_DEVICES=0, 1

如:

# 在 GPUID=0和1 的两块显卡上运行计算.
CUDA_VISIBLE_DEVICES=0,1 python demo.py

[2] - 在代码中指定 os.environ["CUDA_VISIBLE_DEVICES"] = "0"

如:

# 仅使用第 1 块 GPU
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"

3. 自动分配 GPU 显存

import tensorflow as tf 

# 自动分配显存
gpu_options = tf.GPUOptions(allow_growth=True)
# 设定固定显存,如 GPU显存 * 0.6
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.6)

config = tf.ConfigProto(gpu_options=gpu_options)

session = tf.Session(config=config)

# 
with tf.Session(config=config) as sess:
    sess.run()

或:

import tensorflow as tf 

config = tf.ConfigProto()
# 自动分配显存
config.gpu_options.allow_growth = True
# 设置固定显存比例
config.gpu_options.per_process_gpu_memory_fraction=0.6

session = tf.Session(config=config, ...)
Last modification:April 4th, 2019 at 02:43 pm

Leave a Comment