如何在pytorch

时间:2019-11-23 13:05:24

标签: python computer-vision pytorch vgg-net

我对pytorch非常陌生,我正在尝试获取1 * 4096格式的预训练模型VGG16特征向量的输出,该向量由最后一层之前的各层返回。我发现keras中有类似的功能。 pytorch中是否有任何直接命令可以实现同样的目的?

我正在使用的代码:

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)

1 个答案:

答案 0 :(得分:1)

负责创建features的网络的一部分被命名为... features(不仅在VGG中,而且在torchvision内部的大多数预训练网络中也是如此)。

只需使用此字段并像这样传递您的图像即可:

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)

编辑:

要查看任何torchvision模型内发生的情况,可以检查其源代码。对于VGG(任意),this file的顶部有一个基类。

要获得4096的展平特征,您可以进行类似于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

您还可以遍历moduleschildren到您想要的任何地方并输出结果(或结果或您想要的结果)