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, ...)
4 comments
请问同一个工程不同模块，如何使得pytorch和tensorflow使用不同gpu？
同时使用两个框架的话，pytorch 和 tf 好像都支持不同显卡的分配的.
你好，在百度 AI Studio 里，local devices 的输出只有 "/device:XLA_GPU:0" 而没有 "/device:GPU:0"，用 with tf.devices("/xla_gpu:0") 和 with tf.devices("/gpu:0") 都不行，请问有办法使用那块 XLA_GPU:0 吗？谢谢
不好意思，这个好像暂时还没遇到过