python-tensorflow指定运行设备

#%%

import tensorflow as tf
import numpy as np
import os
import hashlib
import matplotlib.pylab as plt
import cv2
from PIL import Image

#%%
with tf.device("/gpu:0"):
    #这个代码块中定义的操作，会在tf.device给定的设备上运行
    #不是所有的操作都可以在GPU上运行的
    a = tf.Variable([1,2,3],dtype=tf.int32,name='a')
    b = tf.constant(2,dtype=tf.int32,name='b')
    c = tf.add(a,b)

#不同的变量可以在不同的设备上操作，但是所有这些节点还是在同一个
#图上，所以可以执行
with tf.device("/cpu:0"):
    d = tf.Variable([1,2,3],dtype=tf.int32,name='d')
    e = tf.constant(3,dtype=tf.int32,name='e')
    f = tf.add(d,e)

h = tf.add(f,c)

init_op = tf.global_variables_initializer()
with tf.Session(config=tf.ConfigProto(log_device_placement=True,allow_soft_placement=True)) as sess:
    sess.run(init_op)
    #执行结果
    print(sess.run([c,f]))
    print(sess.run(h))