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 有效，其它版本的我不知道。