Cuda out of memory when modifying model class

username: bryan123
cooked: <p>cfg = {<br>
‘VGG11’: [64, ‘M’, 128, ‘M’, 256, 256, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’],<br>
‘VGG13’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, ‘M’, 512, 512, ‘M’, 512, 512, ‘M’],<br>
‘VGG16’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, 256, ‘M’, 512, 512, 512, ‘M’, 512, 512, 512, ‘M’],<br>
‘VGG19’: [64, 64, ‘M’, 128, 128, ‘M’, 256, 256, 256, 256, ‘M’, 512, 512, 512, 512, ‘M’, 512, 512, 512, 512, ‘M’],<br>
}</p>
<p>class VGG(nn.Module):</p>
<pre><code>def __init__(self, vgg_name):
    super(VGG, self).__init__()
    self.cfg = cfg[vgg_name]
    self.teacher = self._make_layers()
    self.pool = nn.AvgPool2d(kernel_size=1, stride=1)
    self.linear = nn.Linear(512, 3) # change last layer to 3

def forward(self, x):
    teacher_counter = 0
    student_features = []
    feature = x
    for block in self.cfg:
        if block == 'M':
            feature = self.teacher[teacher_counter](feature)
            student_features.append(feature)
            teacher_counter += 1

    out = self.pool(feature)
    out = out.view(out.size(0), -1)
    out = self.linear(feature)

    print(len(student_features))
    return out, student_features


def _make_layers(self):
    layers = []
    teacher = []
    in_channels = 3
    for x in self.cfg:
        if x == 'M':
            teacher.append(nn.Sequential(*layers).to("cuda"))
            layers = []
        else:
            layers += [
                    nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                    nn.BatchNorm2d(x),
                    nn.ReLU(inplace=True)
                ]
            in_channels = x

    return teacher
</code></pre>
<p>The goal is to return the final feature of VGG16 and get features for each ‘M’ in student_features. Could anyone please let me know if anything is wrong with the above implementation? The error message now says “CUDA out of memory” even though it works perfectly with the original VGG-16 under the same batch size</p>

username: ptrblck
cooked: <p>I would recommend to compare your custom <code>VGG</code> model against the <code>torchvision</code> implementation on the CPU first and check where the difference might be coming from first.<br>
You could start with the general model architecture by checking all layers and then you could continue with the intermediate activations etc.</p>

username: bryan123
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for your advice! Could you confirm whether my implementation is correct to get a final feature and feature for every ‘M’ layer/block?</p>

username: ptrblck
cooked: <p>I cannot verify it as your code is not executable, but execute the code and check the returned <code>student_features</code> to see if the stored features are all different.<br>
I would guess that <code>student_features.append(feature)</code> might append a reference to <code>feature</code> and, since it’s overridden in each iteration, it could store the same reference in the <code>list</code>.<br>
In this case, use <code>student_features.append(feature.clone())</code> and it should work.</p>

username: bryan123
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for the input!</p>
<p>It appears that CUDA is out of memory after calling feature = self.teacher [teacher_counter] (feature) 3 times (it should be 5 times). Does calling nn.Sequential many times make the memory full (even if it’s only a small part of the model)? Is there a way to release the memory?</p>

username: ptrblck
cooked: <p>No, the <code>nn.Sequential</code> container will not use more memory than calling layers manually.<br>
Note that you are storing additional features in the <code>forward</code> pass (along with the computation graph), which will of course increase the memory usage, so the OOM might be expected.</p>