Why does the size of Forward/backward pass differ when using a single class for a model and partitioning the model using different classes and later accumulating it?

Jimut123: <p>I am creating two models using the same stack of layers using a single functional class and another with discrete classes, later accumulating into one class. Even model.summary() gives the same number of parameters for both models; why does the size of Forward/backward pass (618.00 MB vs. 729.38 MB) differ by some amount using torch summary? Even the training results are inconsistent, using the same seed for both the models and every hyper-parameter the same.</p>

ptrblck: <p>Could you post both model architectures and did you measure the actual memory usage yourself or only via <code>torchsummary</code>?</p>

Jimut123: <p>I measured the memory using torchsummary. I am showing an example (since the actual model was too big) here (Note the difference is small here); I want to know what exactly is creating the difference (56.39 vs 53.02); also, the results were inconsistent in the model that I was using during training by just using the class-based method even when using proper seed in everything…</p>
<pre><code class="lang-auto">! pip install -q torchview
! pip install -q -U graphviz



import torch
import torch.nn as nn
from torch import Tensor
import torch.nn.functional as F
from torchsummary import summary
from torchsummary import summary
from torchview import draw_graph
from torchvision.models import resnet18, GoogLeNet, densenet, vit_b_16
# when running on VSCode run the below command
# svg format on vscode does not give desired result
import graphviz
graphviz.set_jupyter_format('png')
import matplotlib.pyplot as plt



use_cuda = torch.cuda.is_available()
print('use_cuda: {}'.format(use_cuda))
device = torch.device("cuda" if use_cuda else "cpu")
print("Device to be used : ",device)


import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class Test_Model_1(nn.Module):

    def __init__(self, inp_ch):
        super().__init__()

        self.conv1 = nn.Conv2d(inp_ch, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv3 = nn.Conv2d(128, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 1, 3, padding=1)

        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv_before_flatten = nn.Conv2d(128, 1, 1)
        self.dropout = nn.Dropout(p=0.2)
        self.maxpool = nn.MaxPool2d(2)

        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(3072, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        conv1 = self.conv1(x)
        x = self.maxpool(conv1)

        conv2 = self.conv2(x)
        x = self.maxpool(conv2)

        conv3 = self.upsample(self.conv3(x))
        conv4 = self.upsample(self.conv4(conv3))

        conv_flatten = self.conv_before_flatten(x)
        flatten = self.flatten(conv_flatten)
        fcns = self.fc3(self.fc2(self.fc1(flatten)))

        out = self.sigmoid(fcns)
        return out, conv4


model1 = Test_Model_1(inp_ch=3)
model1 = model.to(device)

summary(model1, input_size=(3, 256, 192))



##### Here is the assembled class based implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary


class Body1(nn.Module):
    def __init__(self, inp_ch):
        super().__init__()
        self.conv1 = nn.Conv2d(inp_ch, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
        self.maxpool = nn.MaxPool2d(2)

    def forward(self, x):
        conv1 = self.conv1(x)
        x = self.maxpool(conv1)

        conv2 = self.conv2(x)
        x = self.maxpool(conv2)
        return x


class Head1(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv3 = nn.Conv2d(128, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 1, 3, padding=1)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

    def forward(self, x):
        conv3 = self.upsample(self.conv3(x))
        conv4 = self.upsample(self.conv4(conv3))
        return conv4


class Head2(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv_before_flatten = nn.Conv2d(128, 1, 1)
        self.dropout = nn.Dropout(p=0.2)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(3072, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        conv_flatten = self.conv_before_flatten(x)
        flatten = self.flatten(conv_flatten)
        fcns = self.fc3(self.fc2(self.fc1(flatten)))
        out = self.sigmoid(fcns)
        return out



class Test_Model_2(nn.Module):
    def __init__(self, inp_ch):
        super().__init__()

        self.body1 = Body1(inp_ch)
        self.head1 = Head1()
        self.head2 = Head2()

    def forward(self, x):
        body_1 = self.body1(x)
        head_1 = self.head1(body_1)
        head_2 = self.head2(body_1)
        return head_1, head_2



model2 = Test_Model_2(inp_ch=3)
model2 = model.to(device)

summary(model2, input_size=(3, 256, 192))
</code></pre>
<p>model1.summary() gives the following:</p>
<pre><code class="lang-auto">----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 256, 192]           1,792
         MaxPool2d-2          [-1, 64, 128, 96]               0
            Conv2d-3         [-1, 128, 128, 96]          73,856
         MaxPool2d-4          [-1, 128, 64, 48]               0
            Conv2d-5           [-1, 64, 64, 48]          73,792
          Upsample-6          [-1, 64, 128, 96]               0
            Conv2d-7           [-1, 1, 128, 96]             577
          Upsample-8          [-1, 1, 256, 192]               0
            Conv2d-9            [-1, 1, 64, 48]             129
          Flatten-10                 [-1, 3072]               0
           Linear-11                  [-1, 128]         393,344
           Linear-12                   [-1, 64]           8,256
           Linear-13                   [-1, 10]             650
          Sigmoid-14                   [-1, 10]               0
================================================================
Total params: 552,396
Trainable params: 552,396
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.56
Forward/backward pass size (MB): 53.02
Params size (MB): 2.11
Estimated Total Size (MB): 55.69
----------------------------------------------------------------
</code></pre>
<p>While, model2.summary() gives:</p>
<pre><code class="lang-auto">----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 256, 192]           1,792
         MaxPool2d-2          [-1, 64, 128, 96]               0
            Conv2d-3         [-1, 128, 128, 96]          73,856
         MaxPool2d-4          [-1, 128, 64, 48]               0
             Body1-5          [-1, 128, 64, 48]               0
            Conv2d-6           [-1, 64, 64, 48]          73,792
          Upsample-7          [-1, 64, 128, 96]               0
            Conv2d-8           [-1, 1, 128, 96]             577
          Upsample-9          [-1, 1, 256, 192]               0
            Head1-10          [-1, 1, 256, 192]               0
           Conv2d-11            [-1, 1, 64, 48]             129
          Flatten-12                 [-1, 3072]               0
           Linear-13                  [-1, 128]         393,344
           Linear-14                   [-1, 64]           8,256
           Linear-15                   [-1, 10]             650
          Sigmoid-16                   [-1, 10]               0
            Head2-17                   [-1, 10]               0
================================================================
Total params: 552,396
Trainable params: 552,396
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.56
Forward/backward pass size (MB): 56.39
Params size (MB): 2.11
Estimated Total Size (MB): 59.06
----------------------------------------------------------------
</code></pre>

Jimut123: <p>Even the layers are the same when seeing it visually…</p>
<p>So, why does class-based implementation increase the Forward/backward pass size?</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/3/537d36351d189a54facd15143ba78dccb95b94a8.png" data-download-href="https://discuss.pytorch.org/uploads/default/537d36351d189a54facd15143ba78dccb95b94a8" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/3/537d36351d189a54facd15143ba78dccb95b94a8_2_269x500.png" alt="image" data-base62-sha1="bUzTuiNDknloRgqBaBabPLlH9tC" width="269" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/3/537d36351d189a54facd15143ba78dccb95b94a8_2_269x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/3/537d36351d189a54facd15143ba78dccb95b94a8_2_403x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/3/537d36351d189a54facd15143ba78dccb95b94a8_2_538x1000.png 2x" data-dominant-color="DFF1DE"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">621×1152 55.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ae2ea564319ce872eb3f72e37562b83b3c43728.png" data-download-href="https://discuss.pytorch.org/uploads/default/8ae2ea564319ce872eb3f72e37562b83b3c43728" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ae2ea564319ce872eb3f72e37562b83b3c43728_2_284x500.png" alt="image" data-base62-sha1="jOE3oX9q0mWLfLgvG8YbjMSFqXu" width="284" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ae2ea564319ce872eb3f72e37562b83b3c43728_2_284x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ae2ea564319ce872eb3f72e37562b83b3c43728_2_426x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ae2ea564319ce872eb3f72e37562b83b3c43728_2_568x1000.png 2x" data-dominant-color="E1F1E0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">688×1209 66.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The primary motivation for creating the class-based discrete blocks is I want to set require_grad=False/True in the entire class (or block, like head1 or head2 like a switch, and set require_grad=True, in different stages of training), and I don’t have to do that iteratively in the first implementation. <a class="mention" href="/u/ptrblck">@ptrblck</a> Any suggestions on how to set the block’s required grad=True and False, in the first model implementation, without going through each layer’s bias and weights individually in the first (non-discrete class-based) implementation?</p>

ptrblck: <p>The difference in memory is most likely coming from the different allocations of intermediate activations in the <code>forward</code> method. You would have to make sure to assign the same number of outputs to variables or reassign them.<br>
Seeding can be tricky as you would need to guarantee the exact same order or calls is made to the pseudorandom number generator. I’m not on my workstation now but you could debug it already by checking the parameters layer by layer.</p>

wangh09: <p>Having the same issue, did you figure out why it is larger when using more nn.Modules? In my case, I created the following network:</p>
<pre><code class="lang-auto">def __init__():
    layer0 = SomeCustomizedModule1
    layer1 = SomeCustimizedModule2

def foward(self, x):
    x = layer0(x)
    x = layer1(x)
    return x
</code></pre>
<p>and if I move layer1 into SomeCustomizedModule1, the forward/backward passing size dramatically drops. Really weird : (</p>

ptrblck: <p>Would you post the definitions of all used classes?</p>

wangh09: <p>Hi, I figured out why the sizes differ. The forward/backward passing size is calculated by adding all the modules’ output sizes together.</p>
<p>Say we have 2 convs in our network:</p>
<pre><code class="lang-auto">def __init__(self):
  self.conv1 = nn.Conv2d(3, 64, 3, padding=1) 
  self.conv2 = nn.Conv2d(64, 128, 3, padding=1)

def forward(self, x):
   x = self.conv1(x)  --&gt; output1
   x = self.conv2(x)  --&gt; output2
   return x
</code></pre>
<p>The f/bp size would be output1 size + output2 size.<br>
If we wrap them into one nn.Module, for example:</p>
<pre><code class="lang-auto">class XXX(nn.Module):
def __init__(self):
   self.conv1 = ...   
   self.conv2 = ...   
def forward(self, x):
   x = self.conv1(x) --&gt; output1
   x = self.conv2(x) --&gt; output2
   return x
-------------------
def __init__(self):
    self.module = XXX()

def forward(self, x):
   x = self.module(x) --&gt; output3
   return x
</code></pre>
<p>The final f/bp size would be output1 + output2 + output3!<br>
There would be a redundant output for each wrapping Module!<br>
Would this be a torchsummary bug?</p>

Jimut123: <p>No idea; I couldn’t figure this thing out. Even the training had a minor difference in optimization when using proper seeding.</p>

ptrblck: <p>I don’t fully understand the explanation as calling <code>self.module(x)</code> will just return the same output from <code>self.conv2</code>, which you previously tagged as <code>output2</code>.<br>
Also, what does <code>final f/bp size would be output1 + output2 + output3</code> mean in this context? Shapes are not added together so I guess you are using a different syntax to describe intermediate shapes?</p>