【深度学习走进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)])