tensorflow中显存分配

只想快乐

已于 2024-03-21 19:36:59 修改

阅读量395

点赞数 3

分类专栏：科研经验个人笔记文章标签： tensorflow 显存分配

于 2024-03-21 19:36:06 首次发布

本文链接：https://blog.csdn.net/weixin_44738399/article/details/136915385

版权

个人笔记同时被 2 个专栏收录

12 篇文章 0 订阅

订阅专栏

科研经验

10 篇文章 0 订阅

订阅专栏

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