基于CNN的工件毛刺分类：从数据准备到ONNX部署全流程

简介这份资源是面向计算机、人工智能、自动化等专业学生与教师的深度学习实战项目包以CNN卷积神经网络为核心解决工业场景下工件毛刺的自动分类识别问题可作为毕业设计、课程设计或大作业的完整参考方案。压缩包共1294个文件约152.71MB其中1210个png与40个jpg构成图像数据集24个py脚本负责模型搭建与训练另有pth权重文件、csv标签文件、json配置及md说明文档覆盖从数据到推理的全流程。目前已有117人学习下载。项目代码均经过测试运行读者可拿到完整的网络结构定义、训练与预测脚本、已训练模型权重以及带标注的图像数据便于快速复现分类效果并在此基础上修改网络层数、更换数据集或迁移到其他缺陷检测任务适合入门进阶与项目立项演示。1. 工件毛刺分类为什么值得用 CNN 做一遍车间里最常见的质检场景之一就是金属件铣削、车削、冲压之后边缘残留的毛刺。毛刺大小不一、形状随机有的像细丝有的像鱼鳞有的干脆是一小块翻边。传统做法是老师傅拿放大镜看或者用接触式探针抽检效率低不说判定标准还随人波动。把「Python 深度学习 CNN」这套组合拳打到工件毛刺分类上本质是用卷积神经网络替代人眼做二分类或多分类有毛刺 / 无毛刺或者按毛刺严重程度分成轻微、中等、严重。它适合三类人机械设计制造及其自动化专业做毕业设计的学生、想入门深度学习但缺真实工业数据集的工程师、以及产线上想做视觉质检预研的技术员。整套方案的门槛并不高一台带独显的笔记本、一份标注好的工件端面图像、一份能跑通的 CNN 训练脚本就能把准确率做到 90% 以上。真正难的不是网络结构而是数据怎么来、怎么标、怎么让模型在换批次工件时不翻车。2. 从工件图像到 CNN 输入数据准备与增强的完整链路2.1 工件毛刺图像的数据集该怎么攒工业图像和 ImageNet 那种自然图像完全不是一回事。工件毛刺图像的特点是背景高度一致大多是金属灰、机床台面、目标区域小毛刺往往只占图像边缘几个像素宽、类间差异细微轻微毛刺和加工纹理容易混。我一般建议按下面的规格来攒数据项目建议值说明单类样本数不少于 300 张低于这个数 CNN 很容易过拟合图像分辨率原图 1280×960 以上后续统一缩放到 224×224拍摄距离固定工位、固定焦距避免尺度变化引入噪声光照环形 LED 白光避免反光金属表面反光是最大干扰源标注方式按图像打类别标签毛刺分类是图像级分类不需要框类别划分无毛刺 / 轻微 / 严重三分类比二分类更贴近产线需求如果手头没有真实产线图像可以用公开的机械加工表面缺陷数据集做迁移或者自己拿手机在台钳上拍。注意一点同一批工件拍出来的图不能全放进训练集否则验证集准确率虚高这是毕业设计答辩时最容易被老师问倒的地方。2.2 用 Python 做数据清洗与增强拿到原始图像后第一步不是直接喂网络而是清洗。下面这段脚本做三件事统一尺寸、剔除模糊图、按 7:2:1 划分数据集。import os import cv2 import numpy as np import shutil from sklearn.model_selection import train_test_split RAW_DIR raw_images # 原始图像按类别放在子文件夹 OUT_DIR dataset # 输出目录 IMG_SIZE 224 # CNN 标准输入尺寸 BLUR_THRESHOLD 100.0 # 拉普拉斯方差阈值低于此值判为模糊 def is_blurry(img_path): img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) if img is None: return True return cv2.Laplacian(img, cv2.CV_64F).var() BLUR_THRESHOLD def prepare(): classes os.listdir(RAW_DIR) for split in [train, val, test]: for c in classes: os.makedirs(os.path.join(OUT_DIR, split, c), exist_okTrue) for c in classes: c_dir os.path.join(RAW_DIR, c) files [f for f in os.listdir(c_dir) if f.lower().endswith((.jpg, .png, .bmp))] valid [] for f in files: p os.path.join(c_dir, f) if is_blurry(p): continue img cv2.imread(p) img cv2.resize(img, (IMG_SIZE, IMG_SIZE)) valid.append((f, img)) names [v[0] for v in valid] imgs [v[1] for v in valid] # 先 7:3 切训练和临时集再从临时集 2:1 切验证和测试 X_train, X_tmp, n_train, n_tmp train_test_split( imgs, names, test_size0.3, random_state42, stratifyNone) X_val, X_test, n_val, n_test train_test_split( X_tmp, n_tmp, test_size0.33, random_state42) for split, data, names_ in [(train, X_train, n_train), (val, X_val, n_val), (test, X_test, n_test)]: for img, name in zip(data, names_): cv2.imwrite(os.path.join(OUT_DIR, split, c, name), img) if __name__ __main__: prepare()逻辑说明is_blurry用拉普拉斯算子方差判断图像清晰度金属件拍摄时手抖或对焦不准的图必须剔除否则模型会学到「模糊有毛刺」这种伪特征。train_test_split做了两次切分最终比例约 7:2:1。参数BLUR_THRESHOLD100是经验值如果你的图像整体偏暗可以降到 60 左右IMG_SIZE224是为了后续直接套用预训练模型如果自己从零搭网络128 也够用。2.3 在线增强让 300 张图发挥 3000 张的效果工件图像不能随便做水平翻转因为毛刺方向有工艺含义。我一般只用以下几种增强随机旋转 ±15 度、亮度抖动 ±20%、高斯噪声、随机裁剪 90% 区域再缩放回去。用torchvision.transforms写出来是这样from torchvision import transforms train_tf transforms.Compose([ transforms.RandomRotation(15), # 小幅旋转模拟工件摆放偏差 transforms.ColorJitter(brightness0.2, contrast0.2), # 光照波动 transforms.RandomResizedCrop(224, scale(0.85, 1.0)), # 随机裁剪 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], # ImageNet 统计值 std[0.229, 0.224, 0.225]), transforms.RandomErasing(p0.1) # 随机遮挡提升鲁棒性 ]) val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意RandomErasing的概率不要超过 0.2否则毛刺区域被遮掉模型学不到关键特征。Normalize用的均值和方差是 ImageNet 的统计值如果你完全从零训练、不用预训练权重可以改成自己数据集的均值和方差但差别通常不大。3. CNN 模型搭建从零写还是拿预训练权重改3.1 两种路线的选型对比毕业设计里最常见的纠结是自己搭一个 CNN还是用 ResNet18 这类预训练模型微调我的建议很直接——数据少于 1000 张一律用预训练模型微调数据超过 5000 张且和 ImageNet 差异极大才考虑从零训练。工件毛刺图像通常只有几百到一两千张所以走迁移学习路线。路线适用数据量训练时间预期准确率代码复杂度从零搭 4 层 CNN 500短75%~85%低ResNet18 微调500~5000中90%~96%中EfficientNet-B0 微调1000~10000中长92%~97%中从零搭 强增强 5000长88%~94%高3.2 用 PyTorch 搭一个可复现的基线模型下面这个模型结构是我在多个工件分类任务里反复用过的基线主干用 ResNet18 预训练权重把最后的全连接层换成三分类输出前面几层冻结只训练后两层和分类头。import torch import torch.nn as nn from torchvision import models def build_model(num_classes3, freeze_backboneTrue): model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) # 冻结前面所有层只训练 layer4 和 fc if freeze_backbone: for name, param in model.named_parameters(): if not name.startswith(layer4) and not name.startswith(fc): param.requires_grad False # 替换分类头 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, 128), nn.ReLU(inplaceTrue), nn.Linear(128, num_classes) ) return model if __name__ __main__: net build_model(num_classes3) total sum(p.numel() for p in net.parameters()) trainable sum(p.numel() for p in net.parameters() if p.requires_grad) print(f总参数: {total}, 可训练参数: {trainable})逻辑说明freeze_backboneTrue时只训练layer4和fc可训练参数大约占总参数的 30% 左右这样在小数据集上不容易过拟合。Dropout(0.3)放在全连接前是防止分类头记住训练样本的简单手段。如果你的数据量超过 2000 张可以把freeze_backbone设为False全网络微调但学习率要降到 1e-4 以下。3.3 训练循环与关键超参数训练脚本里最容易被忽视的是类别不平衡处理。工件毛刺数据里「无毛刺」样本往往远多于「严重毛刺」直接训练会让模型偏向多数类。解决办法是在损失函数里加权重。import torch from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR def train(): device torch.device(cuda if torch.cuda.is_available() else cpu) train_ds ImageFolder(dataset/train, transformtrain_tf) val_ds ImageFolder(dataset/val, transformval_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) # 按类别频率计算权重缓解不平衡 counts [0] * len(train_ds.classes) for _, label in train_ds.samples: counts[label] 1 weights torch.tensor([1.0 / c for c in counts], dtypetorch.float32).to(device) weights weights / weights.sum() * len(counts) model build_model(num_classeslen(train_ds.classes)).to(device) criterion torch.nn.CrossEntropyLoss(weightweights) optimizer AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr3e-4, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30) best_acc 0.0 for epoch in range(30): model.train() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() loss criterion(model(imgs), labels) loss.backward() optimizer.step() scheduler.step() # 验证 model.eval() correct, total 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs).argmax(dim1) correct (pred labels).sum().item() total labels.size(0) acc correct / total print(fEpoch {epoch1}, Val Acc: {acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) if __name__ __main__: train()参数说明lr3e-4是微调预训练模型的常用起点如果 loss 震荡明显就降到 1e-4weight_decay1e-4是 AdamW 的默认正则强度T_max30要和总 epoch 数一致余弦退火才能完整走完一个周期batch_size32在 6GB 显存上跑 224×224 的 ResNet18 刚好够用显存不够就降到 16。类别权重那段代码是很多毕业设计忽略的但它在三分类不平衡数据上通常能带来 3~5 个百分点的提升。4. 训练完不算完模型评估、导出与推理部署4.1 别只看准确率混淆矩阵和每类召回率工件毛刺分类里把「严重毛刺」误判成「无毛刺」的代价远大于反过来。所以评估时必须看每类召回率而不是总体准确率。import torch import numpy as np from sklearn.metrics import classification_report, confusion_matrix from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder def evaluate(): device torch.device(cuda if torch.cuda.is_available() else cpu) test_ds ImageFolder(dataset/test, transformval_tf) test_loader DataLoader(test_ds, batch_size32, shuffleFalse) model build_model(num_classeslen(test_ds.classes)).to(device) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in test_loader: imgs imgs.to(device) preds model(imgs).argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(confusion_matrix(all_labels, all_preds)) print(classification_report(all_labels, all_preds, target_namestest_ds.classes, digits4)) if __name__ __main__: evaluate()classification_report会输出每类的 precision、recall、f1-score。重点看「严重毛刺」那一行的 recall如果低于 0.85说明模型对这类样本学得不够需要回头补数据或者调类别权重。混淆矩阵能直观看出哪两类最容易混通常是「轻微毛刺」和「无毛刺」之间的边界模糊。4.2 导出 ONNX 并在无 Python 环境推理毕业设计答辩时经常需要现场演示而答辩教室的电脑不一定装了 PyTorch。把模型导出成 ONNX 格式用onnxruntime推理依赖只有几十 MB。import torch def export_onnx(): model build_model(num_classes3) model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) model.eval() dummy torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy, burr_classifier.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11 ) print(导出完成) if __name__ __main__: export_onnx()导出后用下面这段代码做单张图推理不依赖 PyTorchimport onnxruntime as ort import cv2 import numpy as np def predict(img_path): sess ort.InferenceSession(burr_classifier.onnx) img cv2.imread(img_path) img cv2.resize(img, (224, 224)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 mean np.array([0.485, 0.456, 0.406], dtypenp.float32) std np.array([0.229, 0.224, 0.225], dtypenp.float32) img (img - mean) / std img np.transpose(img, (2, 0, 1))[None, ...] out sess.run(None, {input: img})[0] cls int(np.argmax(out)) return [无毛刺, 轻微毛刺, 严重毛刺][cls], out if __name__ __main__: label, logits predict(test_sample.jpg) print(label, logits)注意预处理必须和训练时完全一致包括 BGR 转 RGB、归一化参数、通道顺序。ONNX 推理最常见的翻车就是预处理对不上导致线上结果和验证结果差十几个百分点。5. 避坑与排查工件毛刺分类里最容易翻车的 5 个地方5.1 验证集准确率 98%换一批工件掉到 60%现象训练日志里 val acc 一路涨到 0.98但拿新拍的工件图测试模型几乎全判成同一类。原因训练集和验证集来自同一批拍摄背景、光照、工件型号完全一致模型学到了「背景特征」而不是「毛刺特征」。解决按工件批次划分数据集而不是随机划分。具体做法是给每张图打上批次号同一批次的图只出现在 train 或 val 其中一个里。如果数据量不够至少保证验证集里有不同光照条件下拍的图。5.2 Loss 不下降准确率卡在 33%现象三分类任务训练 loss 从 1.1 降到 1.09 就不动了准确率始终在 0.33 附近。原因学习率太大预训练权重被破坏或者数据标签有问题比如所有图都被标成了同一类。解决先把学习率降到 1e-5 试一个 epoch如果 loss 开始降说明是学习率问题。如果还是不动用ImageFolder打印几个 batch 的标签分布确认标签没有全错。我遇到过最离谱的一次是文件夹名字写错三个类别的图全在一个文件夹里。5.3 显存溢出CUDA out of memory现象训练到第几个 batch 突然报RuntimeError: CUDA out of memory。原因batch_size太大或者num_workers太多导致内存泄漏或者验证时没有加torch.no_grad()。解决先把batch_size减半再把num_workers设为 2 或 0 试试。验证循环一定要包在with torch.no_grad():里否则计算图会一直累积。如果还不行用torch.cuda.empty_cache()在每个 epoch 结束后清一次缓存。5.4 模型对光照变化极度敏感现象同一工件换个角度打光预测结果就变了。原因训练数据光照太单一模型把亮度当成了判别特征。解决在增强里加大亮度抖动范围ColorJitter(brightness0.4, contrast0.3)同时补拍不同光照条件下的样本。如果条件允许在预处理里加一步直方图均衡化把光照影响压下去。5.5 导出 ONNX 后推理结果和 PyTorch 不一致现象PyTorch 里预测是「严重毛刺」ONNX 推理出来是「无毛刺」。原因预处理不一致最常见的是 PyTorch 用 RGB、OpenCV 读进来是 BGR忘了转换或者归一化参数写错。解决写一个对比脚本同一张图分别走 PyTorch 和 ONNX打印预处理后的张量逐像素对比。只要预处理对齐ONNX 和 PyTorch 的输出差异通常在 1e-5 以内。6. 把模型推到产线边缘一个可落地的推理加速技巧前面训练和导出都跑通之后最后一个现实问题是速度。产线节拍通常要求单张图推理在 100ms 以内而 ResNet18 在 CPU 上跑 224×224 大约要 200~300ms。我一般用两个手段压下来一是把输入分辨率从 224 降到 160工件毛刺这种边缘特征在 160 下仍然保留得不错二是用 ONNX Runtime 的量化功能做 INT8 推理。from onnxruntime.quantization import quantize_dynamic, QuantType def quantize(): quantize_dynamic( model_inputburr_classifier.onnx, model_outputburr_classifier_int8.onnx, weight_typeQuantType.QUInt8 ) print(量化完成) if __name__ __main__: quantize()量化后模型体积大约缩小到原来的四分之一CPU 推理速度提升 2~3 倍。代价是准确率可能掉 1~2 个百分点需要重新在测试集上验证一遍。如果掉得太多可以只量化卷积层保留全连接层为 FP32ONNX Runtime 支持按节点指定量化类型。另一个容易被忽视的点是输入尺寸和模型结构的匹配。如果你把输入改成 160×160而模型里还有自适应池化层通常不用改代码就能跑但如果模型里有硬编码的全连接输入维度就需要重新导出。我自己的习惯是训练时就用最终部署的分辨率避免训练和推理两套预处理。最后说一个验证技巧拿 20 张产线实拍图人工标好标签分别用 PyTorch 原模型、ONNX FP32、ONNX INT8 跑一遍记录三者的预测一致率和单张耗时。如果 INT8 和 FP32 的一致率低于 95%就放弃量化改用降分辨率方案。这套对比流程我每次上线前都会走一遍比拍脑袋调参靠谱得多。希望帮到你。本文还有配套的精品资源点击获取