TensorFlow2.x 设置显存使用

记录一下使用TensorFlow遇到的第一个坑——显存管理。最开始不知道这个东西，每次跑完一遍都得重启…学习了Aurelier Geron的<<Hands-on Machine Learning with Scikit-learn, Keras & TensorFlow>>的第十九章才知道是咋回事。
原文截图：

按着我的习惯，先开一个terminal用于监控显存使用，每秒刷新一次。
watch -n 1 nvidia-smi
在训练的时候，TensorFlow默认会用尽所有的显存，所以需要设置memory_limit。（我的是2080单卡，8G显存，显存使用上限设置为4096M）
导入tensorflow之后立刻执行下面两句:
获得当前主机上特定运算设备的列表

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=4096)]
)

如果是单GPU模拟多GPU环境，就这样写

tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
     tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)