本文主要是介绍【深度学习 走进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)])
这篇关于【深度学习 走进tensorflow2.0】TensorFlow 2.0 常用模块tf.config的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!