tensorflow中显存分配
- 问题:使用tensorflow-gpu训练模型,GPU的显存都是占满的。
# GPU 1的显存将占满
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
- 原因:默认情况下,tensorflow会把可用的显存全部占光,无论显卡中剩余多少显存,都会被tensorflow占用。
- 解决方法(针对tensorflow2.x的版本)
# 查看可用的gpu
gpu_list = tf.config.experimental.list_physical_devices('GPU')
# 指定使用的GPU
tf.config.experimental.set_visible_devices(gpu_list[1], 'GPU')
# 指定使用的GPU上,按需分配内存
tf.config.experimental.set_memory_growth(gpu_list[1], True)
参考博客
[1]: https://blog.csdn.net/hxj0323/article/details/121629067
[2]: https://zhuanlan.zhihu.com/p/38512293