PatchCore-opt — 昇腾 NPU 推理优化

PatchCore 工业异常检测模型在 Atlas 800I A2 (Ascend 910B4) 上的推理优化项目。 最终成果：15 类 MVTec AD 平均 77.36 FPS（CPU 基准 0.14 → 552 倍加速），AUROC 零回退（0.9886 vs CPU 0.9883）。

性能基线

CPU 基准（原始算法参考）

基准即 PatchCore 官方实现（WideResNet50 骨干网络）在 CPU（鲲鹏 920） 上运行整条推理流水线的性能：

骨干网络推理：CPU FP32（单线程 torchvision）
kNN 搜索：CPU Faiss
后处理：CPU numpy

这是未施加任何优化的原始算法参考基线。

CPU 基准性能指标（15 类 MVTec AD 平均）：

CPU 基准 FPS：0.14（采样 3 个代表性类 bottle/pill/wood 实测）
CPU 基准 AUROC：0.9883（与 NPU 基线一致，算法确定性保证）

NPU 朴素移植基线（优化起点）

将 PatchCore 按原始代码移植到 Atlas 800I A2（Ascend 910B4）的最小改动版本：骨干网络在 NPU FP32 推理，特征向量搬运回 CPU 后用 Faiss 做 kNN 搜索。

NPU 基线 FPS：17.09（NPU 加速骨干网络后，瓶颈仍在 CPU kNN 和主机-设备拷贝） NPU 基线 AUROC：0.9883

优化后结果

项目	CPU 基准	NPU 基线	优化后 (08_npu_knn_v2)
硬件	鲲鹏 920	Atlas 800I A2 (Ascend 910B4)	同
骨干网络	CPU FP32	NPU FP32	NPU FP16 (AMP)
kNN 方法	CPU Faiss	CPU Faiss	NPU 矩阵乘法 (FP16)
后处理	CPU	CPU 主机 (.cpu().numpy())	NPU 张量全流程驻留
15 类平均 FPS	0.14	17.09	77.36
15 类平均 AUROC	0.9883	0.9883	0.9886 (+0.0003)
加速比 v.s. CPU	—	122 倍	552 倍
加速比 v.s. NPU	—	—	4.53 倍

优化成果：在 AUROC 零回退约束下，将 15 类 MVTec AD 平均推理 FPS 从 CPU 基准 0.14 提升至 77.36，相对 CPU 加速 552 倍；相对 NPU 朴素移植加速 4.53 倍。AUROC 从 0.9883 提升至 0.9886。

目录结构

├── scripts/                          # 核心脚本
│   ├── inference_baseline.py         # 基线 NPU 推理
│   ├── benchmark.py                  # 统一基准测试框架
│   ├── accuracy.py                   # 精度验证
│   ├── profiler.py                   # NPU Profiling (torch_npu)
│   └── logger.py                     # 实验日志系统
├── experiments/                      # 优化实验（step-by-step）
│   ├── 00_profiling/                 # Profiling 基线
│   ├── 01_adaptive_pool/             # MeanMapper reshape+mean
│   ├── 02_convbn_fuse/               # Conv+BN 融合 + AMP
│   ├── 03_task_queue/                # TASK_QUEUE_ENABLE=2
│   ├── 04_tcmalloc/                  # tcmalloc 内存优化
│   ├── 05_knn_batch/                 # kNN 批处理
│   ├── 06_graph_mode/                # torch.compile 图模式
│   ├── 07_npu_knn/                   # NPU kNN (PyTorch MatMul)
│   └── 08_npu_knn_v2/                # NPU 张量后处理 + 全流水线 FP16
├── verification/                     # 验证套件
├── validation_report.md              # 精度验证报告（全量 15 类）
└── SKILL.md                          # 可复用技能文档

快速使用

基线推理

python3 scripts/inference_baseline.py \
    --model-path models/IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1 \
    --data-path data/MVTec_AD \
    --classes bottle

最优方案（08_npu_knn_v2）

export TASK_QUEUE_ENABLE=2
export LD_PRELOAD=/usr/lib64/libtcmalloc.so.4
python3 experiments/08_npu_knn_v2/inference_npu_knn.py \
    --model-path models/IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1 \
    --data-path data/MVTec_AD \
    --amp \
    --output results/08_npu_knn_v2.json

精度验证

python3 scripts/accuracy.py \
    --baseline results/baseline_15cls.json \
    --optimized results/08_npu_knn_v2.json \
    --report validation_report.md

实验结果（15 类 MVTec AD 平均，NPU 优化演进）

实验	FPS	相较于 NPU 朴素移植的加速比	平均 AUROC	关键优化
CPU 基准线（纯 CPU 参考）	0.14	—	0.9883	—
NPU 基准线（朴素移植）	17.09	1.000 倍	0.9883	CPU Faiss kNN，FP32
combined_best (仅 bottle 类别)	19.46	1.139 倍	1.0000	卷积+批归一化融合 + 自动混合精度
07_npu_knn	67.28	3.460 倍	0.9883	NPU kNN + 卷积批归一化融合 + 自动混合精度
★ 08_npu_knn_v2	77.36	4.527 倍	0.9886	NPU 张量后处理 + 全流水线 FP16

08_npu_knn_v2 逐类详情

类别	AUROC	AUROC 变化量	嵌入耗时(毫秒)	kNN 耗时(毫秒)	总耗时(毫秒)	FPS
bottle	1.0000	±0.0000	10.98	4.83	15.81	63.24
cable	0.9964	+0.0005	8.16	4.48	12.64	79.10
capsule	0.9733	-0.0020	8.19	4.38	12.57	79.53
carpet	0.9868	±0.0000	7.99	4.65	12.64	79.13
grid	0.9758	-0.0008	7.86	4.64	12.50	80.00
hazelnut	1.0000	±0.0000	7.96	5.19	13.15	76.04
leather	1.0000	±0.0000	8.02	4.55	12.56	79.59
metal_nut	0.9990	±0.0000	8.10	4.51	12.60	79.32
pill	0.9558	+0.0024	8.18	4.43	12.61	80.23
screw	0.9801	+0.0004	8.19	4.33	12.52	80.39
tile	0.9917	+0.0004	8.03	4.51	12.54	79.72
toothbrush	0.9944	±0.0000	8.16	3.76	11.92	83.86
transistor	0.9983	±0.0000	8.96	4.26	13.22	75.63
wood	0.9886	+0.0035	9.25	4.34	13.60	73.55
zipper	0.9892	±0.0000	8.46	4.31	12.77	78.31
平均	0.9886	+0.0003	8.38	4.50	12.88	77.36

精度验证：全部 15 类 AUROC 下降 < 1%，平均 AUROC 较基准线提升 +0.0003 ✓

性能演进

阶段	优化	FPS	相较于 CPU 的加速比
CPU 基准线（纯 CPU 参考）	—	0.14	—
NPU 基准线（朴素移植）	CPU Faiss kNN + FP32	17.09	122 倍
+ 卷积+批归一化融合	减少内核启动	~17.5	~125 倍
+ 自动混合精度 FP16	骨干网络计算加速	~18	~129 倍
+ TASK_QUEUE_ENABLE=2	调度优化	~19	~136 倍
+ tcmalloc	内存分配优化	~19.5	~139 倍
+ NPU kNN	PyTorch 矩阵乘法替代 CPU Faiss	67.28	481 倍
+ NPU 张量后处理 + 全流水线 FP16	消除 .cpu().numpy() 同步	77.36	552 倍

性能分析发现

阶段	基准版本	优化后	说明
骨干网络（WideResNet50）	~52ms	~8ms	AMP FP16 加速 + NPU 稳态
kNN（CPU Faiss → NPU 矩阵乘法）	~7ms	~5ms	消除 NPU→CPU 数据搬运
后处理（.cpu().numpy()）	隐含	~0ms	张量全流程 NPU 驻留
总计	58.52ms	12.88ms	4.53 倍加速

最终优化栈

CPU Baseline: CPU (Kunpeng 920) Faiss, FP32, 0.14 FPS

NPU Baseline (朴素移植, 17.09 FPS)
├── Conv+BN 融合               → 减少 kernel launch
├── AMP FP16 backbone           → 计算加速
├── _adaptive_avg_pool1d_exact  → 精度安全的池化替代
├── NPU kNN (PyTorch MatMul)    → 消除 NPU→CPU 数据搬运
├── ★ NPU 张量后处理            → 消除 .cpu().numpy() 同步
├── ★ 扩展 AMP 全流水线         → ResizeBilinear 等保持 FP16
├── TASK_QUEUE_ENABLE=2         → NPU 调度优化
└── tcmalloc                     → 内存分配优化
                                        ↓
                       08_npu_knn_v2: 77.36 FPS
                       (552x speedup v.s. CPU, 4.53x v.s. NPU naive)

环境要求

软件依赖

Python 3.10+
PyTorch + torch_npu (Ascend NPU)
numpy、scikit-learn、tcmalloc（可选）
patchcore-inspection（子模块，模型与数据预备）

预训练模型下载

本项目使用 WideResNet50 骨干网络的 PatchCore 预训练模型。

来源	地址
GitHub Release (官方)	patchcore-models.zip
Release 页面	https://github.com/amazon-science/patchcore-inspection/releases

# 下载并解压
wget https://github.com/amazon-science/patchcore-inspection/releases/download/v1.0.0/patchcore-models.zip
unzip patchcore-models.zip -d models/
# 预期目录: models/IM224_WR50_L2-3_P01_D1024-1024_PS-3_AN-1/models/

MVTec AD 数据集下载

来源	地址
MVTec 官网 (官方)	https://www.mvtec.com/company/research/databases/mvtec-ad/
镜像下载 (tar.xz)	mvtec_anomaly_detection.tar.xz
HuggingFace (Synthetic)	`huggingface-cli download siins/AD_MVTec --local-dir data/mvtec_synthetic`

# 方式 A: 镜像下载
wget https://www.mydrive.ch/shares/38536/3830184030e49fe74747669442f0f282/download/420938113-1629952094/mvtec_anomaly_detection.tar.xz
tar -xJf mvtec_anomaly_detection.tar.xz -C data/

# 方式 B: HuggingFace
pip install huggingface-hub
huggingface-cli download siins/AD_MVTec --local-dir data/mvtec_synthetic

一键下载

bash verification/scripts/download_assets.sh --all

精度修复说明

_adaptive_avg_pool1d_exact 采用 cumsum + 索引采集的方式实现，与原始 F.adaptive_avg_pool1d 在数学上等价（最大差异 < 3e-6），解决了使用 F.interpolate() 替代方案导致 AUROC 从 1.0 降至 0.54 的问题。

本项目是赛道2「PatchCore 昇腾NPU推理优化」的参赛提交。完整验证报告参见 validation_report.md。