Deleting the last layer in a network resulting in size mismatch!

username: AlJazari
cooked: <p>I am working on a simple architecture, LeNet, with the following architecture:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output
</code></pre>
<p>I want to delete the last layer in the network, fc2. However, when I do so using the approach shown below I get a size mismatch!</p>
<pre><code class="lang-auto">model = Net().to(device)
new_model = nn.Sequential(*list(model.children())[:-1]).to(device)
data, label = next(iter(train_loader))
data, label = data.to(device), label.to(device)
output = new_model(data)
</code></pre>
<p>This is throwing the following error:<br>
<code>RuntimeError: size mismatch, m1: [32 x 36864], m2: [9216 x 128] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:290</code></p>
<p>How can I delete the last layer without getting this error?</p>

username: KFrank
cooked: <p>Hello Al!</p>
<aside class="quote no-group quote-modified" data-username="AlJazari" data-post="1" data-topic="75150" data-full="true">
<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/aljazari/48/24899_2.png" class="avatar"> AlJazari:</div>
<blockquote>
<pre><code class="lang-auto">class Net(nn.Module):
    ...
    def forward(self, x):
        ...
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = self.flatten(x)
        x = self.fc1(x)
        ...
        return output
</code></pre>
<p>I want to delete the last layer in the network, fc2. However, when I do so using the approach shown below I get a size mismatch!</p>
<pre><code class="lang-auto">new_model = nn.Sequential(*list(model.children())[:-1]).to(device)
</code></pre>
</blockquote>
</aside>
<p>When you create <code>new_model</code> it doesn’t know about the detailed<br>
structure of your class <code>Net</code>.  In particular, it doesn’t know about<br>
<code>Net</code>’s <code>forward()</code> function.  (It does know about <code>model.children()</code>).</p>
<p>So I don’t think any of the <code>torch.nn.functional</code> functions used<br>
in <code>model</code> (an instance of your <code>Net</code>) get called when you execute<br>
<code>new_model (data)</code>.  In particular, <code>F.max_pool2d (x, 2)</code> isn’t<br>
getting called.</p>
<aside class="quote no-group">
<blockquote>
<p>This is throwing the following error:<br>
<code>RuntimeError: size mismatch, m1: [32 x 36864], m2: [9216 x 128] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:290</code></p>
</blockquote>
</aside>
<p>Note that the size of the output of <code>self.conv2</code> is <code>(64*3)**2 = 36864</code>,<br>
the second dimension of <code>m1</code> in your size-mismatch error.  The size of<br>
the input to <code>self.fc1</code> is <code>9216</code>, the first dimension of <code>m2</code>.  Had<br>
<code>max_pool2d()</code> been called it would have reduced <code>36864</code> by a factor<br>
of <code>2**2 = 4</code>, yielding <code>9216</code>, eliminating the size-mismatch error.</p>
<aside class="quote no-group">
<blockquote>
<p>How can I delete the last layer without getting this error?</p>
</blockquote>
</aside>
<p>Either build a <code>New_Net</code> class (of which <code>new_model</code> will be an<br>
instance) with an appropriate <code>forward</code> function, or only use <code>Module</code>s<br>
to build your <code>Net</code> class (and make <code>Net</code> an <code>nn.Sequential</code>). That is,<br>
use <code>Module</code>s like <code>nn.ReLU</code> and <code>nn.MaxPool2d</code> instead of functions<br>
like <code>F.relu()</code> and <code>F.max_pool2d()</code>.</p>
<p>Good luck.</p>
<p>K. Frank</p>

username: AlJazari
cooked: <p>Great. I thought so. Thank you so much for the informative response!</p>