TensorFlow 2.0 硬體資源設置

tags: tensorflow tf2.0

參考來源：

• 簡單粗暴TensorFlow 2.0

預設情況下，tensorflow會為了避免碎片化的記憶體造成效能不佳的情況而一次性的佔用顯卡所有記憶體，這部份可以透過tf.config來設置，這邊說明的是限制硬體資源以及顯卡記憶體的使用。

說明

我們先瞭解硬體資源的限制

首先，導入需求套件：

import tensorflow as tf

接下來我們可以先取得硬體清單：

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'),
PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')] 
[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

上面的清單明白看出硬體資源擁有兩張顯卡以及一顆CPU。

如果，我們這次的訓練過程只需要使用一塊顯卡的話，可以透過tf.config.experimental.set_visible_devices設置：

tf.config.experimental.set_visible_devices(devices=gpus[0], device_type='GPU')

上面的設置方式等價於直接使用環境變數設置，如下：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"

或者你也可以這麼做：

export CUDA_VISIBLE_DEVICES=0

記憶體的部份，有兩種設置方式，分別為依需求動態調整，以及一次性給足。

先看動態調整的設置，這可以透過tf.config.experimental.set_memory_growth設置：

for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu, True)

一次性給足的情況下，如果需求記憶體超過限制，那就會拋出異常。

固定限制的設置可以透過tf.config.experimental.set_virtual_device_configuration設置：

tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])

上面案例設置限制使用記憶體為4g

如果我們只有一塊顯卡，模型不大，希望可以將這單張顯卡模擬為多卡來加快訓練，這可以透過tf.config.experimental.set_virtual_device_configuration來設置：

tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
     tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])

上面案例將單張顯卡模型為2張擁有2g記憶體的顯卡