摘要：钢材是现代制造业重要基础原料，生产过程中会产生裂纹、夹杂、斑块、点蚀、氧化皮、划痕等多种表面缺陷，不同缺陷对应不同生产故障成因，准确识别缺陷类型对工艺改进、产品质量管控意义重大。本文基于Pytorch搭建EfficientNet(b0‑b7)全套工业缺陷分类工程，数据集1700余张钢材表面图像，完成crazing、inclusion、patches、pitted_surface、rolled‑in_scale、scratches共6类缺陷识别。工程支持迁移学习、冻结主干训练、参数量与计算量统计、全套多分类评价指标、混淆矩阵、多指标曲线可视化、批量图片推理；附带可直接渲染的Mermaid网络结构图。本文从研究背景、数据集、EfficientNet原理、网络结构（Mermaid代码）、代码模块解析、训练策略、指标解读、运行流程、不足与优化方向完整展开。
1 研究背景与意义
钢铁板材在轧制、热处理生产过程，受原料、设备、工艺参数影响，会产生各式各样的表面缺陷。
- crazing（裂纹）：表面出现网状开裂纹路；
- inclusion（夹杂）：外来杂质嵌入钢材表面；
- patches（斑块）：局部不规则色斑斑块；
- pitted_surface（点蚀表面）：密集点状凹坑损伤；
- rolled‑in_scale（氧化皮）：轧制带入氧化铁皮；
- scratches（划痕）：长条状摩擦刮伤痕迹。
传统工厂大多依靠人工肉眼分拣，长时间工作后工人疲劳，会造成漏检错检；另外人工检测速度有限，难以匹配高速流水线生产节拍。机器视觉结合深度学习可以24小时不间断工作，自动区分不同缺陷种类，不仅判断有无缺陷，还可以区分缺陷类别，反向指导生产工艺调整。
EfficientNet是Google提出的轻量化高性能CNN，区别于以往单独增加网络深度或者通道数的做法，提出复合缩放（Compound Scaling） ，使用一组复合系数，同时缩放网络深度、网络宽度、输入图像分辨率，在计算量和识别精度之间取得很好平衡，b0为基础最小模型，b1b7逐级放大模型规模。工业场景硬件配置差异大，可以根据算力选择b0b7不同版本，算力弱选b0，追求精度选b7。
本数据集一共1700多张钢材缺陷图像，属于中小规模工业数据集，非常适合使用迁移学习，利用ImageNet预训练学到的边缘、纹理特征，迁移到钢材缺陷纹理识别任务。
2 数据集说明
数据集6类钢材表面缺陷，标签json定义：
{
"0": "crazing",
"1": "inclusion",
"2": "patches",
"3": "pitted_surface",
"4": "rolled-in_scale",
"5": "scratches"
}数据集采用ImageFolder目录格式组织：
data
├─train
│ ├─crazing
│ ├─inclusion
│ ├─patches
│ ├─pitted_surface
│ ├─rolled-in_scale
│ └─scratches
└─val
├─crazing
├─inclusion
└─......其余缺陷类别数据预处理流水线
- 训练集：Resize(256) → RandomRotation(90)随机旋转 → CenterCrop(224) → ToTensor() → ImageNet标准化
mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225] - 验证集：Resize(256) → CenterCrop(224) → ToTensor() → 同样标准化。
说明：钢材缺陷很多是方向性纹理（划痕、裂纹），随机旋转增强可以提升模型对缺陷方向的鲁棒性；本项目使用ImageNet预训练权重，因此沿用ImageNet均值方差。工业实际部署时，如果图像灰度分布差异巨大，可以重新统计数据集自身均值方差。
3 EfficientNet网络整体架构与Mermaid结构图
EfficientNet核心基础模块是MBConv（Mobile Inverted Bottleneck），融合深度可分离卷积、倒残差结构、SE通道注意力。
- 1×1升维卷积提升通道；
- 3×3深度可分离卷积做空间特征提取；
- SE注意力模块对通道特征重标定；
- 1×1降维卷积压缩通道；
- 残差跳跃连接，输入输出相加。
3.1 整体网络文字结构
输入图像 3×224×224
↓
Stem层：3×3 Conv + BN + Swish激活
↓
重复堆叠多组MBConv模块（不同b版本深度、通道数不一样）
↓
Head层：1×1 Conv升维 + BN + Swish
↓
Global Average Pooling全局平均池化
↓
Dropout随机失活
↓
Linear分类头 → 输出6维（6类钢材缺陷）3.2 MBConv模块内部子结构
MBConv输入特征图
↓
1×1 Pointwise卷积（升通道）+ BN + Swish
↓
k×k Depthwise深度卷积 + BN + Swish
↓
SE 通道注意力模块（全局池化‑>两层全连接‑>sigmoid通道权重）
↓
1×1 Pointwise卷积（降通道）+ BN
↓
Dropout
↓
残差shortcut（输入与输出相加，步长=1且通道匹配时启用）3.3 Mermaid绘图代码
3.4 b0~b7版本差异说明
| 模型 | 缩放系数 | 参数量 | 特点 |
|---|---|---|---|
| EfficientNet‑b0 | 基础基准 | 最小 | 适合嵌入式、低算力设备，快速实验 |
| EfficientNet‑b1 | 复合缩放1 | 小 | 精度略高于b0，速度较快 |
| EfficientNet‑b2 | 复合缩放2 | 中小 | 工业项目常用折中选择 |
| EfficientNet‑b3 | 复合缩放3 | 中等 | 精度进一步提升 |
| EfficientNet‑b4 | 复合缩放4 | 中高 | 算力充足推荐 |
| EfficientNet‑b5 | 复合缩放5 | 偏大 | 高精度，显存占用高 |
| EfficientNet‑b6 | 复合缩放6 | 大 | 适合离线检测场景 |
| EfficientNet‑b7 | 复合缩放7 | 最大 | 追求最高识别精度 |
本项目代码通过--model b0/b1/b2...b7参数切换模型；加载官方预训练权重，替换最后classifier层输出维度为6，适配6类钢材缺陷；开启--freeze_layers True可以冻结主干网络，仅训练最后的分类头，保护预训练特征，适合1700张这样的中小数据集。
4 项目核心代码模块解读
项目包含三份核心脚本：train.py训练脚本、utils.py工具库、predict.py批量推理脚本。
4.1 utils.py核心模块
- create_model()：模型工厂函数，根据model参数选择b0‑b7，加载对应官方预训练权重，替换分类头输出通道等于缺陷类别数量。
get_device()自动获取cuda/cpu运行设备；mkdir()自动清空旧实验目录，新建weights权重文件夹；data_trans()训练集验证集两套数据增强。- ConfusionMatrix混淆矩阵类：计算多分类全套指标：准确率、每一类Precision精确率、Recall召回率、Specificity特异度、F1‑score，输出宏平均指标；自动绘制保存混淆矩阵图片。>
工业缺陷检测重点指标：Recall召回率，代表缺陷被找出来的比例，缺陷检测场景希望召回率尽可能高，降低漏检；Precision精确率控制把正常表面误判为缺陷的误检情况。混淆矩阵可以直观看到哪两类缺陷容易混淆，例如patches斑块和pitted_surface点蚀容易发生错分。
train_one_epoch()：一轮训练，CrossEntropy交叉熵损失，tqdm进度条实时打印loss、acc；evaluate()验证集评估。- 绘图工具函数：
plt_loss_acc()损失精度曲线；plot_lr_decay()学习率变化曲线；plot_prsf()每epoch宏平均Precision、Recall、Specificity、F1曲线。 get_data()读取ImageFolder数据集；gen_json()生成class_indices.json类别映射；get_optim()支持SGD、Adam、AdamW三种优化器。
4.2 train.py训练脚本流程
- 接收命令行参数，自定义实验保存目录save_ret；调用mkdir清空旧文件；获取训练设备。
- 加载训练、验证数据集，自动生成类别json文件，打印训练集、验证集样本数量。
- 创建EfficientNet模型，thop工具统计总参数量、可训练参数量、FLOPs计算量，控制台打印输出。
- freeze_layers开启时，冻结主干全部参数，只训练classifier分类头。
- 筛选可更新梯度参数pg；选择优化器；设置LambdaLR余弦退火学习率调度器。
- epoch循环：
- 执行train_one_epoch训练一轮；scheduler更新学习率；evaluate做验证；
- 每epoch全部指标保存到
train_log_results.json； - 对比验证集accuracy，保存最优
best.pth；训练结束保存last.pth。
- 训练结束自动绘制loss‑acc、lr、P‑R‑S‑F1曲线图。
- 加载best最优权重，分别计算、绘制训练集、验证集混淆矩阵并保存图片。
参考训练超参数
--model b0
--pretrained True
--freeze_layers True
--batch‑size 32
--epochs 100
--optim SGD
--lr 0.01
--lrf 0.00014.3 predict.py批量推理脚本
- 读取指定实验目录下
class_indices.json类别映射、best.pth权重； - 读取
./inference/infer_img文件夹下全部待检测钢材图片； - 图像预处理，模型推理，softmax得到概率，取top‑3预测结果；
- OpenCV将缺陷名称+概率绘制到图像上；
- 自动创建
inference/ret文件夹，保存绘制完成的结果图片。
运行示例：
python predict.py --model b0 --save_ret runs --weights best.pth5 多分类评价指标解读（钢材表面6类缺陷）
本任务6分类，代码采用宏平均macro‑average，平等对待6种缺陷类别，不受各类样本数量不均衡干扰。
- Accuracy准确率：全部样本预测正确占总样本比例；
- Precision精确率：预测为某一类缺陷样本中真实属于该缺陷的比例，反映误检；
- Recall召回率：真实该类缺陷样本被识别出来的比例，缺陷检测最重要指标，越低代表漏检越多；
- Specificity特异度：不属于该缺陷的样本被正确判定为非该缺陷；
- F1‑score：精确率、召回率调和平均，综合评价该类别识别性能；
- 混淆矩阵：6×6矩阵，对角线为各类别预测正确数量；非对角线代表类别之间错分，用来定位模型容易混淆的缺陷类型。
过拟合判断：训练集指标很高，验证集指标显著偏低。工业小数据集很容易出现过拟合；缓解手段：开启冻结训练、增加数据增强、适当降低epoch、使用权重衰减。
6 完整项目运行流程
-
环境依赖
pytorch、torchvision、opencv‑python、pillow、matplotlib、thop、tqdm -
数据集按照ImageFolder格式放到
./data/train与./data/val。 -
启动训练示例：
python train.py --model b0 --pretrained True --freeze_layers True --batch-size 32 --epochs 100 --optim SGD --lr 0.01 --lrf 0.0001
-
训练产物全部输出到
‑‑save_ret指定目录（默认runs）
- 权重文件weights/best.pth、last.pth
- class_indices.json类别映射
- train_log_results.json完整训练日志
- loss_accuracy_curve.jpg、LR_decay.jpg、results.jpg各类指标曲线
- trainSet Confusion Matrix.jpg、valSet Confusion Matrix.jpg混淆矩阵
- 批量缺陷推理：将钢材图像放入
inference/infer_img，执行predict.py，结果输出至inference/ret。