tensorflow中显存分配

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