tensorflow计算网络占用内存_记录 TensorFlow Eager 模式下的一次内存泄漏

最新推荐文章于 2022-08-17 20:45:35 发布

weixin_39640262

最新推荐文章于 2022-08-17 20:45:35 发布

阅读量290

点赞数

文章标签： tensorflow计算网络占用内存

本文链接：https://blog.csdn.net/weixin_39640262/article/details/111361864

版权

本文记录了在 TensorFlow Eager 模式下遇到的内存泄漏问题，通过实时监控显存占用，定位到问题代码，并提供了解决方案。在训练 Object Detection 模型时，显存持续增加，最终发现是特定操作导致内存泄漏。通过修改代码，成功解决了内存泄漏问题，能够继续训练大数据集。

摘要由CSDN通过智能技术生成

问题描述：
- 在实现 TensorFlow Eager 版 Object Detection 的过程中碰到一个问题，即训练若干epoch后，占用的显存会大幅增加。
- 如果要训练较多的steps时，会由于GPU不足而导致训练失败。
- 猜测是内存泄漏问题。
搜索到的有关Eager模式下内存泄漏的一些资料：
- issue #19385: Memory leak in Eager execution
- issue #19499: tf.data.Dataset iterators are not cleaned when the loop ends with a break
- issue #19671: Variables may live longer than they suppose to on Eager Execution
  - 可能跟这个issue类似，提到可以通过tf.set_random_seed(1)来解决，但没解决我的问题。
- issue #20062: Memory leak using loss in Eager Execution
- issue #20218: ran out of memory in eager execution
- issue #20436: Memory leaking in tf.data.Dataset in eager mode
TODO List
- 研究Eager模式的底层C++实现思路。
- 研究Eager模式下的性能调优。
身边也没什么人能讨论这个问题，都是自己慢慢google解决，可能不是最优解，但起码解决了问题……

最开始的思路：查看当前占用显存的是哪些对象。
采坑过程：
- 思路很清晰，但最终没找到方法可以直接实现……（后来转念一想，如果真有这样的方法，那不是内存泄漏问题都可以很容易定位了……）
- 替代方案：由于在Graph模式中内存泄漏问题，很多来源于训练过程中多次重复建立Op，所以希望查看训练过程中创建了哪些Op。这可以通过 tf.ConfigProt() 中的 log_device_placement实现，代码如下。
</