获得当前主机上运算设备列表
1 2 3 4
| gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices( device_type='CPU')
|
设置当前程序可见的设备范围
当前程序只会使用自己可见的设备,不可见的设备不 会被当前程序使用,因此可以设置当前程序可见的设 备范围。
1 2 3 4
| gpus=tf.config.experimental.list_physical_devices(device_type='GPU') #可见设备范围 tf.config.experimental.set_visible_devices( devices=gpus[0:2], device_type='GPU')#设置前两快可见
|
指定程序只在显卡 2,3 上运行。
1 2
| import os os.environ['CUDA_VISIBLE_DEVICES'] = "2,3"#即可指定程序只在显卡 2,3 上运行。
|
设置显存使用策略
默认情况下,TensorFlow 将使用几乎所有可用的显存 ,以避免内存碎片化所带来的性能损失
TensorFlow 提供两种显存使用策略,让我们能够更灵 活地控制程序的显存使用方式:
(1)仅在需要时申请显存空间(程序初始运行时消耗 很少的显存,随着程序的运行而动态申请显存;
GPU 的显存使用策略设置为 “仅在需要时申请显存空间”:
1 2 3 4
| #GPU 的显存使用策略设置为 “仅在需要时申请显存空间”: gpus = tf.config.experimental.list_physical_devices(device_type='GPU') for gpu in gpus: tf.config.experimental.set_memory_growth(device=gpu, True)#device gpu名字
|
(2)限制消耗固定大小的显存(程序不会超出限定的 显存大小,若超出的报错)。
设置 TensorFlow 固定消耗 GPU:0 的 1GB 显存:
1 2 3
| #设置 TensorFlow 固定消耗 GPU:0 的 1GB 显存: 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)])
|