PatchCore NPU Inspection Skill
1. 基本信息
| 字段 | 值 |
|---|---|
| Skill 名称 | patchcore-npu-inspection |
| 适用场景 | 工业异常检测 (Industrial Anomaly Detection) |
| 适配硬件 | 昇腾 800I A2 (Atlas 800T A2) |
| 框架依赖 | PyTorch ≥ 1.9, torch_npu ≥ 2.1.0 |
2. 功能描述
本 Skill 基于 PatchCore 算法（Roth et al., CVPR 2022），在昇腾 800I A2 NPU 上实现了 端到端的工业异常检测推理流水线。
主要特性
| 特性 | 说明 |
|---|---|
| 全链路 NPU 迁移 | 从骨干网络推理到最近邻搜索全程运行在 NPU 上，零 CPU 回退 |
| BaihuNN 搜索引擎 | 自研百行级 NN 搜索库，用 torch.mm 矩阵乘替代 Faiss GPU/CUDA 依赖 |
| 自动环境检测 | 自动检测 NPU 可用性、CPU 核数、内存/显存，动态调优 batch_size 和采样率 |
| 三档加速模式 | normal（保精度）、fast（2–3×）、turbo（5–8×） |
| 15 类 MVTec-AD 全覆盖 | bottle, cable, capsule, carpet, grid, hazelnut, leather, metal_nut, pill, screw, tile, toothbrush, transistor, wood, zipper |
3. 使用方法
3.1 环境准备
确保环境中已安装昇腾 NPU 驱动（CANN ≥ 7.0.RC1）及 PyTorch NPU 插件：
# 验证 NPU 可用
python -c "import torch; print(torch.npu.is_available())"
# 应输出: True
3.2 模型推理
# 克隆仓库
git clone https://gitcode.com/quzhi_1981/Patchcore.git
cd Patchcore/src
# 运行推理（全 15 类，自动检测 NPU）
python inference.py --data_dir /path/to/mvtec_ad
# 仅指定类别调试
python inference.py --data_dir /path/to/mvtec_ad --categories bottle cable
3.3 加速模式
# normal 模式 — 保精度（比赛提交推荐）
python inference.py --data_dir /path/to/mvtec_ad --mode normal
# fast 模式 — 约 2–3× 加速，精度微降 ~0.5%
python inference.py --data_dir /path/to/mvtec_ad --mode fast
# turbo 模式 — 约 5–8× 加速，精度降 ~1–2%
python inference.py --data_dir /path/to/mvtec_ad --mode turbo
4. 性能与精度
4.1 测试环境
| 项目 | 配置 |
|---|---|
| 硬件 | 昇腾 800I A2 (Atlas 800T A2) |
| 显存 | 64 GiB |
| CANN | 7.0.RC1 |
| PyTorch | 2.1.0 / torch_npu 2.1.0 |
| 数据集 | MVTec-AD（15 类，~5 GB） |
4.2 优化前后对比
| 指标 | 优化前（Faiss CPU） | 优化后（BaihuNN NPU） | 提升 |
|---|---|---|---|
| 15 类总耗时 | ~1710s (28.5min) | 1363.7s (22.7min) | 快 25% |
| 特征搜索方式 | Faiss CPU 暴力搜索 | BaihuNN (torch.mm 矩阵乘) |
算子级加速 |
| 设备利用率 | CPU 核利用率低 | NPU MATRIX 单元满负荷 | 架构匹配 |
| 精度 (ImgAUROC) | 0.990 | 0.9905 | 持平 |
| 依赖 | faiss-cpu（额外安装） | 零额外依赖，仅 PyTorch | 部署更轻量 |
4.3 精度指标（MVTec-AD 15 类平均）
| 指标 | 论文基线 (GPU) | NPU 实测 | 偏差 |
|---|---|---|---|
| Image AUROC | 0.990 | 0.9905 | +0.0005 |
| Pixel AUROC | 0.980 | 0.9812 | +0.0012 |
| PRO | — | 0.9963 | — |
15 类中有 14 类达到或超过论文水平，最大偏差 < 0.0015，满足赛事精度要求。
5. 适配改造点
| # | 改造点 | 原实现 (CUDA) | NPU 适配后 | 涉及文件 |
|---|---|---|---|---|
| 1 | 设备抽象 | tensor.cuda() / model.to('cuda') |
tensor.to(device) / model.to(device)，由 --device auto 自动选择 npu:0 或 cpu |
inference.py |
| 2 | NN 搜索替换 | Faiss GPU (faiss-gpu) |
BaihuNN — 用 torch.mm 矩阵乘在 NPU 上执行两阶段检索（FP16 粗筛 → FP32 精排），消除 CPU↔NPU 数据搬运 |
src/patchcore/baihu_nn.py |
| 3 | 显存管理 | torch.cuda.empty_cache() |
torch.npu.empty_cache()，兼容两者 |
inference.py |
| 4 | 随机种子 | torch.cuda.manual_seed() |
torch.npu.manual_seed()，兼容两者 |
inference.py |
| 5 | 数据加载 | 无特殊处理 | pin_memory=True 配合 NPU 异步传输 |
inference.py:get_dataloader() |
| 6 | 自动环境检测 | 无 | 检测 torch_npu 是否安装、NPU 设备数和显存、CPU 核数、物理内存 |
inference.py:check_environment() |
| 7 | 参数自动调优 | 无 | 根据 NPU 显存/CPU 内存自动设置 batch_size、num_workers、采样率 | inference.py:auto_tune_params() |
关键技术细节
BaihuNN 两阶段检索（核心优化）：
阶段 1（NPU AI Core, FP16）：
dist² = ||q||² + ||x||² - 2 · q · xᵀ ← torch.mm 矩阵乘，充分利用 MATRIX 单元
选出 Top-K 候选
阶段 2（CPU, FP32）：
对 Top-K 候选执行精确 L2 距离排序 ← 仅 K 条，开销极小
该方案在保持精度的同时，将 KNN 搜索从 CPU 核低效串行转为 NPU 矩阵乘高效并行，是本次调优的核心加速手段。
6. 依赖清单
| 库 | 最低版本 | 说明 |
|---|---|---|
| Python | ≥ 3.8 | — |
| torch | ≥ 1.9.0 | PyTorch 框架 |
| torch_npu | ≥ 2.1.0 | 昇腾 NPU PyTorch 插件 |
| torchvision | ≥ 0.16.0 | 图像处理 |
| timm | ≥ 0.6.0 | 骨干网络（Wide ResNet-50） |
| faiss-cpu | ≥ 1.7.0 | CPU 近邻搜索（仅回退用） |
| scikit-learn | ≥ 1.0 | 评估指标 |
| scikit-image | ≥ 0.19 | 图像后处理 |
| tqdm | ≥ 4.60 | 进度条 |
7. 输出说明
控制台输出
运行结束后逐类打印指标并汇总：
类别 ImgAUROC PixAUROC PRO 训练(s) 推理(s) 总计(s)
--------------------------------------------------------------------------------
bottle 1.0000 0.9848 0.9999 52.4 5.2 57.6
...
平均值 0.9905 0.9812 0.9963 66.4 6.8 73.2
端到端总耗时（挂钟）: 1363.7 秒 (22.7 分钟)
CSV 文件
保存至 results/inference_<timestamp>/results.csv，含逐类和平均结果。
8. 注意事项
- 首次运行会自动下载 Wide ResNet-50 骨干网络权重（需联网）
- MVTec-AD 数据集需提前下载并按
bottle/train/good/等子目录结构存放 - NPU 模式下自动启用 BaihuNN 替代 FaissNN 做最近邻搜索
- 如 NPU 显存不足，脚本自动降低 batch_size 和采样率