TensorFlow 默认会占用设备上所有的 GPU 以及每个 GPU 的所有显存; 如果指定了某块 GPU, 也会默认一次性占用该 GPU 的所有显存. 可以通过以下方式解决:
1 Python 代码中设置环境变量, 指定 GPU
- import os
- os.environ["CUDA_VISIBLE_DEVICES"] = "2" # 指定只是用第三块 GPU
2 系统环境变量中指定 GPU
# 只使用第 2 块 GPU, 在 demo_code.py, 机器上的第二块 GPU 变成 "/gpu:0", 不过在运行时所有的 / gpu:0 的运算将被放到第二块 GPU 上
CUDA_VISIBLE_DEVICES=1 python demo_code.py
- # 只使用第一块 GPU 和第二块 GPU
- CUDA_VISIBLE_DEVICES=0,1 python demo_code.py
3 动态分配 GPU 显存
- # allow_soft_placement=True 没有 GPU 的话在 CPU 上运行
- config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
- config.gpu_options.allow_growth = True # 按需分配显存
- with tf.Session(config=config) as sess:
- sess.run(...)
4 按固定比例分配显存
- # 按照固定的比例分配.
- config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
- # 以下代码会占用所有可使用的 GPU 的 40% 显存
- config.gpu_options.per_process_gpu_memory_fraction = 0.4
- with tf.Session(config=config) as sess:
- sess.run(...)
在我的设备中设置后 GPU 占用情况如下:
- gz_6237_gpu Sat Feb 15 23:01:56 2020 418.87.00
- [0] GeForce RTX 2080 Ti | 43'C, 0 % | 4691 / 10989 MB | dc:python/1641(4681M)
5 通过 tf.device 将运算指定到特定设备上
- with tf.device("/gpu:0"):
- b = tf.Variable(tf.zeros([1]))
- W = tf.Variable(tf.random_uniform([1, 2], -1.0, 1.0))
- y = tf.matmul(W, x_data) + b
这种方式不推荐. TF 的 kernel 中国定义了哪些操作可以跑在 GPU 上, 哪些不可以, 因此强制指定 GPU 会降低程序的可移植性.
推荐的做法是: 在创建会话时, 指定参数 allow_soft_placement=True; 这样如果运算无法在 GPU 上执行, TF 会自动将它放在 CPU 上执行.
- config = tf.ConfigProto(allow_soft_placement=True)
- with tf.Session(config=config) as sess:
- sess.run(...)
来源: https://www.cnblogs.com/zingp/p/12315366.html