TensorFlow在指定CPU/GPU上執行的方法
轉自:http://wiki.jikexueyuan.com/project/tensorflow-zh/how_tos/using_gpu.html
使用 GPUs
支援的裝置
在一套標準的系統上通常有多個計算裝置. TensorFlow 支援 CPU 和 GPU 這兩種裝置. 我們用指定字串 strings
來標識這些裝置. 比如:
"/cpu:0"
: 機器中的 CPU"/gpu:0"
: 機器中的 GPU, 如果你有一個的話."/gpu:1"
: 機器中的第二個 GPU, 以此類推…
如果一個 TensorFlow 的 operation 中兼有 CPU 和 GPU 的實現, 當這個運算元被指派裝置時, GPU 有優先權. 比如matmul
cpu:0
和 gpu:0
中, matmul
operation
會被指派給 gpu:0
.
記錄裝置指派情況
為了獲取你的 operations 和 Tensor 被指派到哪個裝置上執行, 用 log_device_placement
新建一個 session
,
並設定為 True
.
- # 新建一個 graph.
- a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name=‘a’)
- b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3
, 2], name=‘b’) - c = tf.matmul(a, b)
- # 新建session with log_device_placement並設定為True.
- sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
- # 執行這個 op.
- print sess.run(c)
你應該能看見以下輸出:
- Device mapping:
- /job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
- id: 0000:05:00.0
- b: /job:localhost/replica:0
/task:0/gpu:0 - a: /job:localhost/replica:0/task:0/gpu:0
- MatMul: /job:localhost/replica:0/task:0/gpu:0
- [[ 22. 28.]
- [ 49. 64.]]
手工指派裝置
如果你不想使用系統來為 operation 指派裝置, 而是手工指派裝置, 你可以用 with tf.device
建立一個裝置環境, 這個環境下的 operation
都統一執行在環境指定的裝置上.
- # 新建一個graph.
- with tf.device('/cpu:0'):
- a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
- b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
- c = tf.matmul(a, b)
- # 新建session with log_device_placement並設定為True.
- sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
- # 執行這個op.
- print sess.run(c)
你會發現現在 a
和 b
操作都被指派給了 cpu:0
.
- Device mapping:
- /job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K40c, pci bus
- id: 0000:05:00.0
- b: /job:localhost/replica:0/task:0/cpu:0
- a: /job:localhost/replica:0/task:0/cpu:0
- MatMul: /job:localhost/replica:0/task:0/gpu:0
- [[ 22. 28.]
- [ 49. 64.]]
在多GPU系統裡使用單一GPU
如果你的系統裡有多個 GPU, 那麼 ID 最小的 GPU 會預設使用. 如果你想用別的 GPU, 可以用下面的方法顯式的宣告你的偏好:
- # 新建一個 graph.
- with tf.device('/gpu:2'):
- a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
- b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
- c = tf.matmul(a, b)
- # 新建 session with log_device_placement 並設定為 True.
- sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
- # 執行這個 op.
- print sess.run(c)
如果你指定的裝置不存在, 你會收到 InvalidArgumentError
錯誤提示:
- InvalidArgumentError: Invalid argument: Cannot assign a device to node 'b':
- Could not satisfy explicit device specification '/gpu:2'
- [[Node: b = Const[dtype=DT_FLOAT, value=Tensor<type: float shape: [3,2]
- values: 1 2 3...>, _device="/gpu:2"]()]]
為了避免出現你指定的裝置不存在這種情況, 你可以在建立的 session
裡把引數 allow_soft_placement
設定為 True
,
這樣 tensorFlow 會自動選擇一個存在並且支援的裝置來執行 operation.
- # 新建一個 graph.
- with tf.device('/gpu:2'):
- a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
- b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
- c = tf.matmul(a, b)
- # 新建 session with log_device_placement 並設定為 True.
- sess = tf.Session(config=tf.ConfigProto(
- allow_soft_placement=True, log_device_placement=True))
- # 執行這個 op.
- print sess.run(c)
使用多個 GPU
如果你想讓 TensorFlow 在多個 GPU 上執行, 你可以建立 multi-tower 結構, 在這個結構 裡每個 tower 分別被指配給不同的 GPU 執行. 比如:
- # 新建一個 graph.
- c = []
- for d in ['/gpu:2', '/gpu:3']:
- with tf.device(d):
- a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
- b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
- c.append(tf.matmul(a, b))
- with tf.device('/cpu:0'):
- sum = tf.add_n(c)
- # 新建session with log_device_placement並設定為True.
- sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
- # 執行這個op.
- print sess.run(sum)
你會看到如下輸出:
- Device mapping:
- /job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K20m, pci bus
- id: 0000:02:00.0
- /job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: Tesla K20m, pci bus
- id: 0000:03:00.0
- /job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: Tesla K20m, pci bus
- id: 0000:83:00.0
- /job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: Tesla K20m, pci bus
- id: 0000:84:00.0
- Const_3: /job:localhost/replica:0/task:0/gpu:3
- Const_2: /job:localhost/replica:0/task:0/gpu:3
- MatMul_1: /job:localhost/replica:0/task:0/gpu:3
- Const_1: /job:localhost/replica:0/task:0/gpu:2
- Const: /job:localhost/replica:0/task:0/gpu:2
- MatMul: /job:localhost/replica:0/task:0/gpu:2
- AddN: /job:localhost/replica:0/task:0/cpu:0
- [[ 44. 56.]
- [ 98. 128.]]