vgg16提取图片特征向量

关注

原创 2024-04-27 00:14:29 639 阅读 AI 写同款· GEO 优化 ›

Pytorch：利用预训练好的VGG16网络提取图片特征前言这里的提取图片特征特指从VGG网络的最后一个conv层进行提取。虽然下面代码里面给出的是VGG16作为例子，其实也可以用其他的已经经过训练了的神经网络，包括自己训练的。代码模型结构相关基本知识首先说下加载模型，这里用的是torch官方提供的已经训练好的模型，只需要从torchvision模块导入： import torchvision.models as models model =... 阅读详情


from torchvision import transforms
from PIL import Image
import torch
import torchvision.models as models
from torchvision.models import VGG16_Weights

model = models.vgg16(weights=VGG16_Weights.DEFAULT)

# 加载图像
image = Image.open('dog.jpg')

# 数据预处理
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 对图像进行预处理
input_tensor = preprocess(image)
input_batch = input_tensor.unsqueeze(0)

# 将图像输入到模型中
model.eval()
with torch.no_grad():
    features = model(input_batch)
    print(features.data.cpu().numpy()[0].ravel().tolist())
    print(len(features.data.cpu().numpy()[0].ravel().tolist()))

# 获取特征向量
# feature_vector = torch.flatten(features, 1)
# print(len(feature_vector.data.cpu().numpy()[0].ravel().tolist()))