Why is an OOM happening on my model init()?
A single line in my model, tr.nn.Linear(hw_flat * num_filters*8, num_fc), is causing an OOM error on initialization of the model. Commenting it out removes the memory issue.

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)


This is strange, as hw_flat = 96*96 = 9216, and num_filters*8 = 512, so hw_flat * num_filters = 4718592, which is the number of parameters in that layer. I have confirmed this calculation as changing the layer to tr.nn.Linear(4718592, num_fc) results in the same output.

To me this makes no sense as dtype=float32, so the expected size of this would be 32*4718592 = 150,994,944 bytes. This is equivalent to about 150mb.

Error message is:

Traceback (most recent call last):
  File "main.py", line 116, in &lt;module&gt;
    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


I am only running batch sizes of 1 as well (not that that affects this error), with overall input shape to the network being (1, 3, 1536, 1536), and shape after flatten layer being (1, 4718592).

Why is this happening?

Your linear layer is quite large - it does, in fact, need at least 18GB of memory. (Your estimate is off for two reasons: (1) a float32 takes 4 bytes of memory, not 32, and (2) you didn't multiply by the output size.)

From the PyTorch documentation FAQs:


  Don’t use linear layers that are too large. A linear layer nn.Linear(m, n) uses O(n*m)
  memory: that is to say, the memory requirements of the weights scales quadratically with
  the number of features. It is very easy to blow through your memory this way (and 
  remember that you will need at least twice the size of the weights, since you also need
  to store the gradients.)

