使用 Optuna 时 CUDA 内存不足

Question

半个月前，我可以毫无问题地使用 Optuna 进行 48 小时的研究，大约 150+ 次试验。昨天，我在相同的模型、相同的数据集、相同的批量大小和相同的设备（A100 40GB 或 V100 32GB）上再次尝试了 Optuna，但在大约 16 次试验后我总是得到

torch.cuda.OutOfMemoryError: CUDA out of memory

。

按照SO上的一些答案（here和here），我尝试使用

study.optimize(objective, timeout=172800, gc_after_trial=True

启用GC，我还尝试在每个时期后添加

gc.collect()

和

cuda.empty_cache()

，但这些没有帮助。我什至尝试大幅减小隐藏层的规模（例如从 256 到 128），并减小数据集的大小（例如从 50000 到 20000），这些也没有帮助。我没有尝试的是在 Optuna 研究中减少批量大小，但我想这不能解决核心问题，迟早我会得到 OOM 错误。

我可以使用相同的参数和数据集训练（300 epoch 及更多）并评估模型，我最近在进行 Optuna 研究时才遇到这个 VRAM 外问题。

我想知道在 Optuna 研究期间是否有其他通用方法来避免 OOM。这真的很奇怪，因为我第一次使用 Optuna 就很好。

Answer 1

我也有同样的问题。 gc.collect() 和 cuda.empty_cache() 的问题是，这些方法不会从 GPU 中删除模型，它们只是清理缓存。

因此，您需要在每次试验后从 Cuda 内存中删除模型，并且可能还需要清理缓存，如果每次试验不执行此操作，新模型将保留在您的 Cuda 设备上。

所以我将这些行放在“目标”函数的末尾：

del model
torch.cuda.empty_cache()
gc.collect