这篇笔记记录的是我用 PyTorch 做特征提取的完整思路和实操过程。所谓“提取特征”简单说就是让训练好的卷积网络把图片“翻译”成一串有意义的数字向量这串向量既保留了图像的关键信息又能直接喂给下游任务——比如图像检索、聚类、迁移学习、风格迁移。我在做这个的时候发现网上很多教程只讲了怎么加载模型、怎么 forward但真正决定特征质量的东西——从哪一层取、怎么处理输出、模型处于什么模式——往往被一笔带过坑全在后面。这篇笔记就把我实际测试过的方案和踩过的坑一起整理出来适合已经会跑 PyTorch 基础代码、想做特征提取但还没理清思路的读者。1. 特征提取到底在提取什么1.1 从“分类”到“特征”的思维转换先想清楚一个问题我们训练好的分类网络它内部到底学到了什么拿 ResNet 举例网络最后接一个全连接层输出 1000 个类别的概率。但真正有意思的不是这个输出而是全连接层之前的那个向量——比如 ResNet18 里avgpool之后会得到一个 512 维的向量。这个向量就是网络对输入图片的“理解”它把一张图片从像素空间压缩到了一个高维语义空间里。在这个空间里内容相似的图片距离更近内容完全不同的图片距离更远。这就是特征提取的核心思路丢开分类头把中间层输出当作图像的高质量向量表示。整个过程没有训练只有推理所以它也叫“无训练特征提取”或者“固定特征提取”。为什么这种事值得单独写一篇因为实际操作中你会发现PyTorch 的模型结构、hook 机制、预处理流程每一步都有细节会影响最终特征的质量。很多人第一次跑通代码后拿特征去做检索效果一塌糊涂问题基本都出在我后面会讲的几个环节上。1.2 预训练模型选择与版本差异特征提取的效果上限很大程度由预训练模型决定。如果模型是随机初始化的那它提取出的特征完全没意义。所以这里必须用 ImageNet 预训练的权重。我在实验里主要用 ResNet18 和 ResNet50。选它们的原因很简单结构成熟PyTorch 官方支持好不同层级输出的特征图尺寸清楚可控。这里有一个很容易踩的版本坑。不同版本的torchvision加载预训练权重的接口不一样# torchvision 0.13 之前的写法 model torchvision.models.resnet18(pretrainedTrue) # torchvision 0.13 及之后的推荐写法 model torchvision.models.resnet18(weightstorchvision.models.ResNet18_Weights.IMAGENET1K_V1)老写法在新版本里会直接报错提示pretrained被移除了。新写法更规范而且以后换成更强权重如IMAGENET1K_V2只需要改一个枚举值。如果你在别人的旧代码里看到pretrainedTrue记得先确认自己装的是哪个版本。1.3 从哪一层提取全局特征与局部特征的取舍这是整个特征提取里最值得纠结的问题。ResNet 系列网络可以简单理解成几个阶段串起来的网络层位置输出形状输入224x224特征含义layer2 之后28x28x128局部纹理、边缘组合layer3 之后14x14x256部件级特征如轮子、眼睛layer4 之后7x7x512高层语义偏全局avgpool 之后512 维向量全图语义向量分类前最后一层fc 之前512 维向量等价于 avgpool 的输出如果你做图像检索或者相似度匹配推荐从avgpool之后取也就是 512 维或 2048 维的向量。这个位置的特征已经经过了全局平均池化空间信息被压缩了留下来的都是针对分类任务最有效的语义信息鲁棒性最好。如果你做的是目标定位、语义分割、注意力可视化这类需要空间信息的任务那就从layer4的输出取拿到的是一个 7x7x512 的张量保留了空间结构方便后续继续做处理。我的建议是先取全局向量跑通流程确认特征有效之后再有针对性地尝试中间层特征。不要一开始就纠结局部特征很多应用用全局向量就够了。2. 环境准备与模型加载2.1 最小依赖与 torchvision 的版本确认先确认环境里有什么。特征提取对计算资源要求不高用 CPU 也能跑但既然都要上 PyTorch 了有 GPU 就尽量用 GPU速度能快十几倍。建议先跑一下命令确认版本python -c import torch; print(torch.__version__) python -c import torchvision; print(torchvision.__version__)常见的问题是两个版本不匹配。尤其是 torchvision 和 torch 的版本要配套否则 import 就直接报错ImportError: libtorch_cuda.so: cannot open shared object file遇到这种错直接去 PyTorch 官网查对应的安装命令用 pip 重装一次最省事。2.2 加载 ResNet 并去掉分类头加载预训练 ResNet 本身很简单但加载之后必须处理一件事把分类头去掉或绕过。因为 fc 层的输出是 1000 类概率那不是我们想要的特征。我常用的做法是直接替换掉最后一层import torch import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc torch.nn.Identity() # 把全连接层替换成恒等映射 model.eval()Identity()的意思就是输入什么就输出什么。这样model(x)的输出直接从 512 维特征开始不会再有分类概率。注意model.eval()这行不能省后面我会专门解释原因。2.3 预处理管线数据进入模型前的形状与归一化预处理这块很多人随手写但直接影响特征提取效果。ImageNet 预训练模型要求输入是三通道 RGB 图片尺寸被 resize 到 224x224并且要做特定的归一化。归一化的均值和标准差不是随便取的是 ImageNet 数据集的统计值from torchvision import transforms transform 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]) ])为什么要先用Resize(256)再CenterCrop(224)因为直接拉伸到 224 会让图片变形影响特征质量。先放大到 256再从中间裁 224 的区域能在保持宽高比的同时让主体内容尽可能位于画面中心。Normalize这一步把 RGB 像素值从[0,1]区间变成以 0 为中心的分布和 ImageNet 预训练时的输入分布对齐。这一步没做对特征表达会明显劣化。3. 特征输出的两种实现方案3.1 方案一注册 forward hookPyTorch 的 hook 机制就是“钩子函数”在某个模块的前向传播之后自动被调用。我们可以在avgpool后挂一个 hook把中间输出存下来。完整实现大概是这样的import torch import torchvision.models as models from torchvision import transforms from PIL import Image model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.eval() # 用字典保存特征 features {} def hook_fn(name): def hook(model, input, output): features[name] output return hook model.avgpool.register_forward_hook(hook_fn(avgpool)) # 加载并预处理图片 img Image.open(test.jpg).convert(RGB) transform 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]) ]) x transform(img).unsqueeze(0) # 加 batch 维度 # 前向传播 with torch.no_grad(): _ model(x) # 提取特征avgpool 的输出形状是 (1, 512, 1, 1) feat features[avgpool].flatten(1) print(feat.shape) # torch.Size([1, 512])这里有几个关键点要讲清楚第一register_forward_hook的 hook 函数接收三个参数模块本身、模块的输入、模块的输出。我们只需要输出。第二hook 注册在avgpool上前向传播时只要走到了avgpool这一层hook 就会被触发。模型后面的层FC还是会继续执行但我们不关心它。第三with torch.no_grad()必须加。这个上下文管理器告诉 PyTorch 不要构建计算图因为这里只推理不求梯度。不加的话显存占用会高很多速度也慢。注意要以实际运行环境为准。3.2 方案二模型截断如果你确定只需要特征根本不想让模型跑到 FC 层那可以直接把模型“截断”。ResNet18 的children()返回的是各层的列表形式Conv2d, BatchNorm2d, ReLU, MaxPool2d, Layer1, Layer2, Layer3, Layer4, AdaptiveAvgPool2d, Linear我们要的是前 9 个模块也就是去掉最后一个 Linearimport torch.nn as nn backbone nn.Sequential(*list(model.children())[:-1]) with torch.no_grad(): feat backbone(x) # 此时 feat 的形状是 (1, 512, 1, 1) feat feat.flatten(1)这个做法的好处是干净利落整个网络根本不会去碰 FC 层也不会触发后面的计算。结构上更直观也好调试。3.3 两种方案的取舍和我的建议试过两种方案之后我的感受是这样如果你是在写一段临时脚本快速验证特征能不能用用hook 方案因为不用改变模型结构随时可以换层。如果你要把特征提取逻辑封装成工具函数或者部署到服务里用截断方案结构更清晰没有 hook 的隐式触发逻辑调用关系一目了然。如果想多提几层特征比如同时拿 layer3 和 layer4 的特征hook 方案更灵活。我的习惯是开发阶段用 hook封装阶段用截断。两条路都不复杂关键是别把两者混着用容易把自己绕晕。4. 完整实操批量提取特征并验证有效性4.1 批量提取的完整代码特征提取跑单张图片验证没问题之后接下来要面对的场景往往是一批图片。比如一个文件夹里有 5000 张商品图你想把每张都转成 512 维向量存成.npy文件或者 pkl供后续检索用。批量提取的完整流程我贴在这里这套代码我在实际项目里验证过改改路径就能直接用import torch import torchvision.models as models from torchvision import transforms from PIL import Image import os import numpy as np from tqdm import tqdm # 1. 加载模型并截断 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) backbone torch.nn.Sequential(*list(model.children())[:-1]) backbone.eval() # 2. 预处理 transform 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]) ]) # 3. 遍历图片 img_dir path/to/images feats [] names [] for img_name in tqdm(sorted(os.listdir(img_dir))): img_path os.path.join(img_dir, img_name) try: img Image.open(img_path).convert(RGB) x transform(img).unsqueeze(0) with torch.no_grad(): feat backbone(x).flatten(1).numpy() feats.append(feat[0]) names.append(img_name) except Exception as e: print(f[跳过] {img_name}: {e}) # 4. 存储 feats np.array(feats) # shape: (N, 512) np.save(feats.npy, feats) with open(names.txt, w) as f: f.write(\n.join(names))有几点注意说明一下如果图片目录里有非图片文件Image.open会报错我的做法是在try-except里跳过并打印日志避免一个坏文件让整个流程中断。backbone(x).flatten(1)是把(1, 512, 1, 1)展平成(1, 512)取[0]就是单张图的特征向量。想用 GPU 的话加上device torch.device(cuda if torch.cuda.is_available() else cpu)然后把模型和输入都.to(device)即可。注意 GPU 模式下.numpy()之前要先.cpu()否则会报错。4.2 特征归一化决定相似度计算成败的关键这一步是绝大多数人第一次做特征提取时最容易忽略的。如果直接把 512 维特征拿去算欧氏距离或余弦相似度数值大小会受图片亮度、对比度影响导致距离不可靠。正确做法是做L2 归一化把每个向量变成单位向量feats feats / np.linalg.norm(feats, axis1, keepdimsTrue)归一化之后任意两个向量的余弦相似度和内积等价而且距离比较的尺度统一了效果会稳很多。还有一种做法是做标准化也就是每个维度减去均值除以标准差。这种白化处理能进一步去除特征维度间的相关性但对大部分普通任务来说L2 归一化已经够用简单有效。4.3 用最朴素的方式验证特征有效性特征提取完后第一件事不是直接上线上系统而是先验证特征到底有没有意义。最简单的验证方式就是做一个 KNN 检索。假设你有一张查询图先提取它的特征并归一化然后和库里所有特征算余弦相似度取 Top-Kquery_feat query_feat / np.linalg.norm(query_feat) similarities feats query_feat # 所有特征行向量与查询向量做点积 topk_idx np.argsort(similarities)[::-1][:10] print(最相似的10张图) for idx in topk_idx: print(names[idx], similarities[idx])如果提取效果正常返回的 Top-K 结果应该在语义上和查询图一致。比如你拿了一张“红色跑车”查询库里应该把其他“跑车”排在前面而不是返回一堆风景照。如果返回结果乱七八糟、毫不相关说明特征质量有问题。这时候回到前面几个环节排查模型权重是否预训练预处理是否归一化是否在 eval 模式特征向量是否做了 L2 归一化绝大概率是其中一环出了问题。5. 常见问题与排查技巧5.1 问题速查表我把实际操作中常见的异常现象、可能原因和解决方案整理成了一张表遇到问题直接对照着排查会比较快现象可能原因解决方案加载模型报pretrained参数错误torchvision 版本过新改用weightsResNet18_Weights.IMAGENET1K_V1首次下载权重非常慢网络原因用export设置镜像源或手动下载权重文件放到缓存目录提取的特征检索效果差预处理未归一化检查是否用了 ImageNet 的 mean/std提取的特征检索效果差忘了model.eval()加上 eval 模式启用 BatchNorm 的统计量输出是(1,512,1,1)而不是(1,512)忘了展平加flatten(1)或.view(1, -1)GPU 模式下.numpy()报错特征还在显存里先.cpu()再转 numpy批量处理到一半卡死/崩溃有损坏图片文件用try-except跳过异常文件并打印日志两张相同图片的特征差异大预处理随机性强比如RandomCrop特征提取阶段只能使用确定性变换不能带随机关Hook 注册后特征字典为空模型没有走到 hook 所在层确认模块名拼写正确打印model验证5.2 高频坑的排查思路上面表里最值得展开的是两个高频坑第一个是model.eval()的重要性。ResNet 里带 BatchNorm 层它在训练和推理两种模式下的行为完全不同。训练模式下BatchNorm 用当前 batch 的均值和方差做归一化推理模式下用训练阶段累积的全局统计量。如果提取特征时忘了model.eval()你得到的结果会受 batch 内其他图片影响——同一张图片单独跑和跑一整个 batch特征都不一样。这就直接破坏了特征的一致性后续做检索必然翻车。第二个是预处理必须保持可复现。很多人在学分类时习惯了transforms.RandomCrop(224)这类增强手段直接抄到特征提取脚本里。后果就是每次跑同一张图特征向量都不同——因为裁剪位置是随机的。特征提取阶段的预处理一定要用CenterCrop、Resize这类确定性的操作保证“同图同特征”。6. 把特征用起来完整项目中的一点体会特征提取做完我通常会顺手做一件小事把模型输出的特征用 T-SNE 降维到二维然后可视化。这一步非常直观能一眼看出提取的特征是不是把不同类别分开了。如果可视化结果是同类聚在一起、不同类分散开那说明特征是有效的可以放心往下游走。维度压缩的代码很轻量from sklearn.manifold import TSNE import matplotlib.pyplot as plt tsne TSNE(n_components2, random_state42) feats_2d tsne.fit_transform(feats) plt.scatter(feats_2d[:, 0], feats_2d[:, 1], s5) plt.savefig(feat_tsne.png)这个可视化技巧我每次做特征提取都会跑一遍比自己对着数字猜要靠谱得多。另外一点经验是存储格式的选择。提取出来的特征通常用np.save存.npy就够了但如果特征规模特别大比如几十万张图、2048 维.npy加载时会占用大量内存更划算的做法是用faiss建索引直接存到磁盘的索引文件里。这样后续做相似度检索也不是先 load 全部特征再算而是直接查索引内存占用和速度都优化了很多。我做特征提取时最深的体会就是这一套流程代码并不复杂难的是理解模型内部在做什么以及每一步操作对最终特征的影响到底在哪里。把“为什么”想明白了代码也就是照着写一遍的事。如果这篇笔记能让你少走我走过的弯路那目的就达到了。