CODEWITHSUNDEEP
pythoncomputer-visionpytorchvgg-net
MJ2410
MJ2410

Reputation: 518

How to get a particular layer output of a pretrained VGG16 in pytorch

I am very new to pytorch and I am trying to get the output of the pretrained model VGG16 feature vector in 1*4096 format which is returned by the layers just before the final layer. I found that there are similar features available in keras. Is there any direct command in pytorch for the same?

The code I am using:

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import models
from torch.autograd import Variable
from PIL import Image

image1 = Image.open(r"C:\Users\user\Pictures\user.png")

model = models.vgg16(pretrained=True)

scaler = transforms.Resize((224, 224))
to_tensor = transforms.ToTensor()

img = to_tensor(scaler(image1)).unsqueeze(0)

model(img).shape
model(img)

Upvotes: 0

Views: 4468

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24691

Part of the network responsible for creating features is named... features (not only in VGG, it's like that for most of the pretrained networks inside torchvision).

Just use this field and pass your image like this:

import torch

import torchvision

image = Image.open(r"C:\Users\user\Pictures\user.png")

# Get features part of the network
model = models.vgg16(pretrained=True).features
tensor = transforms.ToTensor()(transforms.Resize((224, 224))(image)).unsqueeze(dim=0)

model(tensor)

EDIT:

To see what happens inside any torchvision model you can check it's source code. For VGG (any), there is a base class at the top of this file.

To get 4096 flattened features, you could operations similar to those defined in forward:

# Tensor from previous code snippet for brevity

x = model.avgpool(tensor)
x = torch.flatten(x, 1)
final_x = model.classifier[0](x) # only first classifier layer

You could also iterate over modules or children up to wherever you want and output the result (or results or however you want)

Upvotes: 3

Related Questions