详解TensorFlow报”ResourceExhaustedError: Conv2DSlowBackpropInput: OOM when allocating tensor with shape “的原因以及解决办法

问题描述

在使用 TensorFlow 进行模型训练或推理时，有时候会出现如下错误提示：

ResourceExhaustedError: Conv2DSlowBackpropInput: OOM when allocating tensor with shape ...

这种错误提示一般是由于内存不足造成的，即在运行 TensorFlow 时需要分配的内存超过了系统可用的内存。

错误原因

针对这个错误，可能出现的原因有很多，比如模型设计不合理、数据输入格式不对、模型参数太多等等，但是最主要的原因还是内存不足。

在某些情况下，可能会发现数据集很小，模型也很简单，但是依然会出现 OOM 错误。这个问题很可能是由于 TensorFlow 会自动分配所有可用的 GPU 存储空间，导致其他进程无法访问该 GPU。因此，如果在 GPU 上运行 TensorFlow 时遇到 OOM 错误，请确保 TensorFlow 不会占用所有可用存储空间。

解决办法

针对上述问题，提供以下几个解决办法：

降低 batch size：

如果训练数据集中包含大量的数据，那么可以尝试使用更小的 batch size。这样可以减轻 GPU 的负担，降低内存使用量。

使用更小的模型：

可以尝试通过减少模型复杂度来减轻内存压力。去除一部分卷积层、降低卷积核大小、减少层数等措施都是可行的。

使用 float16：

使用 float16 的计算精度比 float32 低，但可以减少存储需求。在某些情况下，可以通过将模型参数转换为 float16 降低内存占用。

使用 TensorFlow 的 GradientTape：

使用 GradientTape 可以更好地控制 TensorFlow 的内存使用。该方法可以在 TensorFlow 2.0 中使用，并且可以很好地处理内存问题。

设置 allow_growth 为 True：

在运行时，可以将 allow_growth 设为 True，这样 TensorFlow 将只分配所需的 GPU 存储空间，而不是将所有可用空间全部占用。可以使用以下代码来实现：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

使用更大的 GPU：

如果以上方法都无效，则可以考虑使用更大的 GPU。

以上是解决 OOM 问题的几种方法，具体的实践中可以根据自己的情况进行选择。