tensorflow在训练和测试过程中如果不特殊设置会一次占满所有内存然后再计算,这个以空间来换取时间,其实其计算效率并不高。但是在实际工程应用中内存占用过高会导致资源浪费。tensorflow提供了限制GPU分配内存的方法(此方法适用于tf2.x):
gpus = tf.config.experimental.list_physical_devices('GPU')
# 对需要进行限制的GPU进行设置
tf.config.experimental.set_virtual_device_configuration(gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=8192)])
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
tf.config.set_soft_device_placement = False
tf.config.experimental.set_memory_growth = True
gpus = tf.config.experimental.list_physical_devices('GPU')
print("gpus:", gpus)
if gpus:
tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=512)])
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), len(logical_gpus), 'Logical gpus')