1. 程式人生 > >TensorFlow——tensorflow指定CPU與GPU運算

TensorFlow——tensorflow指定CPU與GPU運算

1.指定GPU運算

如果安裝的是GPU版本,在執行的過程中TensorFlow能夠自動檢測。如果檢測到GPU,TensorFlow會盡可能的利用找到的第一個GPU來執行操作。

如果機器上有超過一個可用的GPU,除了第一個之外的其他的GPU預設是不參與計算的。為了讓TensorFlow使用這些GPU,必須將OP明確指派給他們執行。with......device語句能夠用來指派特定的CPU或者GPU執行操作:

import tensorflow as tf
import numpy as np

with tf.Session() as sess:
    with tf.device('/cpu:0'):
        a = tf.placeholder(tf.int32)
        b = tf.placeholder(tf.int32)
        add = tf.add(a, b)
        sum = sess.run(add, feed_dict={a: 3, b: 4})
        print(sum)

裝置的字串標識,當前支援的裝置包括以下的幾種:

cpu:0  機器的第一個cpu。

gpu:0  機器的第一個gpu,如果有的話

gpu:1  機器的第二個gpu,依次類推

類似的還有tf.ConfigProto來構建一個config,在config中指定相關的GPU,並且在session中傳入引數config=“自己建立的config”來指定gpu操作

其中,tf.ConfigProto函式的引數如下:

log_device_placement=True: 是否列印裝置分配日誌

allow_soft_placement=True: 如果指定的裝置不存在,允許TF自動分配裝置

import tensorflow as tf
import numpy as np

config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)

with tf.Session(config=config) as sess:
    a = tf.placeholder(tf.int32)
    b = tf.placeholder(tf.int32)
    add = tf.add(a, b)
    sum = sess.run(add, feed_dict={a: 3, b: 4})
    print(sum)

2.設定GPU使用資源

上文的tf.ConfigProto函式生成的config之後,還可以設定其屬性來分配GPU的運算資源,如下程式碼就是按需分配

import tensorflow as tf
import numpy as np

config = tf.ConfigProto(log_device_placement=True, allow_soft_placement=True)
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:
    a = tf.placeholder(tf.int32)
    b = tf.placeholder(tf.int32)
    add = tf.add(a, b)
    sum = sess.run(add, feed_dict={a: 3, b: 4})
    print(sum)

使用 allow_growth option,剛開始會分配少量的GPU容量,然後按需要慢慢的增加,有與不會釋放記憶體,隨意會導致記憶體碎片。

同樣,上述的程式碼也可以在config建立時指定,

import tensorflow as tf
import numpy as np

gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(gpu_options=gpu_options)


with tf.Session(config=config) as sess:
    a = tf.placeholder(tf.int32)
    b = tf.placeholder(tf.int32)
    add = tf.add(a, b)
    sum = sess.run(add, feed_dict={a: 3, b: 4})
    print(sum)

我們還可以給gpu分配固定大小的計算資源。

gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.5)

上述程式碼的含義是分配給tensorflow的GPU視訊記憶體大小為:GPU的實際視訊記憶體*0.5

&n