How to get output dimension within nn.Sequential()?

alee023: <p>I am writing a flexible VAE that takes in inputs of varying sizes. This is the code that I have so far:</p>
<pre><code class="lang-auto">class Vanilla_3DVAE(nn.Module):
    def __init__(self):
        super(Vanilla_3DVAE, self).__init__()
        # TODO: Parametrize the grid dimensions and channel values
        self.nn_1 = nn.Sequential(
            nn.Conv3d(8, 16, 4, 2, 1 ),        # -4 from dimensions
            nn.ReLU(),
            nn.AvgPool3d(2),            # /2
            nn.BatchNorm3d(16),
            nn.Conv3d(16, 32, 4, 2, 1 ),       # -4
            nn.ReLU(),
            nn.AdaptiveAvgPool3d( 1 ),  # dimension -&gt; 1x1x1
            nn.BatchNorm3d(32))
        self.nn_2 = nn.Sequential(
                                  nn.Linear(32, 16),    #  32x1x1x1
                                  nn.ReLU()
                                )
        self.mu_fc = nn.Linear(16,16)
        self.logvar_fc = nn.Linear(16,16)
        
        self.decoder_nn = nn.Sequential(
                                          nn.ConvTranspose3d(16,8,4,2,0),
                                          nn.BatchNorm3d(8),
                                          nn.ConvTranspose3d(8,16,4,2,0),
                                          nn.BatchNorm3d(16),
                                          nn.ConvTranspose3d(16,8,4,2,1)
                                          )
</code></pre>
<p>Since the inputs will have different dimensions, I’d like to keep track of the dimension size right before nn.AdaptiveAvgPool3d( 1 ), storing this in a variable. I don’t have access to the input dimension size in this .py file. What is the best way to do this?</p>

RaLo4: <p>As far as I know, there are two ways of doing this.</p>
<ol>
<li>If possible I would separate the  <code>AdaptiveAvgPool3d</code> from the <code>self.nn_1</code><br>
Like so:</li>
</ol>
<pre><code class="lang-auto">self.nn_1 = nn.Sequential(
            nn.Conv3d(8, 16, 4, 2, 1 ),        # -4 from dimensions
            nn.ReLU(),
            nn.AvgPool3d(2),            # /2
            nn.BatchNorm3d(16),
            nn.Conv3d(16, 32, 4, 2, 1 ),       # -4
            nn.ReLU())
self.adaptive_pool = nn.Sequential(
            nn.AdaptiveAvgPool3d( 1 ),  # dimension -&gt; 1x1x1
            nn.BatchNorm3d(32))
</code></pre>
<p>And then in your forward:</p>
<pre><code class="lang-auto">def forward(self, x):
        x = self.nn_1(x)
        shape_before_pooling = x.shape
        x = self.adaptive_pool(x)
        [...]
        return x, shape_before_pooling
</code></pre>
<ol start="2">
<li>If for some reason it is not possible for you to separate  <code>AdaptiveAvgPool3d</code> from the <code>self.nn_1</code> you can do this:</li>
</ol>
<pre><code class="lang-auto">shapes = list(range(2))

def get_shapes(self, input, output):
    # input is a tuple of packed inputs
    # output is a Tensor. output.data is the Tensor we are interested
    shapes[0] = input[0].shape
    shapes[1] = output.shape

net = Vanilla_3DVAE()
net.nn_1[6].register_forward_hook(get_shapes)
print(shapes)

</code></pre>

alee023: <p>Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"> Your help is pushing me in the right direction. Since I’m planning to use the variable that holds the pre-pooling shape in the decoder, would something like this work instead? :</p>
<pre><code class="lang-auto">self.nn_1 = nn.Sequential( ... )
self.pre-pool_shape = x.shape
self.nn_adpative_pool = nn.Sequential(...)
self.nn_2 = nn.Sequential(....)

self.mu &amp; logvar_fc layers

self.decoder_nn = nn.Sequential( ... &lt;using variable pre-pool_shape&gt; ...)
</code></pre>
<p><em><strong>Or</strong></em> if using your first method, the dimensions are stored in <code>shape_before_pooling</code> which is in <code>forward()</code>. How can I use this variable in <code>self.decoder_nn</code>? Right now, my forward, encode and decode looks like:</p>
<pre><code class="lang-auto"> def encode(self, x):
        encoded = self.nn_1(x)
        encoded_flat = self.flatten(encoded)
        encoded_flat = self.nn_2(encoded_flat)
        mu = self.mu_fc(encoded_flat)
        logvar = self.logvar_fc(encoded_flat)
        return mu, logvar

    def decode(self, z):
        z_4d = self.restore_dim(z)
        return self.decoder_nn(z_4d)

    def forward(self, x):
        mu,logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        decoded = self.decode(z)
        return (decoded, z, mu, logvar)
</code></pre>
<p>If I make the variable within forward() as you have, create a new parameter in <code>decode()</code> as such: <code>decode(self, z, prePoolShape)</code>, and inside that function, do <code>return self.decoder_nn(z_4d, prePoolShape)</code>, would that work?</p>
<p>Sorry if this seems really beginner. I haven’t had too much experience with VAEs and neural networks in general, and again I really appreciate your help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

RaLo4: <p>Well you can have the variable in the decorator but you want to obviously want to put the value into it in the forward function. Or in your case in the encode function.<br>
This should work:<br>
You have your decorator like I said + the <code>shape_before_pooling</code> variable you wanted</p>
<pre><code class="lang-auto">class Vanilla_3DVAE(nn.Module):
    def __init__(self):
        super(Vanilla_3DVAE, self).__init__()
        self.nn_1 = nn.Sequential(
            nn.Conv3d(8, 16, 4, 2, 1 ),        # -4 from dimensions
            nn.ReLU(),
            nn.AvgPool3d(2),            # /2
            nn.BatchNorm3d(16),
            nn.Conv3d(16, 32, 4, 2, 1 ),       # -4
            nn.ReLU())
        self.adaptive_pool = nn.Sequential(
            nn.AdaptiveAvgPool3d( 1 ),  # dimension -&gt; 1x1x1
            nn.BatchNorm3d(32))
        self.nn_2 = nn.Sequential(....)
        self.mu_fc = (...)
        self.logvar_fc = (...)
        self.decoder_nn = (...)

        self.pre_pool_shape = None        #this is new
</code></pre>
<p>And then have you <code>forward</code>, <code>encode</code> and <code>decode</code> looks like this:</p>
<pre><code class="lang-auto">    def encode(self, x):
        encoded = self.nn_1(x)
        self.shape_before_pooling = encoded.shape        #this is new
        adap_pooled = self.adaptive_pool(encoded)        #this is new
        encoded_flat = self.flatten(adap_pooled)        #this is different
        encoded_flat = self.nn_2(encoded_flat)
        mu = self.mu_fc(encoded_flat)
        logvar = self.logvar_fc(encoded_flat)
        return mu, logvar

    def decode(self, z):
        z_4d = self.restore_dim(z)
        return self.decoder_nn(z_4d)

    def forward(self, x):
        mu,logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        decoded = self.decode(z)
        return (decoded, z, mu, logvar)
</code></pre>
<p>I am fairly certain this should work.</p>