TensorFlow教程之進階指南 3.9 使用 GPUs

本文檔為tensorflow參考文檔，本轉載已得到tensorflow中文社群授權。

在一套标準的系統上通常有多個計算裝置. tensorflow 支援 cpu 和 gpu 這兩種裝置. 我們用指定字元串<code>strings</code> 來辨別這些裝置. 比如:

<code>"/cpu:0"</code>: 機器中的 cpu

<code>"/gpu:0"</code>: 機器中的 gpu, 如果你有一個的話.

<code>"/gpu:1"</code>: 機器中的第二個 gpu, 以此類推...

如果一個 tensorflow 的 operation 中兼有 cpu 和 gpu 的實作, 當這個算子被指派裝置時, gpu 有優先權. 比如<code>matmul</code>中 cpu 和 gpu kernel 函數都存在. 那麼在 <code>cpu:0</code> 和 <code>gpu:0</code> 中, <code>matmul</code> operation 會被指派給 <code>gpu:0</code> .

為了擷取你的 operations 和 tensor 被指派到哪個裝置上運作, 用 <code>log_device_placement</code> 建立一個 <code>session</code>, 并設定為 <code>true</code>.

你應該能看見以下輸出:

如果你不想使用系統來為 operation 指派裝置, 而是手工指派裝置, 你可以用 <code>with tf.device</code> 建立一個裝置環境, 這個環境下的 operation 都統一運作在環境指定的裝置上.

你會發現現在 <code>a</code> 和 <code>b</code> 操作都被指派給了 <code>cpu:0</code>.

如果你的系統裡有多個 gpu, 那麼 id 最小的 gpu 會預設使用. 如果你想用别的 gpu, 可以用下面的方法顯式的聲明你的偏好:

如果你指定的裝置不存在, 你會收到 <code>invalidargumenterror</code> 錯誤提示:

為了避免出現你指定的裝置不存在這種情況, 你可以在建立的 <code>session</code> 裡把參數 <code>allow_soft_placement</code> 設定為 <code>true</code>, 這樣 tensorflow 會自動選擇一個存在并且支援的裝置來運作 operation.

如果你想讓 tensorflow 在多個 gpu 上運作, 你可以建立 multi-tower 結構, 在這個結構 裡每個 tower 分别被指配給不同的 gpu 運作. 比如:

你會看到如下輸出: