【tf学习笔记1】 tensorflow2.0+的相关硬件资源调用和设置问题

一只程序猿林

已于 2022-12-30 17:47:39 修改

阅读量773

点赞数

文章标签： tensorflow 学习人工智能 python 深度学习

于 2022-10-18 12:50:01 首次发布

本文链接：https://blog.csdn.net/qq_53383206/article/details/127386038

版权

1.获得本地计算机的硬件设备名称

首先使用tf.config.list_physical_devices命令获得本地设备列表
代码如下：

import tensorflow as tf
physical_gpus = tf.config.list_physical_devices("GPU")  # 获得本地GPU列表
physical_cpus = tf.config.list_physical_devices("CPU")  # 获得本地CPU列表
print(f"GPU:{physical_gpus}")
print(f"CPU:{physical_cpus}")
print("GPU个数:", len(physical_gpus))
print("CPU个数:", len(physical_cpus))

以下分析以GPU(CPU同理)为例
通过上述代码我们获得了本地计算机的GPU资源信息列表，下一步我们将进行tensorflow代码可见GPU的设置。

2.设置当前tf程序能够调用的GPU资源

以下代码表示设置当前tensorflow代码仅仅使用第一块GPU的资源。

tf.config.set_visible_devices(devices=physical_gpus[0], device_type="GPU")  # 设置使用第一块GPU

在实际工程应用中，还必须对GPU的显存进行可控的设置和管理，不能由tensorflow自身来完成对计算机GPU显存的管理。这是因为tensorflow会把当前能够调用的显存空间全部调用，进而导致显存的过度浪费，甚至导致程序的崩溃和计算机的短暂黑屏，容易对于本地GPU造成一定程度的危害和损伤。下一步将对tensorflow使用GPU的显存作出设置。

3.显存管理

3.1自动显存分配(用多少取多少)

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

使用上面的代码，能够在tensorflow程序运行过程中，动态分配显存。tensorflow会根据当前步骤需要用到的显存空间大小来对显存进行分配，避免程序一开始的盲目分配导致的显存碎片化和浪费。

3.2设置最大使用显存

tf.config.set_logical_device_configuration(physical_gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])

使用上面的代码，设置使用GPU0的资源数最大为1024MB=1GB。

4.单GPU模拟多GPU进行分布式训练

tf.config.set_logical_device_configuration(
    physical_gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=1024),
     tf.config.LogicalDeviceConfiguration(memory_limit=1024)])

上述代码是3.2的代码的扩充，目的是从实际GPU0资源中，分配两块GPU资源区，可以视为产生了两块虚拟GPU，虚拟GPU的生存周期是当前运行的tensorflow程序进程。
使用单机多卡训练MirroredStrategy 策略

strategy = tf.distribute.MirroredStrategy()
print('模拟GPU数量: %d' % strategy.num_replicas_in_sync)  # 输出设备数量变为2个实现模拟多GPU

运行上述代码可以发现，从程序视角来看，已经具有两块GPU可以参与分布式训练。