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

Question

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

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

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

任何建议将不胜感激！ 提前致谢！

更新 1 正在使用的代码在Colab 输入数据来自 UCF-Crime 数据集，它已被预处理为 jpeg 并作为片段存储在 TFRecord 中。 来自同一 TFRecord 的示例是 here，每个 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 上则没有。

Answer 1

Tensorflow Profiler 应该可以帮到你。

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

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

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

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

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

Custom Training的使用方法请参考this articleTensorflow Profiler。

有关Tensorflow Profiler的更多信息，请参考this Tutorial、this Guide和this documentation。