简介：这份资源是面向计算机、人工智能、自动化等专业学生与教师的深度学习实战项目包，以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_ok=True) 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_size=0.3, random_state=42, stratify=None) X_val, X_test, n_val, n_test = train_test_split( X_tmp, n_tmp, test_size=0.33, random_state=42) 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_THRESHOLD=100是经验值，如果你的图像整体偏暗，可以降到 60 左右；IMG_SIZE=224是为了后续直接套用预训练模型，如果自己从零搭网络，128 也够用。
2.3 在线增强：让 300 张图发挥 3000 张的效果
工件图像不能随便做水平翻转，因为毛刺方向有工艺含义。我一般只用以下几种增强：随机旋转 ±15 度、亮度抖动 ±20%、高斯噪声、随机裁剪 90% 区域再缩放回去。用torchvision.transforms写出来是这样：
from torchvision import transforms train_tf = transforms.Compose([ transforms.RandomRotation(15), # 小幅旋转，模拟工件摆放偏差 transforms.ColorJitter(brightness=0.2, contrast=0.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(p=0.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_classes=3, freeze_backbone=True): model = models.resnet18(weights=models.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(inplace=True), nn.Linear(128, num_classes) ) return model if __name__ == "__main__": net = build_model(num_classes=3) 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_backbone=True时只训练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", transform=train_tf) val_ds = ImageFolder("dataset/val", transform=val_tf) train_loader = DataLoader(train_ds, batch_size=32, shuffle=True, num_workers=4) val_loader = DataLoader(val_ds, batch_size=32, shuffle=False, num_workers=4) # 按类别频率计算权重，缓解不平衡 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], dtype=torch.float32).to(device) weights = weights / weights.sum() * len(counts) model = build_model(num_classes=len(train_ds.classes)).to(device) criterion = torch.nn.CrossEntropyLoss(weight=weights) optimizer = AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=3e-4, weight_decay=1e-4) scheduler = CosineAnnealingLR(optimizer, T_max=30) 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(dim=1) correct += (pred == labels).sum().item() total += labels.size(0) acc = correct / total print(f"Epoch {epoch+1}, Val Acc: {acc:.4f}") if acc > best_acc: best_acc = acc torch.save(model.state_dict(), "best_model.pth") if __name__ == "__main__": train()参数说明：lr=3e-4是微调预训练模型的常用起点，如果 loss 震荡明显就降到 1e-4；weight_decay=1e-4是 AdamW 的默认正则强度；T_max=30要和总 epoch 数一致，余弦退火才能完整走完一个周期；batch_size=32在 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", transform=val_tf) test_loader = DataLoader(test_ds, batch_size=32, shuffle=False) model = build_model(num_classes=len(test_ds.classes)).to(device) model.load_state_dict(torch.load("best_model.pth", map_location=device)) model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for imgs, labels in test_loader: imgs = imgs.to(device) preds = model(imgs).argmax(dim=1).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_names=test_ds.classes, digits=4)) 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_classes=3) model.load_state_dict(torch.load("best_model.pth", map_location="cpu")) 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_version=11 ) print("导出完成") if __name__ == "__main__": export_onnx()导出后用下面这段代码做单张图推理，不依赖 PyTorch：
import 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], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.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(brightness=0.4, contrast=0.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_input="burr_classifier.onnx", model_output="burr_classifier_int8.onnx", weight_type=QuantType.QUInt8 ) print("量化完成") if __name__ == "__main__": quantize()量化后模型体积大约缩小到原来的四分之一，CPU 推理速度提升 2~3 倍。代价是准确率可能掉 1~2 个百分点，需要重新在测试集上验证一遍。如果掉得太多，可以只量化卷积层，保留全连接层为 FP32，ONNX Runtime 支持按节点指定量化类型。
另一个容易被忽视的点是输入尺寸和模型结构的匹配。如果你把输入改成 160×160，而模型里还有自适应池化层，通常不用改代码就能跑；但如果模型里有硬编码的全连接输入维度，就需要重新导出。我自己的习惯是：训练时就用最终部署的分辨率，避免训练和推理两套预处理。
最后说一个验证技巧：拿 20 张产线实拍图，人工标好标签，分别用 PyTorch 原模型、ONNX FP32、ONNX INT8 跑一遍，记录三者的预测一致率和单张耗时。如果 INT8 和 FP32 的一致率低于 95%，就放弃量化，改用降分辨率方案。这套对比流程我每次上线前都会走一遍，比拍脑袋调参靠谱得多。
希望帮到你。
本文还有配套的精品资源，点击获取
