待补充图片和1.4
1. tensorflow2
默认情况下,tensorflow2会自动选择设备,如果有GPU,优先使用GPU,且自动占满多块GPU显存。
1.1 设置使用cpu
os.environ['CUDA_VISIBLE_DEVICES'] = "-1"
tf.config.experimental.set_visible_devices(devices=cpus[0], device_type='CPU')
1.2 设置使用GPU个数
首先可以先看看自己设备的GPU和CPU情况
# GPU
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
print(os.environ['CUDA_VISIBLE_DEVICES'])
# CPU
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
- 如果有多块GPU,可以设置使用其中的几块,方法有三种:
- 设置当前程序可见的设备范围:
tf.config.experimental.set_visible_devices(devices=gpus[2:4], device_type='GPU')
,设置设备2,3,4可见,当前程序只会使用自己可见的设备,不可见的设备不会被当前程序使用 - 在终端使用环境变量限制GPU:
export CUDA_VISIBLE_DEVICES=2,3
- 在代码中指定GPU设备(类似于第一种方法):
os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"
- 设置当前程序可见的设备范围:
1.3设置使用GPU显存大小
tensorflow提供了两种限定显存的方法:
- 动态申请内存,需要多少,申请多少
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
- 限制消耗固定大小的显存(程序不会超出限定的显存大小,若超报错)
# 限制使用2G显存
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)
1.4 单GPU模拟多GPU环境
当只有一块GPU时,假如显存为10G,但我们的程序可能只占用5G,那么剩下的5G不就浪费了,这时可以将这块10G的显存分成两块,这样相当于分布式使用2块GPU
tf.config.experimental.set_virtual_device_configuration(
gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)])
按照上述,2块GPU也可以当做4块使用,具体方法待补充