Given vgg16, how to remove pool5 layer and all the classify layers? And how to add new layers to that while keeping the pretrained weights?

username: brisker
cooked: <p><code>pretrained_model = torchvision.models.vgg16(pretrained=True)</code></p>
<ol>
<li>given vgg16, how to remove pool5 layer and all the classify layers?(only keep the conov1_1 to relu5_3 layers)</li>
<li>how to add new layers ,and with new layers with random initialize parameters and the conv1_1 ~ relu5_3 keeping the pretrained weights?</li>
</ol>

username: cyyyyc123
cooked: <p>You can do something like this:</p>
<pre><code class="lang-auto">class EncoderCNN(nn.Module):

    def __init__(self):
        super(EncoderCNN, self).__init__()
        self.vgg = models.vgg16()
        self.vgg.load_state_dict(torch.load(vgg_checkpoint))
        self.vgg.classifier = nn.Sequential(
            *(self.vgg.classifier[i] for i in range(6)))

    def forward(self, images):
        return self.vgg(images)
</code></pre>
<p>VGG net can be viewd as the combination of two sub-nets: feature extracting net and classifying net, and each of them is a nn.Sequential module. I just remove the last fc layer in classifying net by constructing a new nn.Sequential module with the pretrained parameters.</p>
<p>For your requirement, I guess you can do it like this:</p>
<pre><code class="lang-auto">    def __init__(self):
        super(EncoderCNN, self).__init__()
        self.vgg = models.vgg16()
        self.vgg.load_state_dict(torch.load(vgg_checkpoint))
        self.vgg.features = nn.Sequential(
            *(self.vgg.features[i] for i in range(30))

    def forward(self, images):
        return self.vgg.feature(images)
</code></pre>
<p>Sorry that this method looks so ugly. The nn.Sequential object does not support slice so I have to construct a new Sequential by list comprehension.</p>

username: brisker
cooked: <p>I guess the following codes do the same thing, right?<br>
<code>pretrained_model = torchvision.models.vgg16(pretrained=True)</code><br>
<code>modified_pretrained = nn.Sequential(*list(pretrained_model.features.children())[:-1]) # to relu5_3</code></p>

username: cyyyyc123
cooked: <p>Yes! Your implementation looks much more elegant, thank you!</p>

username: Rakshit_Kothari
cooked: <p>For some reason, this method doesn’t seem to work for me. I receive the following error:</p>
<blockquote>
<p>RuntimeError: size mismatch, m1: [3584 x 7], m2: [25088 x 4096] at /pytorch/torch/lib/TH/generic/THTensorMath.c:1293</p>
</blockquote>
<p>Did you face a similar issue? I don’t know why this should happen because after all, it’s using all the layers in a sequential fashion.</p>

username: yashasvi97
cooked: <p>Hi, how can I use this method for pretrained resnet18? Currently the same method gives error as - ‘ResNet’ object has no attribute ‘features’.</p>

username: klory
cooked: <p>Yeah, I got the same result…</p>

username: Aryan_Asadian
cooked: <p>This is because there is no module in the pre-trained model named as <strong>features</strong>. “features” is one of the modules of VGG(the initial example of this thread). To see the module names, just simple print your model.</p>
<pre><code class="lang-auto">my_model = models.ResNet18(pre_trained = False, progressed =True)
print(my_model)
</code></pre>