码迷,mamicode.com
首页 > 其他好文 > 详细

TensorFlow中使用GPU

时间:2020-02-16 01:33:02      阅读:82      评论:0      收藏:0      [点我收藏+]

标签:rand   cuda   会话   device   代码   没有   ble   import   process   

TensorFlow默认会占用设备上所有的GPU以及每个GPU的所有显存;如果指定了某块GPU,也会默认一次性占用该GPU的所有显存。可以通过以下方式解决:

1 Python代码中设置环境变量,指定GPU

本文所有代码在tensorflow 1.12.0中测试通过。

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(...)

 

TensorFlow中使用GPU

标签:rand   cuda   会话   device   代码   没有   ble   import   process   

原文地址:https://www.cnblogs.com/zingp/p/12315366.html

(0)
(0)
   
举报
评论 一句话评论(0
登录后才能评论!
© 2014 mamicode.com 版权所有  联系我们:gaon5@hotmail.com
迷上了代码!