ResNet18 get the output before the last FC (sequential module for FC)

username: S_dB
cooked: <p>I need some help. Basically I want to get the normal output of the ResNet18 that I modified, but also the one from the nn.Linear(512,32) I am able to get the one before nn.Linear(512,32), but not the one coming out, so the one just before the nn.Linear(32,5). I need this to then use an LSTM with these extracted features.</p>
<pre><code class="lang-auto">import torch
from torchvision import models
import torch.nn as nn


class Resnet18(torch.nn.Module):
    def __init__(self):
        super(Resnet18, self).__init__()
        resnet18_pretrained = models.resnet18(pretrained=True)

        self.model = resnet18_pretrained
        self.model.fc =  nn.Sequential(nn.Linear(512, 32),nn.Linear(32, 5))
        self.couches_before_fc = list(self.model.children())[:-1]
        self.resnet_before_fc = nn.Sequential(*self.couches_before_fc)
        self.resnet_before_fc.fc = self.model.fc[0]
        #self.resnet_before_fc = nn.Sequential(*self.couches_before_fc,self.model.fc[0]) #also try this way

    def forward(self, x):

        before_last_fc = self.resnet_before_fc(x)

        x=self.model(x)

        return x, before_last_fc
</code></pre>
<p>This strategy seems logic in my mind but I got the following error:<br>
RuntimeError: CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling <code>cublasSgemm( handle, opa, opb, m, n, k, &amp;alpha, a, lda, b, ldb, &amp;beta, c, ldc)</code> related to the line <strong>before_last_fc = self.resnet_before_fc(x).</strong></p>
<p>By using a hook I am able to get the value of the output before the nn.Linear(32, 5) but I am not able to get it in a tensor form and return it at the end of the forward…</p>
<pre><code class="lang-auto">self.model.fc[0].register_forward_hook(lambda m, input, output: print(output))
</code></pre>
<p>also this additional line make the code crashing after trying to save the model:<br>
torch.save(model, model_out_path)<br>
AttributeError: Can’t pickle local object ‘Resnet18.forward..’</p>
<p>Thanks a lot!</p>

username: ptrblck
cooked: <p>If I execute your code I get a shape mismatch error:</p>
<pre><code class="lang-python">RuntimeError: mat1 and mat2 shapes cannot be multiplied (1024x1 and 512x32)
</code></pre>
<p>since you are not flattening the activation.<br>
This should work:</p>
<pre><code class="lang-python">class Resnet18(torch.nn.Module):
    def __init__(self):
        super(Resnet18, self).__init__()
        resnet18_pretrained = models.resnet18(pretrained=True)

        self.model = resnet18_pretrained
        self.model.fc =  nn.Sequential(nn.Linear(512, 32),nn.Linear(32, 5))
        self.couches_before_fc = list(self.model.children())[:-1]
        self.resnet_before_fc = nn.Sequential(*self.couches_before_fc)
        self.resnet_before_fc.fc = nn.Sequential(
            nn.Flatten(),
            self.model.fc[0]
        )

    def forward(self, x):

        before_last_fc = self.resnet_before_fc(x)

        x=self.model(x)

        return x, before_last_fc
    
model = Resnet18()
x = torch.randn(2, 3, 224, 224)
out = model(x)
</code></pre>
<p>The error is created, since you are wrapping all child modules into an <code>nn.Sequential</code> container, which will be missing all functional API calls from the original <code>forward</code> method (and thus also the flattening operation), so make sure to verify that the <code>nn.Sequential</code> container is indeed calling all needed operations/modules.</p>

username: S_dB
cooked: <p>Thanks a lot this solution is working !</p>