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

(编辑:jimmy 日期: 2024/9/21 浏览:2)

在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  会打印出执行操作所用的设备,以上输出:

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

如果安装的是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上执行:

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

设置使用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设备实现

一句话新闻
微软与英特尔等合作伙伴联合定义“AI PC”:键盘需配有Copilot物理按键
几个月来,英特尔、微软、AMD和其它厂商都在共同推动“AI PC”的想法,朝着更多的AI功能迈进。在近日,英特尔在台北举行的开发者活动中,也宣布了关于AI PC加速计划、新的PC开发者计划和独立硬件供应商计划。
在此次发布会上,英特尔还发布了全新的全新的酷睿Ultra Meteor Lake NUC开发套件,以及联合微软等合作伙伴联合定义“AI PC”的定义标准。