简介：面向AI深度学习与图像检索场景，这套基于视觉变换器（ViT）的代码提供一个完整且轻量的相似图像检索实现，适合开发者快速学习图像特征提取与向量匹配技术，也可直接用于以图搜图、重复图片检测等实际任务。资源包共2个文件，包含一个Python脚本与一份Markdown说明文档，压缩包仅2KB，结构紧凑但功能齐全；脚本负责全部核心逻辑，文档提供配置与使用指引。已有663人学习浏览，说明其实践参考价值不错。脚本完整覆盖图像预处理（调整大小、中心裁剪、转张量、归一化）、基于预训练vit_b_16模型的特征提取、余弦相似度计算以及内存数据库构建与查询等核心环节，查询时按相似度返回最匹配的Top-K结果；Markdown文档则讲解原理与调用示例，可帮助读者快速运行并扩展到自己的图像数据集，整体思路清晰，代码便于二次开发，适用于入门级图像检索实验和轻量级项目集成。
1. 为什么相似图像检索要用ViT：从CNN到Transformer的检索范式变化
基于ViT模型实现相似图像检索，听起来就是把分类模型换成Transformer，真正跑起来才发现，模型只是最短的一块板。先说一个反直觉结论：直接拿预训练ViT提取特征去检索，效果往往打不过微调过的CNN；而把CLS Token的取法、归一化策略和索引结构调对之后，ViT又能反过来领先CNN一截。这套基于ViT的相似图像检索方案，把特征提取、相似度计算、索引构建和结果评估串成一条可复现的流水线，适合正在做AI图像检索、深度学习毕设或者商品同款匹配的开发者。你不需要从零搭建Transformer，只需要踩准特征层、度量和索引三个节点。
2. 把ViT变成特征提取器：CLS Token与Patch Embedding的取舍
ViT本身是分类模型，最后一层输出的是类别概率，检索用的却是图像语义向量。所以第一步要做的不是改网络结构，而是决定从ViT的哪一层取特征、取完怎么加工。这一步如果选错，后面的相似度计算和索引构建都是白搭。
2.1 ViT前向过程：图像怎么变成一连串Token
ViT先把输入图像切成固定大小的Patch，通常是16×16像素，一张224×224的图会被切成14×14共196个Patch，每个Patch展平成向量后经过线性映射变成Token，再和CLS Token拼在一起，加上位置编码，送入标准Transformer Encoder。和CNN不同的是，这个过程里没有卷积的局部归纳偏置，模型对全局关系的建模能力更强，代价是需要更多数据来学特征。
加载模型时最常用的是timm库，它把预训练权重和模型结构封装得比较干净。
import torch import timm # 去掉分类头，只保留特征部分 vit = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=0) vit.eval() dummy = torch.randn(1, 3, 224, 224) with torch.no_grad(): feat = vit(dummy) print(feat.shape) # 输出 (1, 768)这段代码的关键在num_classes=0。这个参数决定模型返回的是特征向量而不是分类logits，对于ViT-Base来说返回维度是768，ViT-Large则是1024。提取特征前必须把模型切到eval模式，否则Dropout路径和归一化层的统计参数会不一致，导致同一张图前后两次提取的特征有细微抖动。torch.randn(1, 3, 224, 224)只是用来确认模型能跑通，后面会替换成真实图像的预处理结果。
这里补充一个容易忽略的点：Patch大小不是随便选的。Patch越小序列越长，模型能看到更细的局部结构，但计算量也同步上涨。vit_base_patch16_224这个名字里，16是Patch大小，224是训练时的输入分辨率。如果换成32的Patch，序列长度会从197降到50，速度更快但细粒度特征明显变弱，检索场景里除非图片本身结构很简单，否则不建议用Patch32。
2.2 提取哪一层特征：CLS Token与池化的对比
模型输出的特征向量里，第0个位置是CLS Token，后面196个位置对应图像各个Patch的Embedding。CLS Token经过整个Transformer的全局交互，可以理解为模型对整张图的语义总结；而Patch Embedding保留了更多局部视觉响应。检索任务要选哪种，取决于你检索的粒度。
# CLS Token：全局语义 cls_feat = feat[:, 0] # 所有Patch Embedding做均值池化：局部响应加权 patch_feat = feat[:, 1:].mean(dim=1) # 两者拼接，维度翻倍，效果不一定更好但特征更冗余 concat_feat = torch.cat([cls_feat, patch_feat], dim=-1)我一般默认先用CLS Token跑通流程，因为它在大多数通用检索场景下表现足够稳定。如果遇到“背景相似但主体不同”或者“纹理相近但类别不同”这类情况，可以换成Patch均值池化再对比一次。需要强调一个常见误用：有的实现直接把整段feat（包括CLS和全部Patch）拉平当作特征，这在ViT上的效果并不好，因为它把196个Patch的局部信息堆成一维向量，会淹没有效语义。上面的对比可以总结成一张表：
| 特征来源 | 维度 | 语义特征 | 适用场景 |
|---|---|---|---|
| CLS Token | 768 | 全局、偏分类语义 | 通用图像检索、同款匹配 |
| Patch均值池化 | 768 | 局部、偏纹理响应 | 材质检索、纹理相似 |
| 中间层特征 | 768×层数 | 多尺度 | 细粒度检索但成本高 |
中间层特征需要改模型前向，一般是用model.forward_features配合特定层的输出，调试成本高，不建议作为第一版方案。
2.3 预处理参数：尺寸与归一化必须和训练一致
ViT对输入尺寸和归一化方式非常敏感，因为位置编码是在224×224下学出来的，如果你用448×448的输入，序列长度会变化，位置编码没办法直接对齐。常见做法是保持和预训练一致，也就是Resize到224×224。如果加载的是384分辨率的模型，Resize目标也要同步改成384。
from torchvision import transforms preprocess = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])这段预处理用的是ImageNet统计出的均值和方差，预训练ViT大多基于这个分布，所以检索阶段必须沿用同一套参数。常见错误是把均值方差填反，或者干脆不归一化，特征分布会整体偏移，检索精度下降得非常明显。另一个容易出问题的地方是图片通道顺序：PIL打开图像默认是RGB，如果用OpenCV读图再转Tensor，得到的是BGR通道，特征直接错位。统一用PIL读图可以避免这类问题。
预处理管线应该作为固定函数封装起来，查询图片和库图片走同一条路径，两个入口各写一份预处理逻辑是典型的翻车前提。后面第5章的避坑部分会再讲这个坑的具体表现。
3. 特征向量计算与相似度度量：让“相似”变成可量化的距离
有了特征向量之后，检索的核心就变成“怎么算两张图相似”。这一步的选型直接决定Top-K结果里是“真相似”还是“看着相似”。很多入门项目在特征层花了大力气，却在这里随便用个距离函数，导致召回结果杂乱。
3.1 余弦相似度 vs 欧氏距离：什么时候用哪个
两个特征向量之间的相似度，最常用的是余弦相似度和欧氏距离。余弦相似度只关注向量方向，对向量模长不敏感；欧氏距离则同时受到方向和模长影响。对于预训练ViT的特征来说，模长差异主要来自图像的纹理丰富度和内容复杂度，这些信息在检索时往往是干扰项，所以更推荐先用余弦相似度。
| 度量方式 | 取值范围 | 对模长敏感 | 等价关系 |
|---|---|---|---|
| 余弦相似度 | [-1, 1] | 否 | 越高越相似 |
| 欧氏距离 | [0, +∞) | 是 | 越小越相似 |
如果特征做了L2归一化，两种度量在排序意义上完全等价：欧氏距离平方 = 2 - 2 × 余弦相似度。也就是说，归一化之后你选哪个都行，真正影响结果的是归一化这步有没有做。这一点很容易被忽视，很多同学跑完发现检索结果不理想，拼命调索引参数，实际上问题出在特征进入距离函数之前。
3.2 特征归一化：不归一化会让大模长图片霸榜
预训练ViT对不同的图输出的特征模长差异很大。某些高纹理图片的模长可能是低纹理图片的两倍以上，直接计算余弦相似度时，虽然公式里有模长归一化，但如果你提前把向量除以模长，结果会更稳定，并且能兼容FAISS里基于内积的索引。
import torch.nn.functional as F # L2归一化，每个特征向量模长变为1 feat_norm = F.normalize(feat, p=2, dim=1)归一化这一步在提取端和查询端都要做。只归一化库特征、不归一化查询特征，算出来的相似度排序会有偏差，因为查询向量的模长没有被消除。把归一化写进特征提取函数里，可以避免两端口径不一致。dim=1是沿着特征维度归一化，也就是每个样本自己除自己的模长；如果误写成dim=0，会把同一维度跨样本归一化，特征分布直接损坏。
3.3 批量提取与检索流水线：从图片目录到Numpy索引
数据准备阶段的典型情况是，你手头有一个图片目录，希望把所有图都过一遍模型，得到特征矩阵。下面这段代码把“扫描目录、预处理、提取特征、保存索引”串起来。
import os import numpy as np from PIL import Image import torch import torch.nn.functional as F from torchvision import transforms preprocess = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def extract_query(image_path): img = Image.open(image_path).convert('RGB') x = preprocess(img).unsqueeze(0) with torch.no_grad(): f = vit(x) return F.normalize(f, p=2, dim=1) image_dir = "./data/images" feature_list, path_list = [], [] for root, _, files in os.walk(image_dir): for name in sorted(files): if not name.lower().endswith(('.jpg', '.jpeg', '.png', '.webp')): continue img_path = os.path.join(root, name) try: f = extract_query(img_path) except Exception as e: print(f"[skip] {img_path}: {e}") continue feature_list.append(f[0].cpu().numpy()) path_list.append(img_path) features = np.vstack(feature_list).astype('float32') np.save("feature_index.npy", features) with open("index_paths.txt", "w", encoding="utf-8") as fp: fp.write("\n".join(path_list))这里有几个参数值得说明。convert('RGB')把灰度图和RGBA图统一转成三通道，避免通道数不一致导致预处理报错。.astype('float32')是给后面接FAISS做铺垫，FAISS要求输入为float32，如果直接存float64，检索时还要再转换一遍。sorted(files)保证多次运行结果顺序可复现。遇到损坏图片时，单张异常不应该中断整个构建流程，所以用try/except跳过并打印原因，这个习惯在几万张图的库上能省很多事。
如果目录里有几万张图，速度瓶颈在图像解码和单张前向。可以先把预处理后的张量收集成一个list，再统一过一个batch，吞吐能提升好几倍。不过生产环境里我一般先用单张跑通验证，再优化成batch，第一版就把batch和异常处理混在一起会很难排查。
4. 让检索跑起来：暴力检索 vs FAISS索引
特征已经就位，接下来面临的是选型问题：直接用Numpy算，还是引入FAISS？答案取决于库的规模。几千张图用Numpy完全没问题，几十万张图就得考虑内存占用和查询延迟。
4.1 暴力检索：数据集不大时三行代码就够
归一化后的特征矩阵正好可以当作一个二维数组，查询向量和它做一次矩阵乘法就能得到全部余弦相似度。
query_feat = extract_query("query.jpg").cpu().numpy() # (1, 768) scores = features @ query_feat.T # (N, 1) flat_scores = scores.flatten() top_k = np.argsort(-flat_scores)[:5] for rank, idx in enumerate(top_k, 1): print(f"Top{rank}: {path_list[idx]} score={flat_scores[idx]:.4f}")因为features已经做过L2归一化，查询向量也做了L2归一化，所以这里的矩阵乘法等价于余弦相似度。np.argsort(-flat_scores)是按相似度从高到低排序，取前5个索引。这个过程的时间复杂度是O(N×768)，一万张图大约几十毫秒，完全够用。暴力检索的最大优点是没有索引构建误差，结果就是精确的Top-K。小数据集上不建议用近似索引，那点速度提升不值得用召回精度来换。
4.2 FAISS索引：IVF与HNSW的选择与参数
当库图片到了百万量级，暴力检索的延迟就不可接受了。FAISS提供了多种索引结构，最常用的是IndexFlatIP、IndexIVFFlat和IndexHNSWFlat。
import faiss d = features.shape[1] # 精确内积索引，相当于Numpy矩阵乘法 index_flat = faiss.IndexFlatIP(d) index_flat.add(features) D, I = index_flat.search(query_feat, k=5) # IVF索引，先聚类再就近搜索 nlist = 100 quantizer = faiss.IndexFlatIP(d) index_ivf = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT) index_ivf.train(features) index_ivf.add(features) index_ivf.nprobe = 10 D, I = index_ivf.search(query_feat, k=5) # HNSW图索引，精度高但构建慢 index_hnsw = faiss.IndexHNSWFlat(d, 32) index_hnsw.hnsw.efConstruction = 200 index_hnsw.add(features) index_hnsw.hnsw.efSearch = 64 D, I = index_hnsw.search(query_feat, k=5)三个索引的适用场景差别很大。IndexFlatIP是精确检索，适合万级以下。IndexIVFFlat通过聚类把搜索空间缩小到几个桶，nlist是聚类数量，nprobe是查询时检查的桶数量，nprobe=10表示只搜索最近的10个聚类，召回率会有轻微损失但速度提升明显，一般从默认值开始调。IndexHNSWFlat用图结构表达特征之间的近邻关系，召回率接近精确检索，M=32是每个节点的最大连接数，efSearch=64是搜索时探索的候选数量，值越大召回越高但查询越慢。需要特别注意：IndexIVFFlat必须先train再add，如果nlist大于等于特征数量，train会直接报错。
这里要提醒一个容易踩的坑：训练IVF索引的nlist一旦设定，库规模变化后需要重新构建。如果你之后往库里加新图，IVF的聚类中心并不会自动更新，新图可能落入不合适的桶。增量更新场景下建议直接用HNSW，或者定期离线重建整个索引。
4.3 检索结果后处理：Top-K、阈值与ID映射
索引返回的是相似度数值和内部索引ID，需要映射回原始图片路径，同时做一步分数过滤。相似度低于某个阈值的查询结果，宁可不要也不要硬排给用户。
MIN_SCORE = 0.65 def search(query_path, k=5, min_score=MIN_SCORE): q_feat = extract_query(query_path).cpu().numpy() D, I = index.search(q_feat, k) results = [] for score, idx in zip(D[0], I[0]): if idx < 0 or score < min_score: continue results.append((path_list[idx], float(score))) return results判定阈值要结合业务来定。同款商品检索在数据集上通常0.8以上才算靠谱，通用内容检索0.6上下也能接受。建议先用一批已知查询做测试，画一个相似度分布图，看真实正样本和硬负样本的分数区间，再切阈值。idx < 0是FAISS返回无效结果时的保护判断，空库或未收敛时会返回-1。
5. 避坑与排查：ViT检索最容易翻车的五个环节
前面几章是正向搭流程，这一章聚焦反向排查。检索效果出问题时，很多人第一反应是调索引参数，其实大部分时候问题出在更早的环节。以下五类问题基本覆盖了ViT检索的大多数事故现场。
5.1 相同图片检索不到自己
现象：拿库里的某张图当作查询图，Top1竟然不是它自己，相似度只有0.9甚至更低。
原因：查询端和入库端的预处理不一致。比如入库时用的是PIL的RGB通道，查询时用了OpenCV的BGR通道，或者两端Resize尺寸不一样。ViT对输入分布极其敏感，通道顺序错误会让特征发生明显偏移。
解决：把第2.3节的预处理封装成一个函数，两端强制使用同一个入口。验证时先用一张图自查询，期望Top1是自己、相似度接近1.0。如果这一步不过，不要往下调索引，先检查预处理差异。
5.2 特征维度对不上
现象：查询特征维度是768，索引使用的是1024维，FAISS直接报错，或者Numpy矩阵乘法的维度不匹配。
原因：入库时加载的是vit_base_patch16_224，查询时换成了vit_large_patch16_224，或者某个环节用了num_classes=1000导致输出了分类logits。
解决：把模型加载封装成统一函数，打印一次feat.shape确认维度。一个实用习惯是保存索引文件时，把特征维度写进文件名，比如features_768.npy，这样查询脚本加载时能立刻发现不匹配。
5.3 小数据集上检索效果不如CNN
现象：库只有几千张图，按CLS Token特征做检索，结果比不过直接在公开数据集上训练的ResNet。
原因：ViT的全局语义特征更适合大数据集下的分布，小数据集上局部细节的区分度反而不如CNN特征突出。这是一个特征适配问题，不是代码实现了Bug。
解决：先尝试把特征从CLS Token换成Patch均值池化，对比两者在这批数据上的检索精度。如果提升不明显，需要考虑第6章的对比学习微调，用库内数据把ViT的特征空间重新塑形。
5.4 索引构建后特征与预处理分离
现象：索引文件构建好后，换一台机器或者换一个脚本加载，检索结果大幅下降。
原因：索引文件和路径列表保存了，但预处理代码没有同步，新环境里可能用了不同的归一化参数或分辨率。
解决：把预处理函数写成一个独立模块，连同模型权重一起随索引文件打包。每次发布检索服务时，跑一遍自查询测试，确认Top1相似度与基线环境一致再上线。
5.5 显存溢出与推理速度慢
现象：批量提取特征时batch size设成64直接OOM，CPU推理一张图要两三秒。
原因：ViT相比CNN计算量更大，尤其在输入分辨率大、序列长度长的情况下。很多开发者沿用CNN场景的batch size习惯，自然扛不住。
解决：显存不够时先用batch=16跑通，再考虑半精度推理。半精度对检索精度的影响非常小，但显存占用直接减半。推理速度方面，把图片固定到224并采用批量预处理，比逐张循环快非常多。如果实时性要求高，可以预先建索引，查询阶段只跑单张前向，这已经把计算量压到最低了。
6. 进阶：用对比学习微调ViT，把检索精度再拉一个档位
6.1 数据增强与对比损失
当库内数据分布和预训练数据集差异较大时，冻结ViT提取的特征不一定够用。常见做法是去掉分类头，在CLS Token后面接一个投影头（非线性层），用对比损失微调。数据增强很关键，同一张图的两个增强视图要最大化相似，不同图的视图要拉开距离。一般用RandomResizedCrop配ColorJitter，太强的增强会把语义破坏掉。
loss = -torch.log( torch.exp(sim_pos / tau) / (torch.exp(sim_pos / tau) + torch.exp(sim_neg / tau)) )loss那段是InfoNCE的核心形式，sim_pos表示正样本对的相似度，sim_neg表示负样本对，tau是温度系数，越小对困难样本的惩罚越强。训练时每个batch里把同一张图的两个增强视图视为正样本对，其他所有视图都当负样本。如果你的显卡规模不大，batch size一般要拉到128以上才能让InfoNCE稳定，不然负样本太少，对比学习的信号会很弱。
6.2 评估指标
微调前后都要跑一遍Recall@K和mAP。把查询集里的每张图都输入模型，统计真正样本出现在Top-K中的比例。这两个指标能量化看出微调带来了多少提升，也能防止你在一个无效方向上调了很多版本。具体到脚本上，我一般直接在自建查询集上跑第4章的search函数，然后按每条查询的标注字段计算命中率；mAP则是对每个查询统计精确率随召回数变化再取平均。微调后特征空间的相似度分布会整体右移，阈值也需要重新切。
我最初做ViT检索时，也踩过不少弯路，最深的教训是：检索效果出问题，第一步永远是检查预处理和特征层，而不是去调索引参数，因为索引只能找回特征里已经存在的语义。从那以后，每次搭检索流水线我都强制走一遍“自查询验证、维度核对、阈值分布分析”三道检查，发现问题的时间从半天缩短到十分钟。希望帮到你。
本文还有配套的精品资源，点击获取
