ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

PatchCore工业缺陷检测:特征嵌入原理、代码复现与避坑实践

PatchCore工业缺陷检测:特征嵌入原理、代码复现与避坑实践 简介面向工业质量检测与缺陷识别场景的PatchCore算法实践资源主要服务于有机器学习与图像处理基础的本科生、研究生及数据科学从业者。资源以一次完整的工业缺陷检测大作业为线索系统梳理了基于嵌入的缺陷检测方法详解特征提取、核心集采样如KCenterGreedy、最近邻搜索、异常热力图生成与可视化等关键步骤从训练阶段构建核心集、测试阶段重叠分块与异常分数计算到像素级/样本级AUC评估和ROC/PR曲线绘制均有清晰说明与实现指引。尤其适合需要完成工业缺陷检测课程设计或希望落地MVTec AD数据集相关算法的读者可据此构建完整可运行的检测流程并掌握关键评估方法。包内为1个PDF文件压缩包大小192KB内容以作业方案和算法说明为主结构紧凑。已有459人浏览学习。1. 工业缺陷检测为什么转向特征嵌入PatchCore要解决的是“没见过的那类缺陷”流水线上最难处理的缺陷不是划痕、污渍这种有明确样本、能靠分类模型学出来的类型而是那些只出现过一两次、甚至压根没进过样本库的新形态异常。传统视觉检测的思路是“先认识缺陷长什么样再在图上找”但工业现场恰恰相反正常产品的外观可以有轻微波动缺陷却往往没有固定形态。基于特征嵌入的工业缺陷检测技术也就是 PatchCore 这类算法的核心思路是把“检测缺陷”重新定义成“找最相似特征距离”只用正常产品参与训练测试时把图像块的嵌入特征和正常特征记忆库做比对距离越远越像缺陷。这个思路避开了缺陷样本不足的问题也是它能替代一部分传统分类检测方案的原因。这篇实践笔记会从特征嵌入原理讲到可复现的训练、推理代码再把我在复现 PatchCore 过程中踩过的坑逐个拆开。2. 特征嵌入三板斧为什么 PatchCore 要取中间层、建记忆库、做核心集采样2.1 为什么是中间层特征而不是分类层特征很多第一次接触 PatchCore 的人会问既然用了 ImageNet 预训练模型为什么不直接拿 ResNet 最后的 1000 维分类特征当图像描述子原因不难理解分类层经过全局池化和全连接之后特征已经把“位置信息”几乎丢光了只剩“这张图整体的语义”。工业缺陷检测恰恰需要保留“缺陷在哪个位置、那块区域的纹理长什么样”的信息。比如一个针尖大小的漏焊点在整张图上占比不到 1%池化之后它的响应会被周围正常区域稀释掉分类特征可能完全看不出差异。PatchCore 的做法是用预训练 ResNet 的中间层输出比如 layer2 和 layer3 的特征图。ResNet 的 layer2 输出尺寸大约是输入的 1/8layer3 是 1/16特征图上的每个像素点都对应着原图上一块区域学术上把这叫 patch也就是图像块。每个位置的 1536 维向量512 1024 拼接就是一个“局部纹理描述子”它同时保住了空间位置和局部细节。这也回答了“特征嵌入”这四个字在 PatchCore 里的含义用预训练网络把图像块映射到一个高维特征空间在这个空间里正常区域的特征会很接近缺陷区域的特征会偏离正常簇。实际选层时我一般会同时挂 layer2 和 layer3 的 forward hook而不是只用单层。只用 layer2 虽然局部细节丰富但对轻微的全局色差、光照变化过于敏感容易误报只用 layer3 对细小缺陷不敏感容易漏检。两层拼在一起算是 PatchCore 里的常见配置后面代码部分会给出具体实现。2.2 特征记忆库把“正常”浓缩成一批向量把训练集里所有正常图像的 patch 特征都收集起来就得到了一个特征记忆库。假设训练集有 600 张 224×224 的正常图经过 layer2layer3 的特征图分辨率大约 14×14每张图能产生约 196 个 patch 特征全量特征就有 600×196 ≈ 11.7 万条。这个量级对近邻检索来说还扛得住但如果生产线上的正常图片更多比如 5000 张特征就会膨胀到百万级别。每条特征 1536 维float32 存储的话 100 万条就要占用约 6GB 内存这还没算推理时的距离计算开销。所以 PatchCore 在训练阶段会对记忆库做一次“压缩”论文里叫核心集采样Coreset Sampling。它通过贪心算法选出一个子集让选出来的特征尽量均匀覆盖原始特征分布。换句话说不是简单随机抽而是在特征空间里每次挑一个“离已选集合最远”的点这样保留的特征多样性最高。采样率可以直接控制0.01 就是把记忆库压缩到 1%十几万条特征最后只剩一两千条计算开销大幅下降精度损失往往只有零点几个百分点。这个设计在硬实时检测场景里很关键毕竟工控机上的 CPU 推理和实验室里的 GPU 显卡是两回事。2.3 推理时的最近邻距离为什么异常分数不用分类器PatchCore 推理时没有任何分类头它把测试图的每个 patch 特征和记忆库里的特征做近邻检索取第 k 近的距离作为这个 patch 的异常得分。这里的直觉是如果某个图像块和记忆库里最相似的正常块都离很远那它八成属于训练分布之外的东西也就是异常。图像级得分则是对所有 patch 得分取最大值用得分最高的那个缺陷区域代表整张图的异常程度。这里有个值得注意的细节距离度量通常用余弦距离而不是欧氏距离。特征归一化之后余弦距离在数值上等价于归一化特征的欧氏距离但余弦对特征向量的模长不敏感能避免某张图对比度偏高导致整体特征模长变化带来的误判。k 的取值论文里一般用 2实际复现时 k 取 2 和取 1 差别不大但 k 太大会把细小的缺陷得分摊平导致漏检。这个参数后面在调参表里会再讲。3. PatchCore 训练代码复现从 MVTec AD 到记忆库构建3.1 环境准备和数据集组织方式复现 PatchCore 不需要特殊依赖核心库就是 PyTorch、torchvision、scikit-learn 和 numpy。PyTorch 1.10 以上都可以torchvision 最好带 ImageNet 预训练权重。数据集方面工业缺陷检测领域最常用的公开基准是 MVTec AD它包含瓶、螺丝帽、电缆、皮革等十几个类别的正常样本和缺陷样本。每个类别目录下会分 train 和 testtrain 里只有正常图test 里同时有正常和缺陷图ground truth 是像素级 mask但做图像级检测时只需要知道这张图是否异常。数据组织方式我建议不要直接按 MVTec AD 原始目录硬编码而是先整理成一个统一入口import os from PIL import Image from torch.utils.data import Dataset class MVTecADDataset(Dataset): def __init__(self, root, category, splittrain, transformNone): self.root root self.transform transform self.paths [] self.labels [] if split train: img_dir os.path.join(root, category, train) for name in os.listdir(img_dir): self.paths.append(os.path.join(img_dir, name)) self.labels.append(0) # 正常样本 else: img_dir os.path.join(root, category, test) for name in os.listdir(img_dir): label 0 if name.startswith(good) else 1 self.paths.append(os.path.join(img_dir, name)) self.labels.append(label) def __len__(self): return len(self.paths) def __getitem__(self, idx): img Image.open(self.paths[idx]).convert(RGB) if self.transform: img self.transform(img) return img, self.labels[idx]这段代码里有两个细节值得注意。第一convert(RGB)必须写很多工业相机采集的是灰度图或带 alpha 通道的 PNG不转的话后续 ResNet 输入维度会出错。第二缺陷样本在 MVTec AD 的 test 目录里是文件名前缀区分我这边的习惯是用startswith(good)判断比读取 mask 文件更直接做图像级评估时够用。预处理策略建议直接用 ImageNet 的标准参数Resize 到 256×256CenterCrop 到 224×224归一化均值[0.485, 0.456, 0.406]标准差[0.229, 0.224, 0.225]。PatchCore 论文里还有一个 bucketing 的缩放技巧简单说就是多尺度缩放后再提特征能小幅提升精度但第一次复现建议先不加跑通基线再回头优化。3.2 用预训练 ResNet 提取 patch 特征核心代码训练阶段的特征提取是整个 PatchCore 的核心。我把提取逻辑拆成一个独立函数方便训练和推理共用import torch import torch.nn.functional as F import torchvision class PatchFeatureExtractor: def __init__(self, backbone_nameresnet50, devicecuda): self.device device self.features {} self.backbone torchvision.models.resnet50( weightstorchvision.models.ResNet50_Weights.IMAGENET1K_V1 ).to(device).eval() # 注册 forward hooklayer2 和 layer3 是 PatchCore 常用特征层 self.backbone.layer2.register_forward_hook( self._make_hook(layer2) ) self.backbone.layer3.register_forward_hook( self._make_hook(layer3) ) def _make_hook(self, name): def hook_fn(module, input, output): self.features[name] output.detach() return hook_fn def extract(self, images): self.features.clear() with torch.no_grad(): _ self.backbone(images.to(self.device)) f2 self.features[layer2] # B, 512, H, W f3 self.features[layer3] # B, 1024, H, W # 把 layer3 上采样到 layer2 分辨率再沿通道拼接 f3 F.interpolate( f3, sizef2.shape[-2:], modebilinear, align_cornersFalse ) patch_feat torch.cat([f2, f3], dim1) # B, 1536, H, W B, C, H, W patch_feat.shape # 每个位置就是一个 patch 特征 patch_feat patch_feat.permute(0, 2, 3, 1).reshape(-1, C) patch_feat F.normalize(patch_feat, p2, dim1) return patch_feat.cpu().numpy()逻辑说明forward hook 可以在不改动 ResNet 源码的情况下拿到中间层输出。layer2 输出通道 512layer3 输出通道 1024拼接后每条特征 1536 维。由于两者的特征图分辨率不同需要先把 layer3 上采样到 layer2 的尺寸再用 permute 把通道放到最后一维得到形状为(patch数量, 1536)的特征矩阵。最后做一次 L2 归一化让后续距离计算等价于余弦距离。调用时需要注意的是extract()接收的是一个 batch 的 Tensor也就是形状为(B, 3, 224, 224)。一张图会返回 256×256 分辨率下的(28, 28)patches 数量吗不对224×224 输入经过 ResNet 的 4 次下采样后layer2 输出是 28×28所以一张图有 784 个 patch 特征。这个数量比论文里常用的 14×14 要多原因是 ResNet 默认 stride而不是额外池化。想减少 patch 数量可以加一层 AdaptiveAvgPool但不建议因为空间分辨率直接影响缺陷定位精度。3.3 记忆库采样greedy coreset 与采样率参数提取完所有正常图的 patch 特征后把矩阵拼起来就得到了原始记忆库。全量记忆库直接用于近邻检索不是不行但推理延迟和内存占用会很难看。这里用贪心核心集采样压缩import numpy as np def greedy_coreset(samples, budget): n samples.shape[0] # 随机选第一个代表点 idx np.random.randint(n) selected [idx] # dists 记录每个点到已选集合的最近距离 dists np.linalg.norm(samples - samples[idx], axis1) while len(selected) budget: # 每次选距离已选集合最远的点 new_idx np.argmax(dists) selected.append(new_idx) # 更新每个点到新集合的最近距离 new_dists np.linalg.norm(samples - samples[new_idx], axis1) dists np.minimum(dists, new_dists) if len(selected) % 500 0: print(fcoreset sampling: {len(selected)}/{budget}) return np.array(selected) memory_bank np.load(memory_bank.npy) ratio 0.01 budget max(1, int(memory_bank.shape[0] * ratio)) coreset_idx greedy_coreset(memory_bank, budget) np.save(coreset_idx.npy, coreset_idx)逻辑说明这个实现是标准的贪心最远点采样核心就一句话“每次从剩余点里挑离已选集合最远的那个”。dists 数组始终保存每个样本到已选集合中最近点的距离所以每选一个新点只需计算它到所有点的距离再和旧 dists 做逐元素最小值。复杂度是 O(budget × n × d)在百万级特征上会有点慢好处是效果好且没有额外依赖。如果数据量实在太大可以用 Faiss 的 IndexFlatIP 加速距离计算但第一次复现用 numpy 版足够。采样率 ratio 是 PatchCore 里最值得调的参数直接决定了推理速度和精度平衡参数推荐值说明ratio0.01压缩到 1%精度损失 0.2~0.5%推荐优先用ratio0.001压缩到 0.1%速度快适合 CPU 部署但细纹缺陷可能被采样掉k2近邻数k1 偏激进k3 以上会把小缺陷得分拉低特征层layer2layer3只选 layer3 省显存但小缺陷漏检距离余弦距离对光照变化鲁棒避免欧氏距离受特征模长影响采样完的 coreset 就是训练阶段的产物后续推理只需加载 coreset 和索引关系。这里有个容易忽略的问题采样前一定要先把特征归一化否则贪心选出来的点会被少数模长很大的离群点带偏采完再归一化就晚了。4. PatchCore 推理与评估计算图像级 AUROC 和生成热图4.1 推理代码计算 patch 得分并还原到图像训练阶段产出一个 coreset推理阶段要做的事就是逐 patch 检索、得距离、聚合。下面这段是推理的核心逻辑from sklearn.neighbors import NearestNeighbors class PatchCoreInference: def __init__(self, coreset, k2): self.knn NearestNeighbors(n_neighborsk, metriccosine) self.knn.fit(coreset) # coreset 必须是归一化后的特征 def predict_patch_scores(self, patch_feat, original_size): # patch_feat: (P, 1536)P 对应 H*W dists, _ self.knn.kneighbors(patch_feat) # 取第 k 近的距离作为 patch 得分 patch_scores dists[:, -1] # 还原成特征图分辨率例如 28x28 h w int(np.sqrt(patch_scores.shape[0])) score_map patch_scores.reshape(h, w) # 放大到原始输入尺寸方便可视化 score_map cv2.resize(score_map, original_size, interpolationcv2.INTER_NEAREST) return score_map def predict_image_score(self, patch_feat): dists, _ self.knn.kneighbors(patch_feat) return float(dists[:, -1].max())逻辑说明NearestNeighbors 的metriccosine会自动对输入行做归一化前提是 coreset 和测试特征都已经归一化过加一道保险不算多余。dists[:, -1]取每个 patch 到 coreset 的第 k 近距离代表这个局部区域的异常程度。图像级得分直接取所有 patch 得分的最大值这也是 PatchCore 的一个特点它不取平均因为平均会把小缺陷区域的异常响应稀释掉。score_map 放大回原图尺寸时我用的是INTER_NEAREST也就是最近邻插值目的是保留 patch 的原始边界不把缺陷区域平滑变形。做热图可视化时可以用 OpenCV 的applyColorMap把 score_map 转成伪彩色图叠加到原图上方便直接定位缺陷区域。4.2 图像级评估AUROC 是怎么算出来的工业缺陷检测里最常见的评估指标是图像级 AUROC它不依赖人工设定阈值只看正常样本和异常样本的分数分布是否分得开。计算方式很简单from sklearn.metrics import roc_auc_score image_scores [] true_labels [] test_dataset MVTecADDataset(root, category, splittest, transformtransform) test_loader DataLoader(test_dataset, batch_size32, shuffleFalse) for images, labels in test_loader: patch_feat extractor.extract(images) # 每张图的 patch 数量需要按 batch 切分 n_patches patch_feat.shape[0] // images.shape[0] for i in range(images.shape[0]): single_patch patch_feat[i * n_patches: (i 1) * n_patches] image_scores.append(inferencer.predict_image_score(single_patch)) true_labels.append(labels[i].item()) auc roc_auc_score(true_labels, image_scores) print(fcategory{category}, image AUROC{auc:.4f})逻辑说明这里的关键是切分 patch 特征。Extractor 会返回整个 batch 的所有 patch形状是(B*P, 1536)所以拿到结果后必须先除以 batch 大小再逐张图去算图像级得分。AUROC 的值可以从两个角度看0.9 以上说明正常和异常特征分布已经明显分离低于 0.8 大概率不是模型不行而是预处理、特征层或数据划分出了问题这在后一章里会展开讲。4.3 PatchCore 的三个关键调参项跑通第一版之后我一般会盯三个调参项优先做增益最大的改动。第一是采样率 ratio。如果推理性能吃紧先看 coreset 压缩。千分之一采样率和百分之五采样率在 MVTec AD 多数类别上差距不到 1%但在 CPU 上的检索耗时能差出好几倍。第二是 k 近邻数。k2 在大部分场景下都稳定但如果检测目标是细小划痕可以试 k1k 值越大异常得分越趋向于“整体不相似”小缺陷反而会被忽略。第三是特征层选择。layer2layer3 是基线如果误报率高可以去掉 layer2 只保留 layer3如果漏检率高就引入 layer1 或增加多尺度 bucketing。这几个调整不像训练超参数那样需要反复跑实验每次改完可以直接在测试集上看 AUROC 变化反馈很快。5. PatchCore 避坑指南数据泄漏、显存溢出、精度不对的五个典型场景5.1 灰度图和 RGB 混合加载导致维度错位现象跑训练特征提取时报RuntimeError: mat1 and mat2 shapes cannot be multiplied或者干脆在backbone前向时报维度不匹配。原因工业数据集里常常混着灰度图和三通道彩色图灰度图用Image.open()读出来是单通道直接ToTensor()后变成 1×H×WResNet 的第一个卷积层要求 3 通道输入自然报错。PyTorch 对单通道输入不会自动复制成三通道这是最容易踩的一处。解决在 Dataset 的__getitem__里统一加.convert(RGB)不要靠外部预处理管道的侥幸。这个操作会把灰度图复制到三个通道代价极小但能省掉一晚上的排查时间。顺带提醒Tiff 格式的工业图也一样最好先转成标准 RGB 再进模型。5.2 记忆库直接进内存OOM 和 memmap现象训练时提取完几千张图的特征后np.concatenate阶段直接MemoryError或者机器开始疯狂换页训练进程被 OOM Killer 杀掉。原因全量记忆库的特征数量是“图片数 × patch 数”一个 5000 张图的数据集patch 特征就是 5000×784 ≈ 390 万条每条 1536 维 float32换算下来要二十多 GB 内存。很多复现脚本把所有特征先拼再采内存自然爆掉。解决两个手段同时用。一是在特征提取循环里先保存到.npy文件不要常驻内存二是用np.memmap做磁盘映射只把需要计算的部分读进内存。贪心核心集采样每次要算全量距离用 memmap 做分块计算可以显著降低峰值内存。如果业务上不需要全量记忆库也可以分段采样再合并但尽量先试 memmap效果更好。5.3 AUROC 接近 1.0训练测试输入尺寸不一致造成的数据泄漏现象某类在训练时 AUROC 达到 0.99但上线后表现远不如预期漏检率明显偏高。原因这是复现里一个很隐蔽的坑。代码里训练集用了T.Resize((256, 256))加T.CenterCrop((224, 224))但测试集只做了T.Resize((224, 224))没有 CenterCrop。两者输入分布不一致模型在测试时看到的图像空间范围和训练时完全不同异常得分自然会失真。更严重的是如果同一张正常图同时出现在训练集和测试集数据泄漏会让 AUROC 虚高。解决训练和测试必须用完全相同的预处理流程尤其是 Resize 和 CenterCrop 的顺序、尺寸不能有半点差异。另外检查 test 目录里是否混入了 train 目录的副本这在拼接多来源数据集时很容易发生。建议写一个自动检查脚本对比所有图像的文件名和 MD5把重复样本剔除后再评估。5.4 复现结果和论文差距大特征层和采样率的影响现象同一份 MVTec AD 数据论文里某个类别 AUROC 是 0.98复现出来只有 0.91甚至更低。原因PatchCore 的结果对特征层组合和采样率很敏感。只取 layer3 时细小缺陷容易被忽略采用随机采样替代贪心核心集采样精度也会有明显下降。还有一个常见问题是特征归一化被遗漏导致余弦距离和欧氏距离混用数值分布完全偏离预期。论文的完整配置是 layer2layer3、贪心核心集、采样率 1%、k2少哪个都会掉点。解决先严格按这组配置复现再逐步替换。每换一个变量就单独记录 AUROC比如把 layer2 去掉对比一次、把随机采样换成贪心采样对比一次。这样不仅能定位差距来源还能积累出自己数据集上的最佳配置。另外MVTec AD 的不同类别难度差异很大像螺丝帽这种纹理复杂类别90% 和 98% 之间可能就差一个 bucketing 多尺度预处理不要急着怀疑实现有 bug。5.5 多进程 DataLoader 卡死和显存占用异常现象训练时DataLoader(num_workers8)跑一段时间后进程卡住或者 GPU 显存被占满batch 明明只有 32 张图却报 CUDA OOM。原因num_workers过高时Windows 环境下 PyTorch 的 DataLoader 偶发卡死而显存异常攀升是因为 forward hook 里保存了中间层输出PyTorch autograd 会把所有梯度计算图保留在显存里直到torch.no_grad()退出后才释放。如果用extract()作为独立函数注意self.features是实例属性会一直持有张量引用导致显存无法回收。解决特征提取阶段务必包在with torch.no_grad():里并且每次提取完主动del self.features或调用.cpu()把张量搬离显存。DataLoader 遇到卡死时先从num_workers0跑一遍确认逻辑没问题再逐步加 worker 数。工业数据量通常没大到需要 8 个 workernum_workers2往往又稳又快。还有一点hook 保存的 output 本身是中间激活图如果显存吃紧可以在 hook 里直接做上采样和拼接降低中间张量驻留时间。6. 用 t-SNE 把 PatchCore 特征嵌入画出来验证效果的一个实用小技巧调参调到一定程度靠 AUROC 一个指标很难判断是“特征没分开”还是“距离阈值不对”。我习惯在项目汇报前把正常样本和异常样本的 patch 特征各抽一部分做 t-SNE 可视化直接看两个簇的重叠程度这比对着数字猜直觉管用得多。import numpy as np from sklearn.manifold import TSNE import matplotlib.pyplot as plt normal_samples memory_bank[np.random.choice(len(memory_bank), 2000, replaceFalse)] abnormal_feat extractor.extract(abnormal_images_loader) abnormal_samples abnormal_feat[np.random.choice(len(abnormal_feat), 2000, replaceFalse)] embedding np.vstack([normal_samples, abnormal_samples]) labels np.array([0] * 2000 [1] * 2000) tsne TSNE(n_components2, perplexity30, n_iter500, random_state42) xy tsne.fit_transform(embedding) plt.figure(figsize(8, 6)) plt.scatter(xy[labels 0, 0], xy[labels 0, 1], s2, alpha0.5, labelnormal) plt.scatter(xy[labels 1, 0], xy[labels 1, 1], s2, alpha0.5, labelabnormal) plt.legend() plt.show()如果图上出现两个清晰的簇说明 PatchCore 的特征嵌入已经把正常和异常分开后续调阈值就行如果两个簇完全重叠就要回到特征层和采样率上去找原因。还有一个小技巧把异常得分最高的一些 patch 特征单独高亮用不同颜色标注在 t-SNE 散点图上能直观看到“哪些特征的异常分数是虚高的”对排查误报来源帮助很大。生产环境里确定阈值时我一般不会直接拿测试集的 AUROC 曲线上找最优点而是用训练集正常图的得分分布画直方图取 99.9 分位数作为初始阈值再用验证集微调。这么做的好处是阈值不依赖缺陷样本的数量因为工业场景里缺陷样本永远是不够的。最后说一个我自己的教训第一次用 PatchCore 时为了省显存只取了 layer3 特征视觉化之后发现异常簇和正常簇贴得很近怎么调阈值都压不住误报换成 layer2layer3 拼接后两个簇才明显分开。特征嵌入这件事信息量比模型复杂度更值钱希望在复现 PatchCore 的路上这篇笔记能帮你少走几步弯路。本文还有配套的精品资源点击获取
返回列表