physical_devices = tf.config.list_physical_devices('GPU')
for dev in physical_devices: # 如果使用多块GPU时
tf.config.experimental.set_memory_growth(dev, True)
将上面代码放在程序的开头, 即可做到限制显存使用的效果.
physical_devices = tf.config.list_physical_devices('GPU')
for dev in physical_devices: # 如果使用多块GPU时
tf.config.experimental.set_memory_growth(dev, True)
将上面代码放在程序的开头, 即可做到限制显存使用的效果.