How could I remove certain layer in torchvision.models

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