我正在使用 tf.keras 训练图像分割模型，使用自定义数据生成器来读取和增强图像。虽然训练模型工作正常（即没有内存问题），但当尝试在我的测试集上预测时，我的 GPU（8GB，稍后参见 nvidia-smi）会耗尽内存。无论是在训练后直接预测还是在重新启动内核后、使用model.load_weights()和使用model.predict()之后加载模型以及使用训练中使用的相同批处理大小（4，在训练期间使用约 6GB 内存）或批处理大小为 1 的情况下，都是这种情况两个批量大小都尝试分配超过 8GB。

在训练期间，内存使用量稳定在 6GB 左右，但在使用时，model.predict()它开始时约为 6GB，但在抛出之前大约 10 秒后跳至 8GB ResourceExhaustedError（请参阅稍后的堆栈跟踪）。这对我来说似乎非常违反直觉，而且我通过谷歌找到的提示（例如重新启动python，从权重加载模型，然后预测释放事先使用的内存）还没有起作用，所以任何帮助都会很棒。

我的数据生成器和训练/预测代码（包括错误消息）的输出!nvidia-smi如下：

英伟达-SMI

Mon Aug  9 14:27:29 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 471.11       Driver Version: 471.11       CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name            TCC/WDDM | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce …