win10 + tf2.1 + cuda10.1 跑神经网络模型时,会出现'Internal: Invoking GPU asm compilation is supported on Cuda non-Windows platforms only','(GPU_0_bfc) ran out of memory'等错误信息,有时候报OOM的错误后奔溃,有时候又能用,这时不需要卸载cuda又重新安装cuda,那就说明GPU的配置出现了问题。有时候一个train能跑了以后,其他train也能跑了,又有时候一个train不能跑了以后其他train也不能跑,最坏的打算就是CPU跑,速度巨慢。
以下是tf2.0版本以上配置GPU的一些语句。
一、获得当前主机上CPU、GPU的列表
查看主机上的CPU以及是否有GPU,以列表的形式输出。
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print(gpus, cpus)
output:
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')] [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]
二、设置当前程序可见的设备范围
默认情况下 TensorFlow 会使用其所能够使用的所有 GPU。设置之后,当前程序只会使用自己可见的设备,不可见的设备不会被当前程序使用。
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
# 若GPU设备只有0, 可以设置os.environ['CUDA_VISIBLE_DEVICES'] = "1"用CPU跑
三、显存的申请与限制
1. 仅在需要时申请显存空间(程序初始运行时消耗很少的显存,随着程序的运行而动态申请显存);
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
2. 限制消耗固定大小的显存(程序不会超出限定的显存大小,若超出的报错)。下面的方式是设置Tensorflow固定消耗GPU:0的2GB显存。
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)
参考文章: