当前位置：首页 > news >正文

保姆级教程：在Windows上用PyCharm配置Anomalib，训练自己的Padim模型并导出ONNX

news 2026/6/3 23:15:44

保姆级教程：在Windows上用PyCharm配置Anomalib，训练自己的Padim模型并导出ONNX

工业缺陷检测一直是制造业中的痛点问题。传统的监督学习方法不仅需要大量标注数据，还常常因为样本不平衡导致模型性能不佳。而无监督学习算法如Padim，仅需正常样本即可训练，大大降低了数据准备的门槛。本文将手把手教你如何在Windows环境下，通过PyCharm配置Anomalib库，完成Padim模型的训练和ONNX导出。

1. 环境准备与项目初始化

1.1 基础环境配置

在开始之前，请确保你的Windows系统满足以下要求：

Windows 10/11 64位系统
Python 3.8或3.9（Anomalib对Python版本有特定要求）
PyCharm 2022.3或更高版本
8GB以上内存（推荐16GB）
NVIDIA显卡（可选，但推荐）

首先创建一个新的Python虚拟环境：

python -m venv anomalib_env .\anomalib_env\Scripts\activate

安装必要的依赖包：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117 pip install anomalib[full] pip install onnx onnxruntime

注意：如果使用CUDA加速，请确保安装对应版本的PyTorch和CUDA工具包。

1.2 PyCharm项目设置

打开PyCharm，选择"New Project"
设置项目位置，选择之前创建的虚拟环境

在项目根目录下创建以下文件夹结构：

├── datasets │ ├── your_dataset │ │ ├── normal │ │ └── abnormal ├── configs └── scripts

从GitHub克隆Anomalib源码：

git clone https://github.com/openvinotoolkit/anomalib.git

2. 理解Anomalib项目结构

Anomalib项目包含多个关键目录，理解这些结构对后续操作至关重要：

目录	作用	关键文件
anomalib/models	包含所有算法实现	padim, fastflow等子目录
anomalib/data	数据处理模块	datamodules, transforms
tools	训练和推理入口	train.py, inference/
docs	文档和示例	examples/

重点关注anomalib/models/padim目录下的内容：

config.yaml：模型配置文件
model.py：Padim算法实现
lightning_model.py：PyTorch Lightning封装

3. 准备自定义数据集

3.1 数据集结构要求

Padim算法对数据集结构有特定要求。假设我们要检测金属表面缺陷，数据集应这样组织：

datasets/ └── metal_defect/ ├── normal/ # 正常样本 │ ├── img1.jpg │ └── img2.jpg └── abnormal/ # 异常样本（可选，仅测试用） ├── defect1.jpg └── defect2.jpg

提示：训练阶段只需要normal文件夹中的样本，abnormal文件夹仅用于测试模型效果。

3.2 数据预处理技巧

工业图像常需要特殊预处理：

from anomalib.data.transforms import PreProcessor pre_process = PreProcessor( image_size=(256, 256), to_tensor=True, normalization="imagenet" ) # 应用预处理 processed_image = pre_process(image)

常见预处理参数说明：

image_size：统一调整图像尺寸
normalization：标准化方法（"imagenet"或"none"）
to_tensor：转换为PyTorch张量

4. 配置与训练Padim模型

4.1 修改配置文件

复制默认配置文件并修改关键参数：

# configs/padim_metal_defect.yaml model: name: padim backbone: resnet18 layers: - layer1 - layer2 - layer3 dataset: name: metal_defect format: folder path: ./datasets/metal_defect normal_dir: normal abnormal_dir: abnormal task: classification normalization: imagenet image_size: 256 metrics: image: - F1Score - AUROC optimization: export_mode: onnx

关键修改点：

将task设为classification（除非有像素级标注）
删除pixel相关指标（除非做分割任务）
添加export_mode: onnx以支持ONNX导出

4.2 启动模型训练

在PyCharm中创建运行配置：

选择tools/train.py为主脚本

设置参数：

--model padim --config configs/padim_metal_defect.yaml

或者直接运行命令：

python tools/train.py --model padim --config configs/padim_metal_defect.yaml

训练完成后，检查结果目录：

results/ └── padim └── metal_defect ├── images # 测试结果可视化 └── weights ├── model.ckpt └── onnx ├── model.onnx └── meta_data.json

5. ONNX模型导出与验证

5.1 模型导出配置

Anomalib会自动根据配置导出ONNX模型。如需手动导出，可使用：

from anomalib.deploy import export_onnx export_onnx( model_path="results/padim/metal_defect/weights/model.ckpt", export_path="custom_model.onnx", input_size=(1, 3, 256, 256) )

5.2 模型验证

使用ONNX Runtime进行推理验证：

import onnxruntime as ort import numpy as np # 加载模型 sess = ort.InferenceSession("results/padim/metal_defect/weights/onnx/model.onnx") # 准备输入数据 input_data = np.random.randn(1, 3, 256, 256).astype(np.float32) # 运行推理 outputs = sess.run(None, {"input": input_data}) # 解析输出 anomaly_map = outputs[0] # 异常热图 pred_score = outputs[1] # 异常分数

5.3 常见问题解决

形状不匹配错误：
- 确保输入图像尺寸与训练时一致
- 检查预处理是否与训练配置相同
精度下降问题：
- 比较PyTorch和ONNX模型的输出差异
- 检查ONNX opset版本（推荐使用opset=13）
内存不足错误：
- 减小批处理大小
- 使用onnxruntime-gpu加速

6. 高级技巧与优化建议

6.1 模型性能优化

通过OpenVINO进一步优化ONNX模型：

mo --input_model model.onnx --output_dir openvino_model --data_type FP16

优化后的模型可获得2-3倍的推理速度提升。

6.2 实际部署注意事项

预处理一致性：
- 确保部署环境的预处理与训练完全一致
- 保存预处理参数到meta_data.json
后处理实现：
- 异常分数标准化（0-1范围）
- 热图生成算法
阈值选择策略：
- 自适应阈值（如95%分位数）
- 基于验证集调整

6.3 扩展应用思路

Padim模型可应用于多种工业场景：

电子元件外观检测
纺织品瑕疵识别
食品品质监控
药品包装检查

每个场景可能需要调整：

图像分辨率
网络骨干（如改用resnet34）
异常判定阈值

7. PyCharm调试技巧

7.1 调试训练过程

在train.py中设置断点
检查DataLoader输出
监控损失函数变化

7.2 性能分析工具

使用PyCharm的Profiler分析瓶颈：

运行train.py性能分析
检查CPU/GPU利用率
优化数据加载流程

7.3 自定义回调

扩展训练过程：

from pytorch_lightning.callbacks import Callback class CustomCallback(Callback): def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): # 监控特定指标 pass

然后在配置中添加：