【深度學習 走進tensorflow2.0】TensorFlow 2.0 常用模塊tf.config

無意中發現了一個巨牛的人工智能教程，忍不住分享一下給大家。教程不僅是零基礎，通俗易懂，而且非常風趣幽默，像看小說一樣！覺得太牛了，所以分享給大家。點這裏可以跳轉到教程。人工智能教程

本篇文章將會教大家如何 合理分配顯卡資源，設置顯存使用策略。主要使用tf.config模塊進行設置。下面我們一起了解下具體用法和例子。

一、指定當前程序使用的 GPU

例如，在一臺具有 4 塊 GPU 和一個 CPU 的工作站上運行以下代碼：

import  tensorflow as tf
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus)
print(cpus)

輸出：

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

下面我需要指定只在0、1 這兩塊顯卡工作。

方法一：可以使用下面代碼：

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

方法二、使用環境變量 CUDA_VISIBLE_DEVICES 也可以控制程序所使用的 GPU

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

即可指定程序只在顯卡 0,1 上運行。

二、設置顯存使用策略

默認情況下，TensorFlow 將使用幾乎所有可用的顯存，以避免內存碎片化所帶來的性能損失。不過，TensorFlow 提供兩種顯存使用策略，讓我們能夠更靈活地控制程序的顯存使用方式：

 1. 僅在需要時申請顯存空間（程序初始運行時消耗很少的顯存，隨着程序的運行而動態申請顯存）；
 2. 限制消耗固定大小的顯存（程序不會超出限定的顯存大小，若超出的報錯）

可以通過 tf.config.experimental.set_memory_growth 將 GPU 的顯存使用策略設置爲 “僅在需要時申請顯存空間”。以下代碼將所有 GPU 設置爲僅在需要時申請顯存空間：

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu, enable=True)

假如我的顯存是4G，我只想用1G的顯存用來跑tensorflow2.0 程序，怎麼辦呢，可以從4G中，虛擬1G出來。

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

然後我們又想用但GPU模擬 多GPU環境，又如何操作呢、

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