How can I use inner module outputs without increasing model size

Phungie: <p>Hi there. I’m gonna keep my explanation short, since I have lots of code to present.</p>
<p>I am trying to develop a new autoencoder architecture, but I am having trouble managing the forward/backward pass size of my model. I’m gonna explain with these test codes:</p>
<p>Imagine we have the two following modules. one returns only the output while the other returns an extra inner output.</p>
<pre><code class="lang-auto">class testNet1(nn.Module):
    def __init__(self, dim_in=3, dim_inner=5, dim_out=3):
        super().__init__()
        self.tube1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)
    
    def forward(self, x):
        x1 = self.tube1(x)
        x2 = self.tube1(x)
        x3 = self.tube1(x)
        x4 = self.tube1(x)

        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.conv_out(x)
        return x, x1
    

class testNet2(nn.Module):
    def __init__(self, dim_in=3, dim_inner=5, dim_out=3):
        super().__init__()
        self.tube1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)
    
    def forward(self, x):
        x1 = self.tube1(x)
        x2 = self.tube1(x)
        x3 = self.tube1(x)
        x4 = self.tube1(x)

        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.conv_out(x)
        return x
</code></pre>
<p>These two modules have the exact same number of parameters, the exact same forward/backward pass size (0.04MB) which you can check for yourself using the <code>torchsummary</code> library.</p>
<p>Now, I made these two models, where one uses the testNet1 and the other uses testNet2. in these examples I haven’t used the inner outputs. But according to my tests, using them or not, doesn’t really affect the F/B pass size.</p>
<pre><code class="lang-auto">class testModule1(nn.Module):
    def __init__(self):
        super().__init__()
        self.net1_1 = testNet1()
        self.net1_2 = testNet1()
        self.net1_3 = testNet1()

    def forward(self, x):
        x, y = self.net1_1(x)
        x, z = self.net1_1(x)
        x, w = self.net1_1(x)
        return x
    

class testModule2(nn.Module):
    def __init__(self):
        super().__init__()
        self.net2_1 = testNet2()
        self.net2_2 = testNet2()
        self.net2_3 = testNet2()

    def forward(self, x):
        x = self.net2_1(x)
        x = self.net2_1(x)
        x = self.net2_1(x)
        return x
</code></pre>
<p>Now these two models have much much different F/B pass sizes. for an arbitrary input size of (3, 64, 64), the module with inner outputs (<code>testModule1</code>) has a F/B pass size of 5.7GB while the second one (<code>testModule2</code>) has a size of 2.44MB.</p>
<p>I have tried detaching the inner output before returning it in hopes of somehow improving the size situation which obviously didn’t help at all.</p>
<p>I want to know how would it be possible to use the inner outputs of a small module in a bigger network without increasing the model size so drastically. I would appreciate any help. Thank you.</p>

ptrblck: <aside class="quote no-group" data-username="Phungie" data-post="1" data-topic="184159">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/phungie/48/57496_2.png" class="avatar"> Phungie:</div>
<blockquote>
<p>Now these two models have much much different F/B pass sizes. for an arbitrary input size of (3, 64, 64), the module with inner outputs (<code>testModule1</code>) has a F/B pass size of 5.7GB while the second one (<code>testModule2</code>) has a size of 2.44MB.</p>
</blockquote>
</aside>
<p>I don’t know what these sizes mean, but assume you have calculated the memory requirements during the forward/backward pass. If so, then these might be expected since you are explicitly storing <code>y</code>, <code>z</code>, and <code>w</code> in the first model, which would need to be kept in memory.<br>
Did you check how large these tensors are and do they match the memory calculation you have performed?</p>

Phungie: <p>Well, these tensors are of shape (5, 64, 64) and I don’t really know what I should be calculating, but 5.7GB does not look right in any shape or form, especially since these values are already created in their respective modules.</p>
<p>I have tried implementing everything in a single class like this:</p>
<pre><code class="lang-auto">class testModule3(nn.Module):
    def __init__(self, dim_in, dim_inner, dim_out):
        super().__init__()
        self.tube1_1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube1_2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube1_3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube1_4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out1 = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)

        self.tube2_1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2_2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2_3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2_4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out2 = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)
        
        self.tube3_1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3_2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3_3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3_4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out3 = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)

    def forward(self, x):
        x1_1 = self.tube1_1(x)
        x1_2 = self.tube1_1(x)
        x1_3 = self.tube1_1(x)
        x1_4 = self.tube1_1(x)

        y = x1_1

        x = torch.cat((x1_1, x1_2, x1_3, x1_4), dim=1)
        x = self.conv_out1(x)
        
        x2_1 = self.tube2_1(x)
        x2_2 = self.tube2_1(x)
        x2_3 = self.tube2_1(x)
        x2_4 = self.tube2_1(x)

        w = x2_1

        x = torch.cat((x2_1, x2_2, x2_3, x2_4), dim=1)
        x = self.conv_out2(x)
        
        x3_1 = self.tube3_1(x)
        x3_2 = self.tube3_1(x)
        x3_3 = self.tube3_1(x)
        x3_4 = self.tube3_1(x)

        z = x3_1

        x = torch.cat((x3_1, x3_2, x3_3, x3_4), dim=1)
        x = self.conv_out3(x)
        return x
</code></pre>
<p>And this works just fine. the F/B pass size of this class is 2.16MB, which is even lower than before.<br>
in this module, I have the <code>y</code>, <code>w</code> and <code>z</code> variables with exactly the same values and I can use them however I choose to. But when I implement my model in a modular way, the model size explodes.</p>
<p>I want to know why this is happening and how I can fix it.</p>

ptrblck: <p>You would need to explain what these sizes mean and how they were calculated.<br>
Checking the max. allocated memory fits my understanding as seen here:</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

class testNet1(nn.Module):
    def __init__(self, dim_in=3, dim_inner=5, dim_out=3):
        super().__init__()
        self.tube1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)
    
    def forward(self, x):
        x1 = self.tube1(x)
        x2 = self.tube1(x)
        x3 = self.tube1(x)
        x4 = self.tube1(x)

        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.conv_out(x)
        return x, x1
    

class testNet2(nn.Module):
    def __init__(self, dim_in=3, dim_inner=5, dim_out=3):
        super().__init__()
        self.tube1 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube3 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube4 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)
        self.tube2 = nn.Conv2d(dim_in, dim_inner, 3, 1, 1)

        self.conv_out = nn.Conv2d(dim_inner*4, dim_out, 3, 1, 1)
    
    def forward(self, x):
        x1 = self.tube1(x)
        x2 = self.tube1(x)
        x3 = self.tube1(x)
        x4 = self.tube1(x)

        x = torch.cat((x1, x2, x3, x4), dim=1)
        x = self.conv_out(x)
        return x

    
print(torch.cuda.max_memory_allocated()/1024**2)
# 0.0
device = "cuda"
h, w = 224, 224

net1 = testNet1().to(device)
x = torch.randn(1, 3, h, w, device=device)
out = net1(x)
print(torch.cuda.max_memory_allocated()/1024**2)
# 9.10302734375

del net1
del x
del out
torch.cuda.reset_max_memory_allocated()
print(torch.cuda.max_memory_allocated()/1024**2)
# 0.0

net2 = testNet2().to(device)
x = torch.randn(1, 3, h, w, device=device)
out = net2(x)
print(torch.cuda.max_memory_allocated()/1024**2)
# 9.10302734375

del net2
del x
del out
torch.cuda.reset_max_memory_allocated()
print(torch.cuda.max_memory_allocated()/1024**2)
# 0.0

class testModule1(nn.Module):
    def __init__(self):
        super().__init__()
        self.net1_1 = testNet1()
        self.net1_2 = testNet1()
        self.net1_3 = testNet1()

    def forward(self, x):
        x, y = self.net1_1(x)
        x, z = self.net1_1(x)
        x, w = self.net1_1(x)
        intermediate_bytes = y.nelement() * y.element_size() + z.nelement() * z.element_size() + w.nelement() * w.element_size()
        print("storing ", intermediate_bytes / 1024**2, " additional MB")
        return x
    

class testModule2(nn.Module):
    def __init__(self):
        super().__init__()
        self.net2_1 = testNet2()
        self.net2_2 = testNet2()
        self.net2_3 = testNet2()

    def forward(self, x):
        x = self.net2_1(x)
        x = self.net2_1(x)
        x = self.net2_1(x)
        return x

mod1 = testModule1().to(device)
x = torch.randn(1, 3, h, w, device=device)
out = mod1(x)
# storing  2.87109375  additional MB
print(torch.cuda.max_memory_allocated()/1024**2)
# 19.83935546875

del mod1
del x
del out
torch.cuda.reset_max_memory_allocated()
print(torch.cuda.max_memory_allocated()/1024**2)
# 0.0

mod2 = testModule2().to(device)
x = torch.randn(1, 3, h, w, device=device)
out = mod2(x)
print(torch.cuda.max_memory_allocated()/1024**2)
# 17.92529296875
</code></pre>