如何在PyTorch中避免“CUDA内存不足”错误

10 浏览
0 Comments

如何在PyTorch中避免“CUDA内存不足”错误

对于GPU内存较低的PyTorch用户来说，我认为这是一个相当常见的信息：

RuntimeError: CUDA内存不足。尝试分配X MiB（GPU X; 总容量X GiB; 已分配X GiB; 剩余X MiB; 缓存X）

我尝试通过将每个层加载到GPU上并将其加载回来来处理图像：

for m in self.children():
    m.cuda()
    x = m(x)
    m.cpu()
    torch.cuda.empty_cache()

但似乎效果不是很好。我想知道是否有一些技巧和窍门可以在使用较少GPU内存的情况下训练大型深度学习模型。

0