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

Miles

发布时间：2026-10-14| 4人浏览过

原创

本文详细介绍如何在TensorFlow中使用tf.device()指定模型运行的具体设备，包括GPU和CPU的选择。通过示例代码展示如何指定特定GPU运行Session，以及如何配置ConfigProto打印设备放置信息。同时解释了默认情况下Session在GPU:0上运行的机制，以及如何将计算指定到CPU。

在TensorFlow框架中，利用 tf.device() 函数可以明确指定模型运算所依托的具体硬件设备，用户能够灵活选择让程序在GPU还是CPU上执行，并进一步指定具体使用哪一块GPU资源。

设置使用GPU

若需指定Session在第二块GPU上运行，可使用 tf.device('/gpu:1') 进行配置：


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)