十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

热轧带钢表面缺陷6分类实战:基于PyTorch的深度学习图像分类

热轧带钢表面缺陷6分类实战:基于PyTorch的深度学习图像分类 简介工业质检中表面缺陷检测是深度学习与计算机视觉的重要落地场景。热轧带钢表面缺陷数据集6分类作为该领域的入门级真实数据集包含氧化皮、划伤、斑块等典型缺陷为图像分类、目标检测与语义分割提供了高质量训练素材。其灰度图像特性、类别不均衡与样本规模适中的特点恰好覆盖了深度学习工程实践中的核心难点。借助迁移学习利用ImageNet预训练权重可显著提升模型收敛速度与泛化能力结合PyTorch框架设计数据增强、类别权重与早停策略能有效应对工业场景下的纹理混淆与标注噪声。从分类模型训练到检测任务改造该数据集为工业视觉项目提供了可复用的技术路线适合作为深度学习和工程落地的练手与验证平台。本文基于该数据集系统拆解6分类任务的数据处理、模型选型、训练调参与评估指标帮助读者完整掌握工业缺陷识别全流程。1. 先说清楚这个数据集到底是什么做工业视觉的同学几乎都绕不过表面缺陷检测这个方向。而热轧带钢表面缺陷数据集6分类基本可以算得上这个领域的入门必修课。它来自热连轧生产线上的真实钢板表面通过高速工业相机拍摄再由质检人员逐张标注最终整理成适合深度学习图像识别任务的公开数据集。很多人第一次接触它是在做缺陷检测分类或者目标检测练手项目的时候。说实话这个数据集的规模不算大但它的价值在于“真实”——它不是那种实验室里拍得干干净净的商品图而是带着光照不均、氧化皮、水渍、划痕干扰的现场数据训练出来的模型也更接近工业落地的状态。这个数据集能做什么最直接的任务是6分类把每一张带钢表面图像归入某一类缺陷或者判定为正常。常见缺陷类型包括轧制氧化皮、划伤、斑块、麻点、夹杂等几种具体类别标签在不同版本里略有差异但核心分类思路是一致的。更进一步你也可以把它改造成目标检测任务定位缺陷在图像中的位置或者做语义分割把缺陷像素级地标出来。也就是说同一份数据换个标签格式就是三个不同任务的训练材料。适合谁来参考如果你是刚入门深度学习图像识别的学生或工程师用这个数据集做第一个分类项目比用MNIST、CIFAR这类通用数据集更有成就感因为你能直观看到模型在复杂工业场景下的表现如果你已经在做工业质检相关项目这个数据集也可以作为预训练验证集测试你的网络结构、数据增强策略或者超参组合好不好用。我后面写的内容基本是按照“拿到数据—理解数据—训练模型—调参避坑”这条路线展开的既讲原理也讲操作你可以直接照着跑。2. 整体设计思路6分类任务的核心难点在哪2.1 类别的含义与图像特征先认识一下这6类缺陷长什么样。热轧带钢在高温状态下经过多道次轧制表面会形成各种不同类型的缺陷。比较常见的有以下几种氧化铁皮压入高温状态下氧化皮被压入钢板表面呈不规则块状或条带状颜色偏深。划伤与设备接触产生的线状痕迹方向通常沿轧制方向宽度不一。斑块表面颜色差异明显的区域边界模糊形状不规则。麻点细小密集的点状凹坑分布可能均匀也可能局部聚集。夹杂非金属夹杂物破损后暴露在表面形态多样常伴随颜色突变。正常表面没有明显缺陷的常规表面用于作为负样本。从图像角度看这些缺陷都拍摄在灰度图上意味着模型没法靠颜色信息做判断只能学习纹理、边缘、形状和灰度分布特征。这一点要特别注意很多用惯了RGB彩色图做分类的人第一次在灰度数据集上训练会明显感觉模型收敛变慢因为可用的特征维度少了一截。应对办法通常是加强纹理类数据增强比如随机裁剪、旋转、缩放、对比度调整等让网络从局部纹理中提取更多判别性信息。2.2 为什么6分类比二分类难得多如果只是“有缺陷/无缺陷”的二分类很多简单网络就能做到比较高的准确率因为正常表面和缺陷表面的差异实在太明显了。但6分类要求模型区分的是“哪一种缺陷”这就难了。举个例子划伤和麻点都是表面损伤类的缺陷有时候在低分辨率下看起来很像氧化铁皮压入和夹杂都伴随局部的灰度突变容易混淆。这就像让一个外行人去区分几种不同类型的皮肤病变普通人能看出“有问题”但说不出具体是哪一种而医生经过专业训练才能准确分辨。所以在这个任务上模型真正要学的不是“异常的普遍特征”而是“每一类缺陷的独有特征”。这要求训练数据质量要高、标注要准确同时也要求网络结构有足够强的特征表达能力。当你发现自己训练出的模型准确率卡在80%左右上不去时大概率不是网络不够深而是类别混淆问题没有针对性地解决。2.3 数据规模与均衡性问题这个数据集的总体样本量不算大平均下来每类大约一两千张总共一万多张。对于深度学习图像识别来说这个规模属于“小样本中的中等水平”——用ResNet从头训练很容易过拟合但用预训练模型做迁移学习又因为灰度图、工业场景与ImageNet分布差异大而效果打折。更麻烦的是6个类别并不是均匀分布的。有的类别样本特别多有的特别少。直接拿原始数据训练模型会对样本量大的类别产生偏好小样本类别会被“淹没”。所以拿到数据后第一件事不是写网络而是做类别分布统计。我一般会先打印每类的样本数画个柱状图看一下再决定后续是加权采样、过采样还是做更多数据增强。3. 数据集与训练方案的选择3.1 数据集的目录结构与标签格式不同来源版本的数据集目录结构略有不同但基本都遵循分类数据集的常见布局。下载解压后通常是一堆子文件夹每个文件夹名就是类别名里面的图片就是该类别的样本。比如data/ ├── train/ │ ├── crazing/ │ ├── inclusion/ │ ├── patches/ │ ├── pitted_surface/ │ ├── rolled_in_scale/ │ └── scratches/ └── test/ ├── crazing/ ├── inclusion/ ├── patches/ ├── pitted_surface/ ├── rolled_in_scale/ └── scratches/如果你拿到的是类似这种结构的原始数据那么恭喜你省去了不少标注转换的功夫。直接用torchvision.datasets.ImageFolder就能加载。如果是那种所有图片放在一个文件夹、另附一张CSV标签表的版本就得多做一步读CSV、按文件名映射标签、自定义Dataset类。我自己比较推荐先整理成ImageFolder结构因为它兼容性最强后续想换模型、换框架都方便。如果将来要做目标检测再转成YOLO格式的txt标签也容易。转换时注意保持文件名唯一性别在不同目录下用重名文件否则后面排查问题会非常痛苦。3.2 模型选型从ResNet到EfficientNet这个数据集用什么网络结构合适我的经验是优先选择在ImageNet上预训练过、但规模适中的模型。ResNet18、ResNet34是入门首选因为参数量小、训练快、不容易过拟合ResNet50效果会好一些但在这个数据量下不加正则化措施也更容易过拟合。EfficientNet系列在算力足够的情况下可以试它的复合缩放策略在工业缺陷这类纹理敏感任务上往往表现不错。如果不想从零训练最常见的做法是把预训练模型最后一层全连接层替换成新的分类层类别数改成6。注意灰度图输入问题预训练模型默认接受3通道RGB输入而带钢数据集是单通道灰度图。两种处理方式一种是把灰度图复制成3通道简单省事也能用预训练权重另一种是修改网络第一层卷积的输入通道数为1但这样做的话预训练权重就对应不上了相当于第一层要重新学。实践经验告诉我复制成3通道几乎总是更好因为预训练权重带来的收益远大于第一层通道数不同带来的损耗。3.3 损失函数与不平衡处理6分类任务的基础损失函数用交叉熵CrossEntropyLoss没问题但类别不均衡时要做调整。最直接的方法是给每个类别设置不同的权重让模型在计算损失时对小样本类别给予更多关注。PyTorch里实现起来很简单import torch.nn as nn class_weights torch.tensor([1.0, 2.0, 1.5, 3.0, 2.5, 1.0]) # 根据各类样本占比反比设定 criterion nn.CrossEntropyLoss(weightclass_weights)具体权重怎么定最简单的做法是用样本数的倒数做归一化。假设六个类别的样本数分别是n1到n6那每个类的权重可以设为total_samples / (num_classes * ni)这样样本少的类权重高样本多的类权重低。这个公式不神秘但实测有效比手动凭感觉调权重靠谱得多。另一个思路是用Focal Loss。它在交叉熵基础上引入调制因子让模型把注意力集中在难分类的样本上。对于缺陷检测这种存在“难例”的场景Focal Loss偶尔能带来惊喜尤其是那些形态上比较接近的类别之间。但注意Focal Loss对超参数gamma很敏感默认gamma2不一定是最优建议在验证集上小范围搜索一下。4. 实操过程完整训练流程拆解4.1 数据加载与预处理先把基本依赖装上PyTorch、torchvision、OpenCV、NumPy、Matplotlib还有tqdm用来显示进度条。环境建议直接用Anaconda创建独立虚拟环境防止依赖冲突。安装命令很简单conda create -n steel python3.10 conda activate steel pip install torch torchvision opencv-python numpy matplotlib tqdm scikit-learn数据加载阶段有几个关键点要注意。第一图片尺寸统一。原数据集的图片分辨率可能不统一或者分辨率很高比如200x200需要缩放到固定尺寸。对于这个数据集我推荐统一缩放到224x224因为这是ResNet等主流模型的标准输入尺寸也不至于因为过度压缩丢失太多纹理细节。第二数据标准化。ImageNet预训练模型要求输入按照特定均值和标准差做归一化灰度图复制成3通道后也要遵守这个规则否则预训练权重效果会打折扣。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) test_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])注意训练集和测试集的变换是不一样的。训练集需要做随机增强增加数据多样性测试集只做缩放和归一化保证评估结果稳定可复现。4.2 数据集划分与加载器构建很多版本的数据集自带了训练集和测试集划分但即便如此我也建议在训练集内再划分一小部分作为验证集。原因很简单测试集只能用来做最终评估如果用它做模型选择或调参就无形中“污染”了测试集最终评估分数会虚高。宁可把训练集再分出10%或15%作为验证集用验证集来监控训练过程、保存最佳模型。加载数据时用torch.utils.data.DataLoader设置合适的batch_size和num_workers。batch_size建议16或32取决于显存大小。num_workers一般设成CPU核数的一半左右太多反而会因为进程切换开销降低效率。训练时设置shuffleTrue验证时不需要shuffle。from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(data/train, transformtrain_transform) val_dataset ImageFolder(data/val, transformtest_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4)顺便提一句ImageFolder会自动把子文件夹名转成类别索引顺序是按文件夹名的字母序排的。这个索引顺序要记住后面画混淆矩阵、算每一类的precision和recall时索引和类别名的对应关系很容易搞混。4.3 模型构建与训练循环以ResNet34为例加载预训练权重替换最后的全连接层import torchvision.models as models import torch.nn as nn model models.resnet34(weightsmodels.ResNet34_Weights.IMAGENET1K_V1) num_features model.fc.in_features model.fc nn.Linear(num_features, 6)优化器一般用Adam或SGD带动量。用SGD时初始学习率可以设0.001到0.01mometum0.9weight_decay1e-4。用Adam时学习率建议小一点1e-4左右。我个人习惯先用Adam快速跑通流程确认没bug后再换SGDCosineAnnealing精细调。学习率调度方面StepLR每5个epoch衰减0.1或者ReduceLROnPlateau根据验证集loss自动降低都是常见做法。训练循环的骨架代码如下for epoch in range(num_epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) # 验证 model.eval() val_correct 0 val_total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) val_correct (predicted labels).sum().item() val_total labels.size(0) val_acc val_correct / val_total print(fEpoch {epoch}: loss{running_loss/len(train_dataset):.4f}, val_acc{val_acc:.4f})这里有几个细节值得展开。第一别忘了把模型切到训练模式model.train()和验证模式model.eval()。Dropout和BatchNorm在两种模式下的行为不一样漏掉这一行会在验证时得到奇怪的结果。第二验证阶段一定要包在torch.no_grad()里否则会计算梯度白白消耗内存还拖慢速度。第三每个epoch结束时保存模型权重不要只保存最后一个epoch的结果因为训练后期可能已经过拟合了验证集上最好的模型往往出现在中途。保存时建议同时保存模型参数和优化器状态torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, val_acc: val_acc, }, fcheckpoint_epoch{epoch}.pth)4.4 训练轮数与早停策略这个数据集规模不大训练轮数不必像ImageNet那样动辄上百。一般20到30个epoch就能看到明显收敛再多就很容易过拟合。我通常会同时监控训练集和验证集的loss曲线如果训练loss持续下降而验证loss开始回升说明模型开始“背答案”而不是“学规律”这时候即使准确率还在涨也要谨慎真实泛化能力可能已经开始下降了。早停是个很实用的策略设定一个patience值比如连续5个epoch验证集准确率没有提升就停止训练并恢复到验证集表现最好的那次权重。配合早停策略训练轮数可以设大一点比如50让它自己决定什么时候停。best_val_acc 0.0 patience 5 wait 0 for epoch in range(num_epochs): # ... 训练验证代码 ... if val_acc best_val_acc: best_val_acc val_acc wait 0 torch.save(model.state_dict(), best_model.pth) else: wait 1 if wait patience: print(fEarly stopping at epoch {epoch}) break4.5 评估指标不要只看准确率6分类任务的评估只看整体准确率是远远不够的。原因还是类别不均衡如果一个类别占样本总数的40%你就算把所有样本都分到这个类也有40%的准确率但这显然是个废物模型。所以要结合混淆矩阵、每一类的precision、recall和F1-score来看。from sklearn.metrics import confusion_matrix, classification_report # 假设已收集所有预测结果和真实标签 cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_namesclass_names))混淆矩阵在缺陷分类任务里特别重要因为它能告诉你哪些类之间容易被混淆。比如如果你发现“划伤”和“麻点”经常互相误判说明这两种缺陷在图像特征上确实接近或者你的数据增强策略导致两者特征混叠了。这时候可以考虑针对性增加一些区分度高的特征预处理或者检查标注是否存在错误。5. 常见问题与排查技巧实录5.1 模型不收敛、loss震荡刚跑起来时loss剧烈震荡甚至不下降这是最常见的问题。先别急着换网络结构按这个顺序排查一是检查数据归一化和预处理是否正确灰度图转三通道时有没有保持像素值范围合理二是学习率是不是太大试试降到原来的十分之一三是检查标签和图像是否对应随机抽几张图看一眼防止文件名错位导致模型在学“随机映射”。还有一种情况是验证集准确率一直不涨但训练集已经很高这是典型的过拟合信号。解决思路有几个方向加大数据增强强度、加Dropout层、减小模型复杂度、或者减小学习率并增加训练轮数。5.2 类别不均衡导致小样本类效果差如果你发现F1-score按类别看时某些类别特别低首先确认是不是因为这类样本本身太少。解决办法我在前面提过类别权重、Focal Loss、过采样。但还有一个容易被忽略的点——数据增强策略应该是“分而治之”的。比如“划伤”这类缺陷具有方向性旋转增强可能反而不合适而“麻点”这类缺陷对旋转不敏感旋转增强就能显著扩充样本。我见过有同学用统一的增强策略把划伤增强出来的图像旋转了90度结果模型学到了“旋转后的划伤”反而看不清真实划伤的特征。这里建议按类别单独设计增强策略。5.3 灰度图还是三通道这个问题经常有人纠结。有些实现把灰度图直接当成单通道输入修改网络第一层有些复制成三通道。我建议复制成三通道。原因是使用ImageNet预训练权重时第一层卷积的输入通道是3如果改成1权重数量对不上只能随机初始化这一层这会拖慢收敛而且第一层提取的基础特征边缘、纹理本来就是通用的用预训练权重利大于弊。复制成三通道的代价只是多了一点点内存几乎没有性能损失。5.4 标注质量问题的应对公开数据集虽然经过整理但标注错误依然存在。怎么应对一种做法是“软标签”不把标注当成绝对真理在损失计算中引入标签平滑Label Smoothing让模型对错误标注的容忍度更高。另一种做法是“难样本分析”在训练完成后把预测错误且置信度高的样本挑出来人工检查是不是标注本身就有问题。我曾经在一个类似的数据集上发现过约3%的标注错误修正后模型准确率提升了近2个百分点这个提升幅度相当可观。5.5 过拟合的进一步处理如果用了预训练模型、数据增强、早停之后还是过拟合可以考虑这几个手段一是冻结预训练模型的前几层只训练后面的层这在小数据集上非常有效二是在全连接层前加Dropout概率设0.3到0.5三是对模型参数做L2正则化weight_decay设大一点。还有一种技巧是用MixUp或CutMix做数据增强把不同样本混合起来训练能给网络带来隐式正则化的效果在缺陷分类上我试过确实能稳定提升泛化能力。6. 从分类到检测这个数据集还能怎么玩6.1 转成目标检测任务如果你不满足于分类想体验一下目标检测这个数据集也可以改造成简单的检测任务。方法是用LabelImg或X-AnyLabeling标注工具在每张图上框出缺陷的边界框导出成YOLO格式或COCO格式然后丢给YOLOv8或Faster R-CNN训练。由于缺陷在图像中占比往往较小建议用YOLOv8的m或l型号小目标检测效果会好一些。数据集划分思路也和分类类似训练集、验证集、测试集按比例划分建议7:2:1。如果原始数据集图片尺寸较大注意是否需要对大图做滑窗裁剪避免把缺陷缩得太小导致检测器根本瞧不见。6.2 迁移学习与预训练技巧在工业缺陷检测场景下从零训练一个模型很少是明智的。除了使用ImageNet预训练还有一个思路是先在大型钢材表面缺陷数据集上做预训练再微调到你的小数据集上。这种“领域内预训练”的方法虽然多一道工序但效果往往比直接用ImageNet预训练更好因为上游任务和下游任务的图像分布非常接近迁移效率更高。如果没有那么大的数据集可以预训练也可以尝试“无监督预训练”用SimCLR或BYOL之类的对比学习方法先在未标注的数据上学习特征表示再拿这些特征做下游分类初始化。这种方式在小样本场景下越来越流行值得一试。6.3 部署到实际生产环境的思考训练出一个高准确率的模型只是第一步真正落地到工业现场还要面对很多挑战。首先是推理速度生产线上的质检往往需要毫秒级响应所以模型量化、TensorRT加速、模型蒸馏这些技术都可能派上用场。其次是分布偏移换了产线、换了相机、换了光照条件模型效果都可能下降所以需要定期用新数据微调。在这个数据集上训练出来的模型可以作为一个很好的“缺陷检测常识库”。当你在一个新的钢铁产线上做项目时不需要从零开始用这里的预训练权重做初始化再用新产线的少量标注数据微调往往能快速达到可用的效果。这也是为什么建议大家即使不做钢铁表面缺陷也值得拿来作为迁移学习的练手数据。7. 个人实操体会与避坑小结最后说几句实在话。这个数据集看着简单但真的能出高分的方案不多原因不是网络架构不够花哨而是很多人在基础的训练细节上栽了跟头。我见过太多同学拿ResNet直接开跑40个epoch后准确率停在85%就上不去了然后开始疯狂换模型换完发现还是85%。其实问题不在模型而在数据端和训练策略端。我个人的建议是这样一套“保底组合拳”复制成三通道 ResNet34预训练 随机裁剪/翻转/颜色抖动增强 类别权重交叉熵 AdamW优化器 余弦退火学习率 早停 验证集模型选择。这套组合在这个数据集上稳定能达到90%以上的测试准确率。如果你追求极致再试试EfficientNet-B3或ConvNeXt配合TTA测试时增强能再往上涨一点但涨幅有限要权衡投入产出比。还有一个小技巧训练过程中把每个epoch的混淆矩阵打印出来而不是只看loss和acc。这样你能实时发现哪个类别在变好、哪个类别在变坏而不是等训练完了才醒悟。我自己做这个数据集时最早发现“斑块”和“夹杂”混淆率居高不下就是靠观察中途混淆矩阵发现的后来通过增加这两类数据的裁剪缩放增强把混淆率压下去不少。总之这个数据集是一个很好的“照妖镜”能把你训练流程中所有不规范的细节都暴露出来。把它完整跑一遍你对图像分类、迁移学习、数据不平衡、模型评估这些东西的理解会比看十篇教程都深。本文还有配套的精品资源点击获取
返回列表