TensorFlow 限制显存使用量

最近更新于 2024-05-05 14:19

对于高配的计算机,可能都不用考虑限制显存这个问题,但是我要考虑。我轻薄本就 2GB 独显,默认加载数据时会一次性全部往显存里写,结果动不动就提示显存不够,然后程序 game over。

另外就是对于开发板,我这里也在用 Jetson Nano 的开发板,GPU 没有独立显存,是和 CPU 共用 4GB 的内存,只要一跑 tf 程序,内存蹭蹭往上飙,着不住啊。带来的就是 CPU 可用的内存减少,那么就卡起来了,所以在 CPU 和 GPU 共用内存的时候,要限量使用,而不是让 GPU 无限制的占用内存,反而性能不高。

下面就是限制内存的方法

gpus= tf.config.list_physical_devices('GPU')
if gpus:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])  # 这里是限制 GPU 可用内存为 1024M = 1GB

针对 TensorFlow 2.x 有效,其它版本的我不知道。

TensorFlow 限制显存使用量
Scroll to top