How could I remove certain layer in torchvision.models

1

I would like to use the torchvision models without some of their layers. For example, I may need to use the model without its last nn.Linear layer, but I stiil need the remain part to work as the original torchvision.models. How could I do it ?

2

Oneway to do it would be pull the source code from pytorch and change the layers as you want.

Eg: https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
pull the resent and change resnet as you need

1 Like
3

The code given below is to download pre-trained Resnet152 and use it till second last layer.

import torch
import torch.nn as nn
import torchvision.models as models
from torch.autograd import Variable

# Pretrained resnet152 model till second last layer as feature extraction.
resnet152 = models.resnet152(pretrained=True)
modules=list(resnet152.children())[:-1]
resnet152=nn.Sequential(*modules)
for p in resnet152.parameters():
    p.requires_grad = False

# Get resnet features for random image_
img = torch.Tensor(3, 224, 224).normal_() # random image
img = torch.unsqueeze(img, 0)  # Add dimension 0 to tensor  
img_var = Variable(img) # assign it to a variable
features_var = resnet152(img_var) # get the output from the last hidden layer of the pretrained resnet
features = features_var.data # get the tensor out of the variable

For more variants of pretrained models go to Models.

7 Likes
5

I tried and got error:

TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not tuple.

Please help, thanks!