Pytorch RuntimeError：CUDA错误：loss.backward()时内存不足，使用CPU时没有错误

Question

我正在具有超过 11G 内存的 Tesla K80 上训练用于语义分割的全卷积网络（FCN32）。

输入图像非常大：352x1216。网络结构如下图所示。我使用了batch_size=1，但仍然遇到 内存不足错误。

标准是 nn.BCEWithLogitsLoss()

当我在 CPU 上运行时，网络工作正常。

---

    Layer (type)               Output Shape       #  Param 
        Conv2d-1        [-1, 64, 352, 1216]           1,792
        Conv2d-2        [-1, 64, 352, 1216]          36,928
     MaxPool2d-3         [-1, 64, 176, 608]               0
        Conv2d-4        [-1, 128, 176, 608]          73,856
        Conv2d-5        [-1, 128, 176, 608]         147,584
     MaxPool2d-6         [-1, 128, 88, 304]               0
        Conv2d-7         [-1, 256, 88, 304]         295,168
        Conv2d-8         [-1, 256, 88, 304]         590,080
        Conv2d-9         [-1, 256, 88, 304]         590,080
    MaxPool2d-10         [-1, 256, 44, 152]               0
       Conv2d-11         [-1, 512, 44, 152]       1,180,160
       Conv2d-12         [-1, 512, 44, 152]       2,359,808
       Conv2d-13         [-1, 512, 44, 152]       2,359,808
    MaxPool2d-14          [-1, 512, 22, 76]               0
       Conv2d-15          [-1, 512, 22, 76]       2,359,808
       Conv2d-16          [-1, 512, 22, 76]       2,359,808
       Conv2d-17          [-1, 512, 22, 76]       2,359,808
    MaxPool2d-18          [-1, 512, 11, 38]               0
       Conv2d-19         [-1, 4096, 11, 38]     102,764,544
       Conv2d-20         [-1, 4096, 11, 38]      16,781,312
       Conv2d-21          [-1, 1, 11, 38]           4,097  ConvTranspose2d-22         [-1, 1, 352, 1216]             4,096

---

错误信息：

RuntimeError                              Traceback (most recent call last) <ipython-input-15-2b2b487c2e58> in <module>()
     36         print (loss)
     37         #torch.cuda.empty_cache()
---> 38         loss.backward()
     39         optimizer.step()
     40 

/anaconda/envs/py35/lib/python3.5/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
     91                 products. Defaults to ``False``.
     92         """
---> 93         torch.autograd.backward(self, gradient, retain_graph, create_graph)
     94 
     95     def register_hook(self, hook):

/anaconda/envs/py35/lib/python3.5/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     88     Variable._execution_engine.run_backward(
     89         tensors, grad_tensors, retain_graph, create_graph,
---> 90         allow_unreachable=True)  # allow_unreachable flag
     91 
     92 

RuntimeError: CUDA error: out of memory

Answer 1

通常发生这种情况是因为 GPU 上的内存。如果您有更强大的 GPU，您的问题就可以解决（正如您在回答中提到的）。
但如果您没有，您可以将图像缩小到大约 256*x

 尺寸。
为了

性能，这也是一个很好的做法。

Answer 2

我找到原因了...与硬件有关。我换了另一台机器，错误消失了。