tf.device()指定tensorflow运行的GPU或CPU设备

最新推荐文章于 2025-01-23 10:38:38 发布

苍蓝儿

最新推荐文章于 2025-01-23 10:38:38 发布

阅读量1.7k

点赞数 1

分类专栏：深度学习 Python TensorFlow

Python 同时被 3 个专栏收录

61 篇文章

订阅专栏

深度学习

32 篇文章

订阅专栏

TensorFlow

20 篇文章

订阅专栏

本文详细介绍如何在TensorFlow中控制模型运行的设备，包括指定GPU或CPU进行运算，以及如何选择特定的GPU。通过实例展示了使用tf.device()函数来指定运算在GPU还是CPU上执行的方法，并解释了配置参数log_device_placement的作用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在tensorflow中，我们可以使用 tf.device() 指定模型运行的具体设备，可以指定运行在GPU还是CUP上，以及哪块GPU上。

设置使用GPU
使用 tf.device('/gpu:1') 指定Session在第二块GPU上运行：

    import tensorflow as tf

    with tf.device('/gpu:1'):
        v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
        v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
        sumV12 = v1 + v2

        with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
            print sess.run(sumV12)

ConfigProto() 中参数 log_device_placement=True 会打印出执行操作所用的设备，以上输出：

如果安装的是GPU版本的tensorflow，机器上有支持的GPU，也正确安装了显卡驱动、CUDA和cuDNN，默认情况下，Session会在GPU上运行：

    import tensorflow as tf

    v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
    v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
    sumV12 = v1 + v2

    with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
        print sess.run(sumV12)

默认在GPU:0上执行：

设置使用cpu
tensorflow中不同的GPU使用/gpu:0和/gpu:1区分，而CPU不区分设备号，统一使用 /cpu:0

    import tensorflow as tf

    with tf.device('/cpu:0'):
        v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
        v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
        sumV12 = v1 + v2

        with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
            print sess.run(sumV12)

---------------------
作者：-牧野-
来源：优快云
原文：https://blog.youkuaiyun.com/dcrmg/article/details/79747882