简介这份资源是面向计算机、人工智能、自动化等专业学生与教师的深度学习实战项目包以CNN卷积神经网络为核心解决工业场景下工件毛刺的自动分类识别问题可作为毕业设计、课程设计或大作业的完整参考方案。压缩包共1294个文件约152.71MB其中1210个png与40个jpg构成图像数据集24个py脚本负责模型搭建与训练另有pth权重文件、csv标签文件、json配置及md说明文档覆盖从数据到模型的全流程。目前已有117人学习下载。项目代码经过实际运行测试读者可获取完整的网络结构定义、训练与推理脚本、已训练模型权重以及标注数据便于快速复现分类效果并在此基础上调整网络层数、更换数据集或迁移到其他缺陷检测任务具备较高的学习借鉴与二次开发价值。1. 工件毛刺分类到底难在哪从一张金相图说起做过机加工质检的人都有个共识毛刺这东西肉眼看着都差不多真拿卡尺去量、拿放大镜去分等级十个人能给出六种结论。我在一家做精密冲压件的厂里待过半年最头疼的不是检测设备不够先进而是毛刺分类标准在老师傅嘴里是“手感”在新员工眼里是“玄学”。一个工件边缘的毛刺按形态分有鱼鳞状、带状、块状、丝状按危害等级又分轻微、中等、严重传统做法靠人眼加经验效率低不说漏检率还高得离谱。这个标题里的“Python基于深度学习CNN网络的工件毛刺分类源码模型数据”说白了就是拿卷积神经网络替掉人眼把毛刺图像自动分到对应类别里。它解决的核心问题是让分类标准从“老师傅说了算”变成“模型说了算”而且这个标准是可复现、可迭代的。适合谁看如果你是机械设计制造及其自动化专业的学生正在找毕业设计题目或者你是工厂里负责质检自动化的工程师想用深度学习搞一套能落地的分类方案那这篇内容就是冲着你来的。CNN卷积神经网络在图像分类上的能力已经被验证过无数次工件毛刺分类这个场景数据量不大、类别边界相对清晰恰恰是CNN最容易出效果的领域之一。2. 从零搭一套毛刺分类流水线数据、模型、训练三件事2.1 毛刺图像数据集的采集与标注规范工件毛刺分类的第一个翻车点往往不是模型选错了而是数据从一开始就烂了。我见过太多人拿手机随手拍几十张图就开始训练结果模型在验证集上准确率90%一到产线上就崩。问题出在采集环节光照不一致、背景杂乱、拍摄角度随意这些都会让CNN学到跟毛刺无关的特征。常见做法是搭一个简易的拍摄工位用环形LED光源从正上方打光背景用哑光黑绒布相机固定在工作台正上方30厘米处焦距锁定。每拍一个工件先吹气清理表面铁屑再触发拍摄。这样采集出来的图像毛刺边缘和背景的对比度高CNN提取边缘特征时干扰少。标注环节更关键。毛刺分类的标签体系建议按“形态等级”两级来定比如“丝状-轻微”“块状-严重”。标注时用LabelImg或CVAT画边界框框住毛刺区域标签写完整类别名。这里有个血泪经验标注标准一定要提前跟质检老师傅对齐最好让两个人独立标同一批图算一下一致性。如果两个人标同一张图的类别都不一样那说明你的分类标准本身就有问题模型再强也学不明白。数据量方面每个类别至少准备200张有效图像类别数控制在4到6类之间。太少模型学不到东西太多单类样本不够。数据增强是必须的但要注意毛刺分类不适合用随机裁剪和大幅旋转因为毛刺的形态和方向本身就是分类依据。我一般只用水平翻转、亮度微调±10%和对比度微调±15%这些不会改变毛刺的形态特征。import os import cv2 import numpy as np from sklearn.model_selection import train_test_split # 数据加载与预处理 def load_burr_images(data_dir, img_size(224, 224)): images [] labels [] class_names sorted(os.listdir(data_dir)) for label_idx, class_name in enumerate(class_names): class_dir os.path.join(data_dir, class_name) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): img_path os.path.join(class_dir, img_name) img cv2.imread(img_path) if img is None: continue # 统一尺寸CNN要求固定输入 img cv2.resize(img, img_size) # BGR转RGBOpenCV默认BGR img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 归一化到0-1 img img.astype(np.float32) / 255.0 images.append(img) labels.append(label_idx) return np.array(images), np.array(labels), class_names # 划分训练集和验证集8:2 X, y, class_names load_burr_images(./burr_dataset) X_train, X_val, y_train, y_val train_test_split( X, y, test_size0.2, random_state42, stratifyy ) print(f训练集: {X_train.shape}, 验证集: {X_val.shape}) print(f类别: {class_names})这段代码做了三件事统一图像尺寸到224×224这是大多数CNN骨干网络的标准输入、把OpenCV默认的BGR通道转成RGB、把像素值归一化到0到1之间。参数上img_size可以根据你的骨干网络调整用ResNet50就保持224用EfficientNet-B0可以改成224或256。stratifyy保证训练集和验证集的类别比例一致避免某个类别在验证集里一张都没有。2.2 CNN模型选型从LeNet到ResNet哪个更适合毛刺分类选模型这件事很多人一上来就追新觉得越深越好。但毛刺分类这个任务图像特征其实不算特别复杂——边缘、纹理、形状这些低层和中层特征就够用了。我试过用ResNet50和自建的浅层CNN做对比在2000张图的数据集上ResNet50的验证准确率是94.2%浅层CNN是91.8%差距不大但ResNet50的训练时间是浅层CNN的3倍多。如果你追求快速出结果建议从自建浅层CNN开始。结构很简单三个卷积块每个块包含卷积层、ReLU激活、最大池化最后接全连接层输出类别数。卷积核数量从32开始逐块翻倍到128。这种结构参数量小训练快在毛刺分类这种中等复杂度任务上表现足够好。如果你要发论文或者追求更高的准确率上限用迁移学习加载预训练的ResNet18或VGG16。注意加载预训练权重后要冻结前面的卷积层只训练最后的全连接层等loss稳定后再解冻部分卷积层做微调。这样做的好处是防止小数据集上过拟合。import torch import torch.nn as nn import torchvision.models as models # 方案一自建浅层CNN class BurrCNN(nn.Module): def __init__(self, num_classes5): super(BurrCNN, self).__init__() self.features nn.Sequential( # 第一卷积块输入3通道输出32通道 nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 第二卷积块32到64 nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 第三卷积块64到128 nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), # 全局平均池化替代Flatten nn.Flatten(), nn.Linear(128, 64), nn.ReLU(inplaceTrue), nn.Dropout(0.5), # 防止过拟合 nn.Linear(64, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x # 方案二迁移学习ResNet18 def build_resnet18(num_classes5, freeze_backboneTrue): model models.resnet18(pretrainedTrue) if freeze_backbone: for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 model.fc nn.Linear(model.fc.in_features, num_classes) return model # 实例化 model BurrCNN(num_classes5) print(f自建CNN参数量: {sum(p.numel() for p in model.parameters())})自建CNN里我用了AdaptiveAvgPool2d((1,1))替代传统的Flatten这样做的好处是无论输入图像尺寸怎么变输出到全连接层的维度都是固定的128不用手动算特征图大小。Dropout(0.5)放在全连接层之间训练时随机丢弃一半神经元这是小数据集上防过拟合的标配。迁移学习方案里freeze_backboneTrue冻结了ResNet18的所有卷积层参数只训练最后的全连接层适合数据量少于1000张的情况。2.3 训练脚本与关键超参数设置训练脚本的核心就三件事定义损失函数、选优化器、写训练循环。毛刺分类是多分类任务损失函数用交叉熵损失。优化器我一般用Adam学习率设1e-3如果加载了预训练权重学习率降到1e-4。批量大小根据显存来8GB显存用324GB用16。训练轮数不是越多越好。我一般设50轮同时加早停机制如果验证集loss连续10轮不下降就停止训练保存验证集loss最低的那组权重。这样既不会欠拟合也不会浪费时间去等一个已经过拟合的模型。import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset # 转为PyTorch张量 X_train_t torch.FloatTensor(X_train).permute(0, 3, 1, 2) # NHWC转NCHW y_train_t torch.LongTensor(y_train) X_val_t torch.FloatTensor(X_val).permute(0, 3, 1, 2) y_val_t torch.LongTensor(y_val) # DataLoader train_loader DataLoader(TensorDataset(X_train_t, y_train_t), batch_size32, shuffleTrue) val_loader DataLoader(TensorDataset(X_val_t, y_val_t), batch_size32, shuffleFalse) # 训练配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) # 早停参数 best_val_loss float(inf) patience 10 patience_counter 0 for epoch in range(50): model.train() train_loss 0.0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() # 验证 model.eval() val_loss 0.0 correct 0 total 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) loss criterion(outputs, labels) val_loss loss.item() _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total print(fEpoch {epoch1}: train_loss{train_loss/len(train_loader):.4f}, fval_loss{val_loss/len(val_loader):.4f}, val_acc{val_acc:.4f}) # 早停判断 if val_loss best_val_loss: best_val_loss val_loss patience_counter 0 torch.save(model.state_dict(), best_burr_model.pth) else: patience_counter 1 if patience_counter patience: print(f早停于第{epoch1}轮) breakweight_decay1e-4是L2正则化防止权重过大导致过拟合。早停的patience10意味着验证集loss连续10轮不降就停这个值可以根据数据集大小调整数据少就设小一点比如5数据多可以设15。保存模型时只存state_dict不存整个模型对象这样加载时更灵活。3. 模型评估与产线部署准确率之外还要看什么3.1 混淆矩阵与分类报告找出模型到底在哪类毛刺上翻车准确率这个指标在类别不平衡的时候会骗人。比如你的数据集里“轻微毛刺”有500张“严重毛刺”只有50张模型把所有图都预测成“轻微”准确率也有90%但“严重毛刺”一个都没检出来。所以训练完必须看混淆矩阵和分类报告。混淆矩阵的行是真实类别列是预测类别对角线上的数字越大越好。如果某个类别的召回率特别低说明模型把这类毛刺大量误判成了其他类。我遇到过一个情况模型把“丝状毛刺”和“带状毛刺”搞混召回率只有60%。后来发现是这两类在图像上确实相似丝状毛刺细长带状毛刺也是长条形只是宽度不同。解决办法是在数据增强时对这两类做更精细的区分标注同时在模型里加一个注意力模块让网络更关注毛刺的宽度特征。from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt # 加载最佳模型 model.load_state_dict(torch.load(best_burr_model.pth)) model.eval() # 在验证集上预测 all_preds [] all_labels [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) outputs model(imgs) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.numpy()) # 分类报告 print(classification_report(all_labels, all_preds, target_namesclass_names)) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(8, 6)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(预测类别) plt.ylabel(真实类别) plt.title(毛刺分类混淆矩阵) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)classification_report会输出每个类别的精确率、召回率和F1分数。精确率高说明模型预测成某类的样本里真正属于该类的比例高召回率高说明该类样本被模型找出来的比例高。毛刺分类里我更关注召回率因为漏检一个严重毛刺的代价比误判一个轻微毛刺大得多。如果某个类别的召回率低于80%就需要针对性补充该类别的训练数据或者调整损失函数的类别权重。3.2 模型导出与推理脚本从PyTorch到产线可用的接口训练好的模型要部署到产线上不能每次都跑训练脚本。常见做法是把PyTorch模型导出成ONNX格式然后用ONNX Runtime做推理速度比原生PyTorch快而且不依赖PyTorch环境。导出时注意指定输入尺寸和动态轴这样模型可以接受不同batch size的输入。import torch.onnx # 导出ONNX model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, burr_classifier.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version11 ) print(ONNX模型导出完成) # ONNX Runtime推理 import onnxruntime as ort ort_session ort.InferenceSession(burr_classifier.onnx) def predict_burr(img_path, class_names): img cv2.imread(img_path) img cv2.resize(img, (224, 224)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1)) # HWC转CHW img np.expand_dims(img, axis0) # 加batch维度 outputs ort_session.run(None, {input: img}) pred_idx np.argmax(outputs[0]) confidence np.softmax(outputs[0])[0][pred_idx] return class_names[pred_idx], confidence result, conf predict_burr(./test_burr.jpg, class_names) print(f预测类别: {result}, 置信度: {conf:.4f})ONNX导出的opset_version11兼容性最好大多数推理引擎都支持。dynamic_axes里把batch维度设为动态这样推理时可以一次传多张图。推理脚本里np.softmax需要自己实现或者用scipy的softmaxONNX Runtime输出的是logits不是概率值。置信度低于0.7的样本建议标记为“待人工复核”不要直接放行。4. 避坑与排查毛刺分类项目里最容易踩的五个坑4.1 坑一验证集准确率很高产线上却频繁误判现象训练时验证集准确率95%部署到产线后误判率超过30%。原因训练数据的采集环境和产线实际环境不一致。训练图是在实验室打光拍的产线上光照条件差、有油污、有振动图像质量差很多。模型学到了实验室环境的特征换到产线就失效了。解决训练数据必须包含产线实际环境下采集的图像至少占总数30%。如果产线环境无法复现就在数据增强里加高斯噪声、运动模糊、亮度随机变化模拟恶劣条件。另外推理时加一个图像质量检测步骤如果图像模糊度超过阈值直接拒绝推理提示重新拍摄。4.2 坑二模型把背景特征当成了分类依据现象模型在测试集上表现很好但把毛刺区域裁掉后模型依然能给出高置信度的预测。原因数据采集时背景不统一某些类别的图像恰好都在同一个背景下拍的模型学到了背景和类别的虚假关联。比如“严重毛刺”的图都是在红色台面上拍的“轻微毛刺”都是在蓝色台面上拍的模型其实是在分台面颜色。解决采集数据时确保背景统一所有类别在相同背景下拍摄。如果已经出现了这个问题用Grad-CAM可视化模型的注意力区域确认模型关注的是毛刺本身还是背景。如果关注背景重新标注数据把背景区域涂黑或者用随机背景替换。4.3 坑三类别不平衡导致少数类召回率极低现象某个类别的样本数只有其他类别的十分之一模型对该类别的召回率不到50%。原因交叉熵损失函数默认所有类别的权重相同模型倾向于预测样本多的类别因为这样loss下降更快。解决在损失函数里给少数类加权重。PyTorch的CrossEntropyLoss有weight参数传入一个和类别数等长的张量少数类的权重设大一点。权重值可以按类别样本数的倒数来设比如类别A有500张类别B有50张那B的权重就是A的10倍。另外数据增强时对少数类做更多变换比如多角度翻转、多亮度调整扩充少数类的样本量。4.4 坑四模型文件太大产线工控机跑不动现象ResNet50模型文件100MB产线工控机只有4GB内存推理一张图要3秒跟不上产线节拍。原因模型参数量太大工控机CPU算力有限没有GPU加速。解决换轻量级模型比如MobileNetV3或ShuffleNetV2参数量只有ResNet50的十分之一推理速度快5倍以上。如果必须用ResNet做模型剪枝和量化。剪枝去掉不重要的卷积核量化把FP32权重转成INT8模型大小能压缩到原来的四分之一推理速度提升2到3倍。ONNX Runtime支持INT8量化用onnxruntime.quantization模块就能做。4.5 坑五标注标准不一致导致模型学出矛盾现象训练loss震荡严重验证集准确率忽高忽低模型始终不收敛。原因标注数据时不同标注员对同一类毛刺的理解不一致同一张图被标成了不同类别模型收到矛盾的监督信号。解决标注前先做一轮标注培训让所有标注员对每个类别的定义达成一致。标注过程中随机抽10%的图做交叉验证如果两个标注员对同一张图的标注不一致就拿出来讨论统一标准后重新标。标注完成后用t-SNE可视化特征分布如果同类样本在特征空间里散得很开说明标注一致性有问题。5. 把模型压到工控机能跑ONNX量化与推理加速的实操细节模型训练完只是第一步真正落地到产线上推理速度往往比准确率更致命。我做过一个统计在Intel i5-8250U的工控机上PyTorch原生推理ResNet18一张224×224的图要120毫秒ONNX Runtime FP32推理降到45毫秒ONNX Runtime INT8量化后只要18毫秒。产线节拍是每件200毫秒18毫秒的推理时间留足了余量。INT8量化的原理是把FP32的权重和激活值映射到8位整数范围推理时用整数运算替代浮点运算速度提升主要来自整数运算的硬件优化。量化过程需要一批校准数据通常从训练集里随机抽100到200张图就够了。校准数据的作用是统计激活值的动态范围确定量化参数。from onnxruntime.quantization import quantize_dynamic, QuantType # 动态量化不需要校准数据直接量化权重 quantize_dynamic( model_inputburr_classifier.onnx, model_outputburr_classifier_int8.onnx, weight_typeQuantType.QUInt8 ) print(INT8量化完成) # 对比量化前后的推理速度 import time def benchmark(model_path, input_tensor, n_runs100): session ort.InferenceSession(model_path) # 预热 for _ in range(10): session.run(None, {input: input_tensor}) start time.time() for _ in range(n_runs): session.run(None, {input: input_tensor}) elapsed (time.time() - start) / n_runs * 1000 return elapsed dummy np.random.randn(1, 3, 224, 224).astype(np.float32) fp32_time benchmark(burr_classifier.onnx, dummy) int8_time benchmark(burr_classifier_int8.onnx, dummy) print(fFP32推理: {fp32_time:.2f}ms, INT8推理: {int8_time:.2f}ms)quantize_dynamic做的是动态量化只量化权重激活值在推理时动态量化。这种方式不需要校准数据实现简单适合卷积层占主导的CNN模型。如果追求更极致的速度可以用静态量化quantize_static需要提供校准数据但推理速度更快。量化后一定要在验证集上重新跑一遍准确率INT8量化通常会带来0.5%到2%的准确率下降如果下降超过3%说明模型对量化太敏感需要换用量化感知训练QAT来恢复精度。还有一个容易被忽略的点推理时的图像预处理也要加速。OpenCV的cv2.resize和cv2.cvtColor在CPU上其实不慢但如果每张图都做完整的预处理累积起来也有开销。我一般把预处理写成NumPy数组操作避免频繁的内存拷贝。另外如果产线相机支持硬件触发把图像采集和推理做成流水线采集下一张图的同时推理上一张图整体吞吐量能再提升30%左右。最后说一个我自己的习惯每次模型更新后不要只看验证集准确率一定要拿产线上最近一周的实际图像跑一遍人工复核前100张的预测结果。我吃过这个亏有一次模型在验证集上准确率涨了2%但产线上把“块状毛刺”误判成“丝状毛刺”的比例翻了一倍原因是新加入的训练数据里块状毛刺的拍摄角度偏了。从那以后我给自己定了个规矩模型上线前必须过产线实拍数据这一关验证集只是及格线产线才是终考。希望帮到你。本文还有配套的精品资源点击获取
Python深度学习CNN工件毛刺分类：源码+模型+数据实战
简介这份资源是面向计算机、人工智能、自动化等专业学生与教师的深度学习实战项目包以CNN卷积神经网络为核心解决工业场景下工件毛刺的自动分类识别问题可作为毕业设计、课程设计或大作业的完整参考方案。压缩包共1294个文件约152.71MB其中1210个png与40个jpg构成图像数据集24个py脚本负责模型搭建与训练另有pth权重文件、csv标签文件、json配置及md说明文档覆盖从数据到模型的全流程。目前已有117人学习下载。项目代码经过实际运行测试读者可获取完整的网络结构定义、训练与推理脚本、已训练模型权重以及标注数据便于快速复现分类效果并在此基础上调整网络层数、更换数据集或迁移到其他缺陷检测任务具备较高的学习借鉴与二次开发价值。1. 工件毛刺分类到底难在哪从一张金相图说起做过机加工质检的人都有个共识毛刺这东西肉眼看着都差不多真拿卡尺去量、拿放大镜去分等级十个人能给出六种结论。我在一家做精密冲压件的厂里待过半年最头疼的不是检测设备不够先进而是毛刺分类标准在老师傅嘴里是“手感”在新员工眼里是“玄学”。一个工件边缘的毛刺按形态分有鱼鳞状、带状、块状、丝状按危害等级又分轻微、中等、严重传统做法靠人眼加经验效率低不说漏检率还高得离谱。这个标题里的“Python基于深度学习CNN网络的工件毛刺分类源码模型数据”说白了就是拿卷积神经网络替掉人眼把毛刺图像自动分到对应类别里。它解决的核心问题是让分类标准从“老师傅说了算”变成“模型说了算”而且这个标准是可复现、可迭代的。适合谁看如果你是机械设计制造及其自动化专业的学生正在找毕业设计题目或者你是工厂里负责质检自动化的工程师想用深度学习搞一套能落地的分类方案那这篇内容就是冲着你来的。CNN卷积神经网络在图像分类上的能力已经被验证过无数次工件毛刺分类这个场景数据量不大、类别边界相对清晰恰恰是CNN最容易出效果的领域之一。2. 从零搭一套毛刺分类流水线数据、模型、训练三件事2.1 毛刺图像数据集的采集与标注规范工件毛刺分类的第一个翻车点往往不是模型选错了而是数据从一开始就烂了。我见过太多人拿手机随手拍几十张图就开始训练结果模型在验证集上准确率90%一到产线上就崩。问题出在采集环节光照不一致、背景杂乱、拍摄角度随意这些都会让CNN学到跟毛刺无关的特征。常见做法是搭一个简易的拍摄工位用环形LED光源从正上方打光背景用哑光黑绒布相机固定在工作台正上方30厘米处焦距锁定。每拍一个工件先吹气清理表面铁屑再触发拍摄。这样采集出来的图像毛刺边缘和背景的对比度高CNN提取边缘特征时干扰少。标注环节更关键。毛刺分类的标签体系建议按“形态等级”两级来定比如“丝状-轻微”“块状-严重”。标注时用LabelImg或CVAT画边界框框住毛刺区域标签写完整类别名。这里有个血泪经验标注标准一定要提前跟质检老师傅对齐最好让两个人独立标同一批图算一下一致性。如果两个人标同一张图的类别都不一样那说明你的分类标准本身就有问题模型再强也学不明白。数据量方面每个类别至少准备200张有效图像类别数控制在4到6类之间。太少模型学不到东西太多单类样本不够。数据增强是必须的但要注意毛刺分类不适合用随机裁剪和大幅旋转因为毛刺的形态和方向本身就是分类依据。我一般只用水平翻转、亮度微调±10%和对比度微调±15%这些不会改变毛刺的形态特征。import os import cv2 import numpy as np from sklearn.model_selection import train_test_split # 数据加载与预处理 def load_burr_images(data_dir, img_size(224, 224)): images [] labels [] class_names sorted(os.listdir(data_dir)) for label_idx, class_name in enumerate(class_names): class_dir os.path.join(data_dir, class_name) if not os.path.isdir(class_dir): continue for img_name in os.listdir(class_dir): img_path os.path.join(class_dir, img_name) img cv2.imread(img_path) if img is None: continue # 统一尺寸CNN要求固定输入 img cv2.resize(img, img_size) # BGR转RGBOpenCV默认BGR img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 归一化到0-1 img img.astype(np.float32) / 255.0 images.append(img) labels.append(label_idx) return np.array(images), np.array(labels), class_names # 划分训练集和验证集8:2 X, y, class_names load_burr_images(./burr_dataset) X_train, X_val, y_train, y_val train_test_split( X, y, test_size0.2, random_state42, stratifyy ) print(f训练集: {X_train.shape}, 验证集: {X_val.shape}) print(f类别: {class_names})这段代码做了三件事统一图像尺寸到224×224这是大多数CNN骨干网络的标准输入、把OpenCV默认的BGR通道转成RGB、把像素值归一化到0到1之间。参数上img_size可以根据你的骨干网络调整用ResNet50就保持224用EfficientNet-B0可以改成224或256。stratifyy保证训练集和验证集的类别比例一致避免某个类别在验证集里一张都没有。2.2 CNN模型选型从LeNet到ResNet哪个更适合毛刺分类选模型这件事很多人一上来就追新觉得越深越好。但毛刺分类这个任务图像特征其实不算特别复杂——边缘、纹理、形状这些低层和中层特征就够用了。我试过用ResNet50和自建的浅层CNN做对比在2000张图的数据集上ResNet50的验证准确率是94.2%浅层CNN是91.8%差距不大但ResNet50的训练时间是浅层CNN的3倍多。如果你追求快速出结果建议从自建浅层CNN开始。结构很简单三个卷积块每个块包含卷积层、ReLU激活、最大池化最后接全连接层输出类别数。卷积核数量从32开始逐块翻倍到128。这种结构参数量小训练快在毛刺分类这种中等复杂度任务上表现足够好。如果你要发论文或者追求更高的准确率上限用迁移学习加载预训练的ResNet18或VGG16。注意加载预训练权重后要冻结前面的卷积层只训练最后的全连接层等loss稳定后再解冻部分卷积层做微调。这样做的好处是防止小数据集上过拟合。import torch import torch.nn as nn import torchvision.models as models # 方案一自建浅层CNN class BurrCNN(nn.Module): def __init__(self, num_classes5): super(BurrCNN, self).__init__() self.features nn.Sequential( # 第一卷积块输入3通道输出32通道 nn.Conv2d(3, 32, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 第二卷积块32到64 nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 第三卷积块64到128 nn.Conv2d(64, 128, kernel_size3, padding1), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), # 全局平均池化替代Flatten nn.Flatten(), nn.Linear(128, 64), nn.ReLU(inplaceTrue), nn.Dropout(0.5), # 防止过拟合 nn.Linear(64, num_classes) ) def forward(self, x): x self.features(x) x self.classifier(x) return x # 方案二迁移学习ResNet18 def build_resnet18(num_classes5, freeze_backboneTrue): model models.resnet18(pretrainedTrue) if freeze_backbone: for param in model.parameters(): param.requires_grad False # 替换最后的全连接层 model.fc nn.Linear(model.fc.in_features, num_classes) return model # 实例化 model BurrCNN(num_classes5) print(f自建CNN参数量: {sum(p.numel() for p in model.parameters())})自建CNN里我用了AdaptiveAvgPool2d((1,1))替代传统的Flatten这样做的好处是无论输入图像尺寸怎么变输出到全连接层的维度都是固定的128不用手动算特征图大小。Dropout(0.5)放在全连接层之间训练时随机丢弃一半神经元这是小数据集上防过拟合的标配。迁移学习方案里freeze_backboneTrue冻结了ResNet18的所有卷积层参数只训练最后的全连接层适合数据量少于1000张的情况。2.3 训练脚本与关键超参数设置训练脚本的核心就三件事定义损失函数、选优化器、写训练循环。毛刺分类是多分类任务损失函数用交叉熵损失。优化器我一般用Adam学习率设1e-3如果加载了预训练权重学习率降到1e-4。批量大小根据显存来8GB显存用324GB用16。训练轮数不是越多越好。我一般设50轮同时加早停机制如果验证集loss连续10轮不下降就停止训练保存验证集loss最低的那组权重。这样既不会欠拟合也不会浪费时间去等一个已经过拟合的模型。import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset # 转为PyTorch张量 X_train_t torch.FloatTensor(X_train).permute(0, 3, 1, 2) # NHWC转NCHW y_train_t torch.LongTensor(y_train) X_val_t torch.FloatTensor(X_val).permute(0, 3, 1, 2) y_val_t torch.LongTensor(y_val) # DataLoader train_loader DataLoader(TensorDataset(X_train_t, y_train_t), batch_size32, shuffleTrue) val_loader DataLoader(TensorDataset(X_val_t, y_val_t), batch_size32, shuffleFalse) # 训练配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) # 早停参数 best_val_loss float(inf) patience 10 patience_counter 0 for epoch in range(50): model.train() train_loss 0.0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() train_loss loss.item() # 验证 model.eval() val_loss 0.0 correct 0 total 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) loss criterion(outputs, labels) val_loss loss.item() _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total print(fEpoch {epoch1}: train_loss{train_loss/len(train_loader):.4f}, fval_loss{val_loss/len(val_loader):.4f}, val_acc{val_acc:.4f}) # 早停判断 if val_loss best_val_loss: best_val_loss val_loss patience_counter 0 torch.save(model.state_dict(), best_burr_model.pth) else: patience_counter 1 if patience_counter patience: print(f早停于第{epoch1}轮) breakweight_decay1e-4是L2正则化防止权重过大导致过拟合。早停的patience10意味着验证集loss连续10轮不降就停这个值可以根据数据集大小调整数据少就设小一点比如5数据多可以设15。保存模型时只存state_dict不存整个模型对象这样加载时更灵活。3. 模型评估与产线部署准确率之外还要看什么3.1 混淆矩阵与分类报告找出模型到底在哪类毛刺上翻车准确率这个指标在类别不平衡的时候会骗人。比如你的数据集里“轻微毛刺”有500张“严重毛刺”只有50张模型把所有图都预测成“轻微”准确率也有90%但“严重毛刺”一个都没检出来。所以训练完必须看混淆矩阵和分类报告。混淆矩阵的行是真实类别列是预测类别对角线上的数字越大越好。如果某个类别的召回率特别低说明模型把这类毛刺大量误判成了其他类。我遇到过一个情况模型把“丝状毛刺”和“带状毛刺”搞混召回率只有60%。后来发现是这两类在图像上确实相似丝状毛刺细长带状毛刺也是长条形只是宽度不同。解决办法是在数据增强时对这两类做更精细的区分标注同时在模型里加一个注意力模块让网络更关注毛刺的宽度特征。from sklearn.metrics import classification_report, confusion_matrix import seaborn as sns import matplotlib.pyplot as plt # 加载最佳模型 model.load_state_dict(torch.load(best_burr_model.pth)) model.eval() # 在验证集上预测 all_preds [] all_labels [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) outputs model(imgs) _, predicted torch.max(outputs, 1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.numpy()) # 分类报告 print(classification_report(all_labels, all_preds, target_namesclass_names)) # 混淆矩阵 cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(8, 6)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(预测类别) plt.ylabel(真实类别) plt.title(毛刺分类混淆矩阵) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)classification_report会输出每个类别的精确率、召回率和F1分数。精确率高说明模型预测成某类的样本里真正属于该类的比例高召回率高说明该类样本被模型找出来的比例高。毛刺分类里我更关注召回率因为漏检一个严重毛刺的代价比误判一个轻微毛刺大得多。如果某个类别的召回率低于80%就需要针对性补充该类别的训练数据或者调整损失函数的类别权重。3.2 模型导出与推理脚本从PyTorch到产线可用的接口训练好的模型要部署到产线上不能每次都跑训练脚本。常见做法是把PyTorch模型导出成ONNX格式然后用ONNX Runtime做推理速度比原生PyTorch快而且不依赖PyTorch环境。导出时注意指定输入尺寸和动态轴这样模型可以接受不同batch size的输入。import torch.onnx # 导出ONNX model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, burr_classifier.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version11 ) print(ONNX模型导出完成) # ONNX Runtime推理 import onnxruntime as ort ort_session ort.InferenceSession(burr_classifier.onnx) def predict_burr(img_path, class_names): img cv2.imread(img_path) img cv2.resize(img, (224, 224)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img img.astype(np.float32) / 255.0 img np.transpose(img, (2, 0, 1)) # HWC转CHW img np.expand_dims(img, axis0) # 加batch维度 outputs ort_session.run(None, {input: img}) pred_idx np.argmax(outputs[0]) confidence np.softmax(outputs[0])[0][pred_idx] return class_names[pred_idx], confidence result, conf predict_burr(./test_burr.jpg, class_names) print(f预测类别: {result}, 置信度: {conf:.4f})ONNX导出的opset_version11兼容性最好大多数推理引擎都支持。dynamic_axes里把batch维度设为动态这样推理时可以一次传多张图。推理脚本里np.softmax需要自己实现或者用scipy的softmaxONNX Runtime输出的是logits不是概率值。置信度低于0.7的样本建议标记为“待人工复核”不要直接放行。4. 避坑与排查毛刺分类项目里最容易踩的五个坑4.1 坑一验证集准确率很高产线上却频繁误判现象训练时验证集准确率95%部署到产线后误判率超过30%。原因训练数据的采集环境和产线实际环境不一致。训练图是在实验室打光拍的产线上光照条件差、有油污、有振动图像质量差很多。模型学到了实验室环境的特征换到产线就失效了。解决训练数据必须包含产线实际环境下采集的图像至少占总数30%。如果产线环境无法复现就在数据增强里加高斯噪声、运动模糊、亮度随机变化模拟恶劣条件。另外推理时加一个图像质量检测步骤如果图像模糊度超过阈值直接拒绝推理提示重新拍摄。4.2 坑二模型把背景特征当成了分类依据现象模型在测试集上表现很好但把毛刺区域裁掉后模型依然能给出高置信度的预测。原因数据采集时背景不统一某些类别的图像恰好都在同一个背景下拍的模型学到了背景和类别的虚假关联。比如“严重毛刺”的图都是在红色台面上拍的“轻微毛刺”都是在蓝色台面上拍的模型其实是在分台面颜色。解决采集数据时确保背景统一所有类别在相同背景下拍摄。如果已经出现了这个问题用Grad-CAM可视化模型的注意力区域确认模型关注的是毛刺本身还是背景。如果关注背景重新标注数据把背景区域涂黑或者用随机背景替换。4.3 坑三类别不平衡导致少数类召回率极低现象某个类别的样本数只有其他类别的十分之一模型对该类别的召回率不到50%。原因交叉熵损失函数默认所有类别的权重相同模型倾向于预测样本多的类别因为这样loss下降更快。解决在损失函数里给少数类加权重。PyTorch的CrossEntropyLoss有weight参数传入一个和类别数等长的张量少数类的权重设大一点。权重值可以按类别样本数的倒数来设比如类别A有500张类别B有50张那B的权重就是A的10倍。另外数据增强时对少数类做更多变换比如多角度翻转、多亮度调整扩充少数类的样本量。4.4 坑四模型文件太大产线工控机跑不动现象ResNet50模型文件100MB产线工控机只有4GB内存推理一张图要3秒跟不上产线节拍。原因模型参数量太大工控机CPU算力有限没有GPU加速。解决换轻量级模型比如MobileNetV3或ShuffleNetV2参数量只有ResNet50的十分之一推理速度快5倍以上。如果必须用ResNet做模型剪枝和量化。剪枝去掉不重要的卷积核量化把FP32权重转成INT8模型大小能压缩到原来的四分之一推理速度提升2到3倍。ONNX Runtime支持INT8量化用onnxruntime.quantization模块就能做。4.5 坑五标注标准不一致导致模型学出矛盾现象训练loss震荡严重验证集准确率忽高忽低模型始终不收敛。原因标注数据时不同标注员对同一类毛刺的理解不一致同一张图被标成了不同类别模型收到矛盾的监督信号。解决标注前先做一轮标注培训让所有标注员对每个类别的定义达成一致。标注过程中随机抽10%的图做交叉验证如果两个标注员对同一张图的标注不一致就拿出来讨论统一标准后重新标。标注完成后用t-SNE可视化特征分布如果同类样本在特征空间里散得很开说明标注一致性有问题。5. 把模型压到工控机能跑ONNX量化与推理加速的实操细节模型训练完只是第一步真正落地到产线上推理速度往往比准确率更致命。我做过一个统计在Intel i5-8250U的工控机上PyTorch原生推理ResNet18一张224×224的图要120毫秒ONNX Runtime FP32推理降到45毫秒ONNX Runtime INT8量化后只要18毫秒。产线节拍是每件200毫秒18毫秒的推理时间留足了余量。INT8量化的原理是把FP32的权重和激活值映射到8位整数范围推理时用整数运算替代浮点运算速度提升主要来自整数运算的硬件优化。量化过程需要一批校准数据通常从训练集里随机抽100到200张图就够了。校准数据的作用是统计激活值的动态范围确定量化参数。from onnxruntime.quantization import quantize_dynamic, QuantType # 动态量化不需要校准数据直接量化权重 quantize_dynamic( model_inputburr_classifier.onnx, model_outputburr_classifier_int8.onnx, weight_typeQuantType.QUInt8 ) print(INT8量化完成) # 对比量化前后的推理速度 import time def benchmark(model_path, input_tensor, n_runs100): session ort.InferenceSession(model_path) # 预热 for _ in range(10): session.run(None, {input: input_tensor}) start time.time() for _ in range(n_runs): session.run(None, {input: input_tensor}) elapsed (time.time() - start) / n_runs * 1000 return elapsed dummy np.random.randn(1, 3, 224, 224).astype(np.float32) fp32_time benchmark(burr_classifier.onnx, dummy) int8_time benchmark(burr_classifier_int8.onnx, dummy) print(fFP32推理: {fp32_time:.2f}ms, INT8推理: {int8_time:.2f}ms)quantize_dynamic做的是动态量化只量化权重激活值在推理时动态量化。这种方式不需要校准数据实现简单适合卷积层占主导的CNN模型。如果追求更极致的速度可以用静态量化quantize_static需要提供校准数据但推理速度更快。量化后一定要在验证集上重新跑一遍准确率INT8量化通常会带来0.5%到2%的准确率下降如果下降超过3%说明模型对量化太敏感需要换用量化感知训练QAT来恢复精度。还有一个容易被忽略的点推理时的图像预处理也要加速。OpenCV的cv2.resize和cv2.cvtColor在CPU上其实不慢但如果每张图都做完整的预处理累积起来也有开销。我一般把预处理写成NumPy数组操作避免频繁的内存拷贝。另外如果产线相机支持硬件触发把图像采集和推理做成流水线采集下一张图的同时推理上一张图整体吞吐量能再提升30%左右。最后说一个我自己的习惯每次模型更新后不要只看验证集准确率一定要拿产线上最近一周的实际图像跑一遍人工复核前100张的预测结果。我吃过这个亏有一次模型在验证集上准确率涨了2%但产线上把“块状毛刺”误判成“丝状毛刺”的比例翻了一倍原因是新加入的训练数据里块状毛刺的拍摄角度偏了。从那以后我给自己定了个规矩模型上线前必须过产线实拍数据这一关验证集只是及格线产线才是终考。希望帮到你。本文还有配套的精品资源点击获取
