Getting and using output dimensions in a neural network by code (not manual calculation)

alee023: <p>I’m working on a VAE that can take in a flexible-dimension input, so I’ll like to store the pre-adaptive pooling dimensions in a variable. This is so that I can use the dimension in the decoder to help get back to the original size. I’m struggling to find a place to create this variable correctly. This is some of my code:</p>
<pre><code>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(
        &lt; layers &gt;
        nn.AdaptiveAvgPool3d( 1 ),  # 4 -&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>Following the suggestion of someone who helped, I separated the adaptive pooling  layer from <code>self.nn_1</code>, and stored the pre-pooling shape in a variable called <code>dim</code> in the <code>encode()</code> function which <code>forward()</code> calls upon. My new forward() code is:</p>
<pre><code>def forward(self, x):
    mu,logvar, dim = self.encode(x)
    z = self.reparametrize(mu, logvar)
    decoded = self.decode(z)
    return (decoded, z, mu, logvar)
</code></pre>
<p>My new issue is - how can I use <code>dim</code> in <code>self.decoder_nn</code>? The <code>__init__</code> cannot change - I can only take in <code>self</code> as the parameter. I assume that this will cause issues in using <code>dim</code> in <code>self.decoder_nn</code>.</p>
<p>Something else that I’ve thought of doing was</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>This had issues because no variable <code>x</code> exists in the model. If using this method, how do I declare <code>self.pre-pool_shape</code>?</p>
<p>Thank you so much <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

arya47: <p>I would suggest creating a method which takes as input the desired shape and returns a Sequential layer.<br>
You can call this in  forward.</p>

alee023: <p>I don’t have too much experience with NNs so I apologize if I’m asking simple questions <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=10" title=":frowning:" class="emoji" alt=":frowning:">. Suppose that <code>forward()</code> returns a Sequential layer. How can I use that with the <code>self.decoder_nn</code> so that the decoder can actually use the dimension?</p>
<p>I wasn’t sure if tagging you, <a class="mention" href="/u/ptrblck">@ptrblck</a>, is okay (let me know), but do you have an idea on how to approach the issue in the post? Thank you.</p>

ptrblck: <p>I wouldn’t recommend to recreate new modules in the <code>forward</code>, as they would be reinitialized in each forward pass and thus not trained.<br>
I’m not sure what the actual use case is and where these shapes would be used. Are you trying to pass them to the transposed conv layers?</p>

alee023: <p>Thanks for replying :). Yes I’d like to somehow pass the dimension back into <code> self.decoder_nn = nn.Sequential( ... )</code> and use the dimension to help get to the original size. The code that I had to get the dimension in in encode():</p>
<pre><code class="lang-auto">class Vanilla_3DVAE(nn.Module):
    def __init__(self):
        super(Vanilla_3DVAE, self).__init__()
        self.nn_1 = nn.Sequential( ... ) # same as before except missing last 2 layers, moved below
        self.nn_adaptivePool = nn.Sequential(
            nn.AdaptiveAvgPool3d( 1 ),  # 4 -&gt; 1x1x1
            nn.BatchNorm3d(32))
        self.nn_2 = nn.Sequential( ... ) # same as before
        self.mu_fc = nn.Linear(16,16)
        self.logvar_fc = nn.Linear(16,16)
        
        self.decoder_nn = nn.Sequential( ... ) # same as before but would like to pass the dimension here

    def encode(self, x):
        encoded = self.nn_1(x)
        prePoolShape = list( encoded.size())
        prePoolDim = prePoolShape[ 2 ] # &lt;&lt;&lt;&lt;&lt; dimension &lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;
        # print( prePoolDim ) 
        encoded1= self.nn_adaptivePool( encoded )
        encoded_flat = self.flatten(encoded1 )
        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>

ptrblck: <p>I’m still unsure how <code>prePoolDim</code> would be used, but I would probably write custom modules to get the needed flexibility. If you want to reuse the sequential pattern, you could derive a custom <code>nn.Sequential</code> module and allow multiple input arguments, which would then be passed to the submodules.</p>

alee023: <p>I was planning to make my decoder_nn like this:</p>
<pre><code class="lang-auto">self.decoder_nn = nn.Sequential(
                                      nn.ConvTranspose3d(16,8, prePoolDim,2,0), # &lt;&lt;&lt;&lt;&lt; dimension &lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;
                                      nn.BatchNorm3d(8),
                                      nn.ConvTranspose3d(8,16,4,2,0),
                                      nn.BatchNorm3d(16),
                                      nn.ConvTranspose3d(16,8,4,2,1)
                                      )
</code></pre>
<p>The output dimensions from this first <code>ConvTranspose3d</code> layer should essentially be <code>8xDimxDimxDim</code>.</p>
<p>Since I’m only allowed <code>self</code> in <code>def __init__(self)</code>, would the custom <code>nn.Sequential</code> module idea with multiple input args work? Also, as mentioned at the bottom of the post, I had this idea but wasn’t sure if it would work:</p>
<pre><code class="lang-auto">self.nn_1 = nn.Sequential( ... )
self.prePoolShape = list( x.size()) # &lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt;
self.prePoolDim = prePoolShape[ 2 ] # &lt;&lt;&lt;&lt;&lt;&lt;&lt;&lt; 
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 prePoolDim&gt; ...)
</code></pre>
<p>Would this idea work? I had issues trying this because no variable <code>x</code> exists in the model but I wasn’t sure what to put in place of <code>x</code>.</p>

ptrblck: <p>No, this approach won’t work, as you are trying to use a dynamic variable (<code>prePoolDim</code> depends on the actual input size) as the kernel size.<br>
The <code>nn.ConvTranspose3d</code> layer will use the value to initialize the kernel in the desired shape, so changing the kernel shape afterwards would work, if you are directly manipulating the kernel in the <code>forward</code> method. However, to do so, you would need to come up with a strategy how the additional filter values should be filled or how the filter should be sliced.</p>
<p>Transposed conv layers accept an <code>output_size</code> argument to specify the desired spatial output size in case multiple values would be valid.<br>
In case the desired output size is fixed, you could write a wrapper to set it, and then use this custom module in the <code>nn.Sequential</code> container.<br>
However, since it’s dynamic I would avoid using <code>nn.Sequential</code> for the lack of flexibility.</p>

alee023: <p>Thanks for getting back! In my case, I don’t want my desired output to be fixed – I’d like it to match the dimensions of the input, which was why I wanted to keep track of the dimension before the adaptive pooling layer.</p>
<p>Keeping the arguments in <code>__init__()</code> the same (<code>self</code>) – so no adding an argument for input dimension or anything – what do  you suggest is the best way to write a flexible VAE that will output something with the same dimensions? Sorry if this is a really general question</p>

ptrblck: <p>I would claim it depends a bit on the expected input shapes.<br>
I.e. in case you are using input shapes as powers of 2, reduce the spatial size by 2, and increase it later by the same factor in the transposed conv layers, I don’t think you would be running into a lot of issues.<br>
However, for arbitrary input shapes, the upsampling layers could cause ambiguous output shapes, so that you could then pass the desired output shape as an argument to the <code>forward</code> method of the transposed conv layers.<br>
To do so I would remove the <code>nn.Sequential</code> container and use the layers manually in the <code>forward</code> of the main model, which would allow you to pass the <code>output_size</code> argument.</p>