在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)
到此这篇关于详解tf.device()指定tensorflow运行的GPU或CPU设备实现的文章就介绍到这了,更多相关tensorflow运行GPU或CPU内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！
您可能感兴趣的文章:- tensorflow 利用expand_dims和squeeze扩展和压缩tensor维度方式
- TensorFlow低版本代码自动升级为1.0版本
- TensorFlow的环境配置与安装方法
- TensorFlow2.0使用keras训练模型的实现
- 解决tensorflow模型压缩的问题_踩坑无数,总算搞定
咨 询 客 服