Inferring shape via flatten operator

domluna: <p>Is there a flatten-like operator to calculate the shape of a layer output. An example would be transitioning from a conv layer to linear layer. In all the examples I’ve seen thus far this seems to be manually calculated, ex:</p>
<pre><code class="lang-python">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        return F.log_softmax(x)
</code></pre>
<p>What would be idiomatic torch to calculate 320?</p>

fmassa: <p>You can use <code>-1</code> in view so that the remaining dimension is automatically calculated, but instead of using it in the first dimension, which you know and is the batch size, you use it the the other one.<br>
For example</p>
<pre><code class="lang-auto">bs = 5
x = torch.rand(bs, 3, 224, 224)
x = x.view(x.size(0), -1)
</code></pre>

domluna: <p>Thanks, that works for the forward method but I’m more concerned with the network definition, <code>self.fc1 = nn.Linear(320, 50)</code>. How do I calculate the 320 there using torch?</p>

fmassa: <p>Hum, I’m afraid you can’t calculate that in <code>__init__</code> without prior knowledge of the input shape.<br>
You could imagine passing an input shape as argument to the <code>__init__</code>. In this situation, you can infer the shape by performing a forward pass over the convolutional blocks. Something like</p>
<pre><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self, input_shape=(1, 28, 28)):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()

        n_size = self._get_conv_output(input_shape)
        
        self.fc1 = nn.Linear(n_size, 50)
        self.fc2 = nn.Linear(50, 10)

    # generate input sample and forward to get shape
    def _get_conv_output(self, shape):
        bs = 1
        input = Variable(torch.rand(bs, *shape))
        output_feat = self._forward_features(input)
        n_size = output_feat.data.view(bs, -1).size(1)
        return n_size

    def _forward_features(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        return x

    def forward(self, x):
        x = self._forward_features(x)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = F.relu(self.fc2(x))
        return F.log_softmax(x)
</code></pre>

ypxie: <p>you can try functional api which is truly dynamic.</p>

ncullen93: <p>Hm yeah I did something similar. You definitely have to know the input size. The only other option is to write a general function which calculates the shape using conv rules without having to actually run the graph…</p>
<pre><code class="lang-python">class Network(nn.Module):
    
    def __init__(self, input_size=(3,32,32)):
        super(Network, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3,32,3),
            nn.ReLU(),
            nn.Conv2d(32,32,3),
            nn.ReLU(),
            nn.MaxPool2d((3,3))
        )
        self.flat_fts = self.get_flat_fts(input_size, self.features)
        self.classifier = nn.Sequential(
            nn.Linear(self.flat_fts, 100),
            nn.Dropout(p=0.2),
            nn.ReLU(),
            nn.Linear(100,10),
            nn.LogSoftmax()
        )
    
    def get_flat_fts(self, in_size, fts):
        f = fts(Variable(torch.ones(1,*in_size)))
        return int(np.prod(f.size()[1:]))
    
    def forward(self, x):
        fts = self.features(x)
        flat_fts = fts.view(-1, self.flat_fts)
        out = self.classifier(flat_fts)
        return out
</code></pre>

fmassa: <p><a class="mention" href="/u/ncullen93">@ncullen93</a> when you do <code>fts(Variable(torch.ones(1,*in_size)))</code>, you are actually performing a forward pass on your network.</p>

ncullen93: <p><a class="mention" href="/u/fmassa">@fmassa</a> right… That’s what you do in your example as well, I think? I’m saying it would be nice to not have to do that.</p>

fmassa: <p>Yes, that’s what I do in my example.</p>
<p>Note that you can add a global operation (like global max/average pooling) just before your view layer, so that you know precisely the number of inputs that the linear layer will receive (as you can see in the <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py#L149">resnet model definition</a>, where the kernel size for the pooling can be computed on the fly using the functional interface).</p>
<p>We could eventually add another method to each <code>Function</code> that, given an input shape and a set of parameters, returns an output shape, but is it really worth it?</p>
<p><a class="mention" href="/u/apaszke">@apaszke</a> what do you think?</p>

apaszke: <p><a class="mention" href="/u/fmassa">@fmassa</a> we’ll need that for lazy execution, but I don’t know what the API’s going to look like yet.</p>

apaszke: <p>Actually no, we don’t need that. I think it’s too much work and maintenance for very small gains. You probably want to do that extra forward once.</p>

abhigenie92: <p>This should be a feature, seems weird to do a forward pass.</p>

xueliang_liu: <p>Keras and Lasagne can do it very easily.</p>

parichehr: <p>Can you please tell me what is bs here?</p>

lkhphuc: <p>bs is for batch_size, which is then below refered to as x.size(0).</p>

Xingdong_Zuo: <p><a class="mention" href="/u/apaszke">@apaszke</a> If I got it correctly, the additional forward pass does not have any effect on the gradient computation in later training phase right ? And will the memory be release automatically (since we do not do backward pass, the node will not be diminished ) ?</p>

mlearnx: <p>guys any update about this feature?</p>
<p>I also wonder to know the answer to Xingdong_Zuo question.</p>
<p>Thanks!</p>

iacob: <p>As of 1.8, PyTorch now has <a href="https://pytorch.org/docs/stable/generated/torch.nn.LazyLinear.html" rel="noopener nofollow ugc"><code>LazyLinear</code></a> which <a href="https://stackoverflow.com/a/68284577/9067615" rel="noopener nofollow ugc">infers the input dimension</a>:</p>
<blockquote>
<p>A <code>torch.nn.Linear</code> module where <em>in_features</em> is inferred.</p>
</blockquote>