tf.device() 查看和指定 GPU或CPU 设备

GPU_device.py：

# https://blog.csdn.net/lemon_tree1994/article/details/79206790
# https://blog.csdn.net/dcrmg/article/details/79747882
import tensorflow as tf
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"   #指定gpu
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))   #打印gpu信息
# os.system("pause")

def examine():
    # TensorFlow可以通过设置 log_device_placement 参数来打印运行每一个运算的设备。
    a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
    b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
    c = a + b
    # 通过log_device_placement参数来输出运行每一个运算的设备。
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    print(sess.run(c))

def allocate():
    # TensorFlow会自动优先将运算放置在GPU上
    # 通过tf.device将运算指定到特定的设备上
    # 如果需要将某些运算放到不同的GPU或者CPU上，就需要通过tf.device来手工指定
    with tf.device('/cpu:0'):
       a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
       b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
    with tf.device('/gpu:1'):
        c = a + b

    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    print(sess.run(c))

if __name__=='__main__':
    examine()
    # allocate()

 

或者在pycharm 中使用 CUDA_VISIBLE_DEVICES：

参考资料： https://www.cnblogs.com/darkknightzh/p/8576825.html