为什么在我的模型 init() 上发生 OOM？

Question

我的模型中的一行 tr.nn.Linear(hw_flat * num_filters*8, num_fc) 导致模型初始化时出现 OOM 错误。将其注释掉可以消除内存问题。

import torch as tr
from layers import Conv2dSame, Flatten

class Discriminator(tr.nn.Module):
    def __init__(self, cfg):
        super(Discriminator, self).__init__()
        num_filters = 64
        hw_flat = int(cfg.hr_resolution[0] / 2**4)**2
        num_fc = 1024

        self.model = tr.nn.Sequential(
            # Channels in, channels out, filter size, stride, padding
            Conv2dSame(cfg.num_channels, num_filters, 3),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters, num_filters, 3, 2),
            tr.nn.BatchNorm2d(num_filters),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters, num_filters*2, 3),
            tr.nn.BatchNorm2d(num_filters*2),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters*2, num_filters*2, 3, 2),
            tr.nn.BatchNorm2d(num_filters*2),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters*2, num_filters*4, 3),
            tr.nn.BatchNorm2d(num_filters*4),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters*4, num_filters*4, 3, 2),
            tr.nn.BatchNorm2d(num_filters*4),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters*4, num_filters*8, 3),
            tr.nn.BatchNorm2d(num_filters*8),
            tr.nn.LeakyReLU(),
            Conv2dSame(num_filters*8, num_filters*8, 3, 2),
            tr.nn.BatchNorm2d(num_filters*8),
            tr.nn.LeakyReLU(),
            Flatten(),
            tr.nn.Linear(hw_flat * num_filters*8, num_fc),
            tr.nn.LeakyReLU(),
            tr.nn.Linear(num_fc, 1),
            tr.nn.Sigmoid()
        )
        self.model.apply(self.init_weights)

    def forward(self, x_in):
        x_out = self.model(x_in)
        return x_out

    def init_weights(self, layer):
        if type(layer) in [tr.nn.Conv2d, tr.nn.Linear]:
            tr.nn.init.xavier_uniform_(layer.weight)

这很奇怪，因为hw_flat = 96*96 = 9216，而num_filters*8 = 512，所以hw_flat * num_filters = 4718592，就是那个层的参数个数。我已确认此计算，因为将图层更改为 tr.nn.Linear(4718592, num_fc) 会产生相同的输出。

对我来说，dtype=float32 没有任何意义，因此它的预期大小为 32*4718592 = 150,994,944 字节。这相当于大约 150mb。

错误信息是：

Traceback (most recent call last):
  File "main.py", line 116, in <module>
    main()
  File "main.py", line 112, in main
    srgan = SRGAN(cfg)
  File "main.py", line 25, in __init__
    self.discriminator = Discriminator(cfg).to(device)
  File "/home/jpatts/Documents/ECE/ECE471-SRGAN/models.py", line 87, in __init__
    tr.nn.Linear(hw_flat * num_filters*8, num_fc),
  File "/home/jpatts/.local/lib/python3.6/site-packages/torch/nn/modules/linear.py", line 51, in __init__
    self.weight = Parameter(torch.Tensor(out_features, in_features))
RuntimeError: $ Torch: not enough memory: you tried to allocate 18GB. Buy new RAM! at /pytorch/aten/src/TH/THGeneral.cpp:201

我也只运行 1 的批量大小（不会影响此错误），网络的整体输入形状为 (1, 3, 1536, 1536)，展平层后的形状为 (1, 4718592 )。

为什么会这样？

Answer 1

您的线性层非常大 - 事实上，它确实需要至少 18GB 的​​内存。 （您的估计有两个原因：（1）float32 占用 4 个字节的内存，而不是 32 个字节，以及（2）您没有乘以输出大小。）

来自PyTorch documentation FAQs：

不要使用太大的线性图层。线性层nn.Linear(m, n) 使用O(n*m) 内存：也就是说，权重的内存需求与 特征的数量。以这种方式很容易破坏你的记忆（并且 请记住，您至少需要两倍的重量，因为您还需要 存储梯度。）