TensorFlow指定GPU或CPU執行相應運算w
阿新 • • 發佈:2018-12-13
一、環境
TensorFlow API r1.12
CUDA 9.2 V9.2.148
Python 3.6.3
二、官方說明
https://www.tensorflow.org/api_docs/python/tf/device
使用預設圖(default graph)是Graph.device()的包裝器(wrapper)
詳細參照tf.Graph.device
輸入:
在上下文中使用的裝置名稱或者函式
輸出:
為新建立的操作指定預設裝置的上下文管理器
三:例項
(1)使用CPU
with tf.device("/cpu:0"): embedding = tf.get_variable("embedding", [vocab_size, size], dtype=tf.float32) inputs = tf.nn.embedding_lookup(embedding, input_.input_data)
(2)使用GPU
with tf.device("/gpu:0"):
embedding = tf.get_variable("embedding", [vocab_size, size], dtype=tf.float32)
inputs = tf.nn.embedding_lookup(embedding, input_.input_data)