python - 如何在 TensorFlow 2 上通过 Eager Execution 调试内存不足问题？

Question

我正在尝试在 TensorFlow 2.2 中拟合模型。我写了一个自定义训练循环。但是，由于 GPU 内存不足，训练很快就会崩溃。该模型在使用 model.fit() API 时使用相同的参数可以正常工作，但我想使用自定义训练循环，因为它为我的需求提供了更大的灵活性。

调试此类内存问题的一般方法是什么？

我四处搜索，但官方 TF 文档大多谈论如何调试逻辑错误。很多页面都在讨论基于 Graph 模式的调试。

任何意见，将不胜感激！提前致谢！

更新 1 正在使用的代码在Colab 上输入数据来自 UCF-Crime 数据集，它已被预处理为 jpeg 并作为片段存储在 TFRecord 中。这里是一个 TFRecord 示例，每个 TFRecord 包含 500 个片段，其中每个片段是视频的 16 个连续 jpeg 编码帧，缩小为 128x128 RGB 图像。

为了重现问题，您可以通过将 TFRecord 的文件路径放入来构建数据集mod_build_dataset()

代码将立即在 Colab 中的 CPU 上崩溃（因为可用的最大 RAM 约为 13GB），GPU 也是如此。

在 Kaggle 上它可以在 CPU 上正常工作（RAM 使用最大为 14.2 GB）但会在 GPU 上崩溃（使用所有 16GB 的图形内存）

我猜这是因为 CPU 的计算速度较慢，因此 GC 在内存不足之前有时间启动，而在 GPU 上则没有。

标签： pythontensorflowtensorflow2.0

Answer 1

Tensorflow Profiler应该可以帮助你。

Profiling帮助您了解各种硬件资源消耗（time和memory）并解决性能瓶颈，最终使模型执行得更快。TensorFlow operations (ops)model

这Tensorflow Profiler使得精确定位bottleneck变得training process更加容易，因此您可以决定应该在哪里进行优化。

它还为您提供了有关您可以遵循以优化模型性能的潜在后续步骤的建议。

使用 TensorFlow Profiler 进行自定义训练循环的步骤如下所示：

from tensorflow.python.profiler import profiler_v2 as profiler
profiler.warmup()
profiler.start(logdir='logdir')
# Train the model here
profiler.stop()

请参阅这篇文章了解如何使用Tensorflow Profilerfor Custom Training。

有关更多信息，Tensorflow Profiler请参阅本教程、本指南和本文档。