出现错误“资源耗尽：在分配具有形状 [1800,1024,28,28] 的张量并在 /job:localhost/... 上键入 float 时出现 OOM”

Question

我在开始训练我的对象检测 Tensorflow 2.5 GPU 模型时遇到资源耗尽错误。我正在使用 18 个训练图像和 3 个测试图像。我使用的预训练模型是来自 Tensorflow zoo 2.2 的 Faster R-CNN ResNet101 V1 640x640 模型。 我正在使用具有 8 GB 专用内存的 Nvidia RTX 2070 来训练我的模型。

我感到困惑的是，为什么在训练集很小的情况下，训练过程会占用我 GPU 的大量内存。这是我与错误一起获得的GPU内存摘要：

Limit:                      6269894656
InUse:                      6103403264
MaxInUse:                   6154866944
NumAllocs:                        4276
MaxAllocSize:               5786902272
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0

我还将训练数据的批量大小减少到 6，测试数据的批量大小减少到 1。

Answer 1

我在我在 gpu 上运行的所有笔记本中使用以下代码，以防止此类错误：

    import tensorflow as tf

    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
      try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
          tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
      except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)

默认情况下，TensorFlow 会将几乎所有 GPU 的所有 GPU 内存（受 CUDA_VISIBLE_DEVICES 限制）映射到进程可见。

More information about using gpu with tensorflow here

也许它会解决错误

希望对你有帮助

Answer 2

训练期间的最大内存使用量受多种因素影响，减少批量大小通常是解决内存限制的方法。 Alexandre Leobons Souza 的建议也可以帮助 Tensorflow 在分配内存方面具有更大的灵活性，但如果您继续看到 OOM 错误，那么我建议进一步减少批量大小。或者，您可以尝试限制模型中的可训练变量，这也会降低训练期间的内存使用量。

您提到，“我感到困惑的是，为什么在训练集如此之小的情况下，训练过程会从我的 GPU 中占用如此多的内存。”。需要记住的是，在训练期间，您的训练数据将用于模型的正向传递，然后您将在反向传递中计算每个可训练变量的梯度。即使您的训练数据很小，中间计算（包括梯度）也需要内存。这些计算与您的批量大小和模型大小成线性关系。通过减少批量大小或减少可训练变量的数量，训练将需要更少的内存。

另一个建议，如果您的输入张量的大小在您的训练数据中发生变化（即，如果地面实况边界框的数量从 1 变为 2，并且您没有填充输入张量），这可能会导致 Tensorflow在训练期间回溯计算图，您将看到警告。我不确定在这种情况下对内存的影响，但怀疑每次回溯实际上都需要内存中的重复模型。如果是这种情况，您可以尝试使用@tf.function(experimental_relax_shapes=True)。