如何使用VGG16提取数据集图像特征？求相关实现方法与参考文章

VGG16的图像分类实现本身就包含特征提取逻辑，不需要找单独的特征提取教程——分类任务用到的卷积层、池化层输出就是通用图像特征，只需要移除末尾的分类输出头，指定对应层的输出即可，具体实现如下：

前置依赖安装 #

• 运行以下命令安装所需依赖：
  pip install torch torchvision pillow numpy

核心实现代码 #

1. 初始化特征提取模型 #

加载ImageNet预训练的VGG16权重，根据需求选择要输出的特征层，切换为评估模式：

import torch
import os
import numpy as np
from PIL import Image
from torchvision import models, transforms

# 优先用GPU加速，没有GPU自动切CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载官方预训练VGG16权重
vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)

# 两种特征输出方案二选一即可
# 方案A：输出最后一个卷积块特征（维度7*7*512，适合细粒度匹配、检索任务）
feature_extractor = vgg16.features.to(device)
# 方案B：输出倒数第二层全连接特征（维度4096，适合聚类、下游分类任务，计算量更小）
# feature_extractor = torch.nn.Sequential(*list(vgg16.children())[:-1]).to(device)

# 必须切换为评估模式，关闭训练态的dropout、批量归一化更新
feature_extractor.eval()

2. 配置图像预处理流程 #

预处理逻辑必须和VGG16预训练时的处理逻辑完全对齐，否则特征分布偏移会导致效果下降：

preprocess = transforms.Compose([
    transforms.Resize((224, 224)),  # VGG16固定输入尺寸为224*224
    transforms.ToTensor(),
    # ImageNet数据集的归一化参数，不要自行替换为自有数据集的统计值
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

3. 批量提取自有数据集特征 #

遍历数据集目录下的所有图像，逐张提取特征后保存为本地文件，后续可直接加载使用：

# 替换为你自己的数据集存放路径
dataset_root = "your_custom_dataset_path"
feature_map = {}

# 关闭梯度计算，减少显存占用、提升推理速度
with torch.no_grad():
    for file_name in os.listdir(dataset_root):
        # 过滤非图像文件
        if file_name.lower().endswith((".jpg", ".jpeg", ".png", ".bmp")):
            img_path = os.path.join(dataset_root, file_name)
            # 读取图像并转为RGB格式，避免通道数异常
            img = Image.open(img_path).convert("RGB")
            # 预处理后增加batch维度，送入对应设备
            input_batch = preprocess(img).unsqueeze(0).to(device)
            # 前向计算得到特征
            raw_feature = feature_extractor(input_batch)
            # 如需一维特征向量，打开下面注释将特征拉平
            flatten_feature = torch.flatten(raw_feature, start_dim=1).squeeze().cpu().numpy()
            feature_map[file_name] = flatten_feature

# 保存提取完成的特征到本地
np.save("vgg16_custom_features.npy", feature_map, allow_pickle=True)

常见注意事项 #

• 推理阶段必须搭配torch.no_grad()和eval()使用，否则输出特征会存在随机误差，还会占用不必要的显存
• 如果自有数据集图像长宽比差异较大，可将预处理中的Resize替换为transforms.Resize(256)+transforms.CenterCrop(224)，避免硬拉伸导致图像内容畸变
• 不需要对自有数据集做单独预训练，直接用ImageNet预训练权重提取的特征就具备通用的语义表达能力，如果后续要做特定下游任务，可以在特征基础上接小的全连接网络微调即可

内容的提问来源于stack exchange，提问作者awatif