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可以参与分布式训练。