PyTorch cuda 内存不足问题

Question

在 PyTorch 中训练模型时，我不断收到以下错误。我什至在代码开头添加了以下内容，但我不断收到此消息。我通过 Jupyter Notebook 运行它。

import gc
gc.collect()
torch.cuda.empty_cache()

我可以做什么来解决这个问题？

---

OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-6-2b42038d1b55> in <module>
     29 
     30         loss_mask = torch.mean((predicted_img - input_tensor) ** 2 * mask / mask_ratio)
---> 31         loss.backward()
     32 
     33         optim.step()

~/anaconda3/envs/ssenv/lib/python3.8/site-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    490                 inputs=inputs,
    491             )
--> 492         torch.autograd.backward(
    493             self, gradient, retain_graph, create_graph, inputs=inputs
    494         )

~/anaconda3/envs/ssenv/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    249     # some Python versions print out the first line of a multi-line function
    250     # calls in the traceback and some print out the last line
--> 251     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    252         tensors,
    253         grad_tensors_,

OutOfMemoryError: CUDA out of memory. Tried to allocate 146.00 MiB. GPU 0 has a total capacty of 9.62 GiB of which 100.94 MiB is free. Process 1485727 has 200.00 MiB memory in use. 

包括非 PyTorch 内存，该进程使用了 9.49 GiB 内存。在分配的内存中，8.96 GiB 由 PyTorch 分配，385.16 MiB 由 PyTorch 保留但未分配。 如果保留但未分配的内存很大，请尝试设置 max_split_size_mb 以避免碎片。请参阅内存管理和 PYTORCH_CUDA_ALLOC_CONF 的文档

Answer 1

您有几个选项可以尝试。

1. 减少批量大小 - 您也可以使用梯度累积来帮助这里。
2. 更改为混合精度训练 - 例如将张量从 FP32 变为 FP16。
3. 如果您可以使用多个 GPU，则可以探索分布式训练。确保在这里选择适当的策略来减少内存使用，而不仅仅是加快批处理速度。
4. 减小模型尺寸。
5. 根据您的数据，也许您可以使输入更小，例如降低图像分辨率、裁剪图像或使用灰度。