computer-vision - 评估模型导致 CUDA 内存不足错误

Question

我的模型接收一个序列 RGB 图像 (512 x 512 x 3) 并据此进行预测。当我收到此错误时，我正在测试一系列 224 张图像：

运行时错误：CUDA 内存不足。尝试分配 4.00 GiB（GPU 0；15.78 GiB 总容量；10.21 GiB 已分配；1.06 GiB 空闲；PyTorch 总共保留 13.43 GiB）

这发生在 Colab 上，但我也在 AWS ml.p2.xlarge 笔记本实例上遇到了类似的错误。

我见过的所有解决方案都指向减少批量大小，但这对我来说没有意义，因为我的批量大小已经是 1：

火炬大小（[1, 224, 3, 512, 512]）

我该如何解决这个问题？看来我可能加载了一些错误的内容，需要 4.00 gigs

512 * 512 * 3 * 224 = 176,160,768

编辑：还意味着我尝试将工人数量更改为 1 并没有解决问题

标签： computer-visionpytorchgpu