python - 如何在 Eager 模式下分析 TensorFlow 1.15 内存分配

Question

我正在训练 YOLOv4。训练循环接收一批图像和几个描述预期目标检测框的 NumPy 数组。YOLOv4 有一个复杂的自定义损失函数。

我的训练循环（简化）如下所示：

for batch in data:
  with tf.GradientTape() as tape:
    pred = model(batch.images, training=True)
    loss = compute_loss(pred, batch.target_outputs)
    if not tf.is_nan(loss).numpy().any():
      gradients = tape.gradient(loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))

据我所知，形状batch.images没有改变。

问题是在几个 epochs 之后我开始得到，ResourceExhaustedError因为 TensorFlow 耗尽了 GPU 内存pred = model(batch.images, training=True)。

我想弄清楚是什么导致分配新内存。有没有办法获取整个 GPU 堆的快照，或者至少获取当前分配的字节数？

标签： pythontensorflow