1. 程式人生 > >TensorFlow指定GPU或CPU執行相應運算w

TensorFlow指定GPU或CPU執行相應運算w

一、環境

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)