基于ShuffleNet的电力巡检图像缺陷分类实战：从模型结构到部署验证

简介本资源面向深度学习入门者与电力巡检图像分析方向的开发者提供一套基于PyTorch的ShuffleNet卷积神经网络缺陷分类识别代码用于解决杆塔倾斜、导线脱落、绝缘子破损等典型电力线路缺陷的自动分类问题。压缩包共8个文件约215KB包含3个Python脚本分别负责数据列表生成、CNN模型训练与PyQt可视化界面另附说明文档、环境依赖文本及示例提示图结构精简、上手门槛低。代码逐行配有中文注释便于初学者理解网络搭建、数据加载与训练流程。数据集图片需自行搜集按类别放入对应文件夹即可开始训练类别可灵活增删。目前已有104人学习适合希望快速掌握图像分类实战、搭建巡检缺陷识别原型的读者参考。1. 电力线路巡检图像缺陷分类ShuffleNet 为什么值得先跑一遍输电线路巡检从人工登塔拍照转向无人机挂载可见光相机之后真正的瓶颈早就不是拍不到而是拍回来没人看得完。一个地市局一年积累几十万张杆塔、绝缘子、金具、导地线的巡检图靠人眼逐张筛漏检和疲劳几乎是必然的。缺陷分类识别要解决的就是这件事把图像喂给模型让它输出正常、破损、锈蚀、异物、销钉缺失这类标签人只看模型标出来的可疑图。标题里的 ShuffleNet 是这条链路里比较务实的一个选择。它属于轻量级 CNN 卷积神经网络核心是分组卷积加通道混洗用极小的参数量和计算量换到接近大模型的分类精度。电力巡检的部署环境通常是边缘盒子、工控机或者无人机机载算力显存和功耗都紧张ShuffleNet 这种结构比直接上 ResNet 更贴合现场。这份资料不含数据集图片意味着它给的是模型代码、逐行注释和说明文档图像得你自己按目录规范放进去。适合两类人一类是想把 CNN 缺陷分类真正落到巡检业务里的工程师另一类是拿它当毕设或课程设计、需要一份能读懂每一行的 ShuffleNet 源码的人。下面按结构怎么立住、数据怎么组织、训练怎么跑、坑在哪、怎么验证的顺序讲透。2. ShuffleNet 的结构拆解与缺陷分类的适配逻辑2.1 分组卷积和通道混洗到底解决了什么普通卷积在每个输出通道上都要和所有输入通道做乘加计算量随通道数平方增长。ShuffleNet 的做法是先分组把输入通道切成 g 组每组只在本组内做卷积计算量直接降到约 1/g。但分组带来一个副作用——不同组之间的信息被隔断了输出通道只看到自己那组输入特征表达会退化。通道混洗就是补这个洞的。它把分组卷积输出的通道重新排列让下一层的每个组都能拿到上一层不同组的特征。实现上不是真的搬数据而是 reshape 成 (g, n) 再转置再 flatten等价于一次无参数的通道重排。这个操作在代码里通常就三行但它是 ShuffleNet 区别于普通分组网络的关键。放到电力缺陷分类上这个设计的意义在于绝缘子破损、销钉锈蚀这类缺陷的判别特征往往集中在局部纹理上不需要全局大感受野分组卷积的局部性反而合适而通道混洗保证了局部特征之间还能跨组融合不会因为省算力把判别力省没了。2.2 ShuffleNet 单元的两条分支怎么读一个基本 ShuffleNet 单元stride1是残差结构主分支走 1x1 分组卷积 → 通道混洗 → 3x3 深度卷积 → 1x1 分组卷积旁支直接恒等映射最后相加。stride2 的下采样单元则把旁支改成 3x3 平均池化主分支的深度卷积步长设为 2最后用 concat 而不是相加通道数翻倍。读源码时重点盯三个地方一是 1x1 卷积后面有没有接 BN 和 ReLU顺序错了精度会掉二是通道混洗发生在第一个 1x1 分组卷积之后、深度卷积之前位置不能挪三是下采样单元里 concat 的通道分配主分支和旁支各占一半写错会导致后续层通道数对不上。import torch import torch.nn as nn def channel_shuffle(x, groups): # x: (N, C, H, W)C 必须能被 groups 整除 N, C, H, W x.size() x x.view(N, groups, C // groups, H, W) # 拆成 groups 组 x x.transpose(1, 2).contiguous() # 组间通道重排 return x.view(N, C, H, W) # 还原形状 class ShuffleUnit(nn.Module): def __init__(self, in_ch, out_ch, stride1, groups3): super().__init__() self.stride stride self.groups groups mid_ch out_ch // 2 if stride 2 else out_ch # 主分支1x1 分组卷积 - 混洗 - 3x3 深度卷积 - 1x1 分组卷积 self.gconv1 nn.Conv2d(in_ch, mid_ch, 1, groupsgroups, biasFalse) self.bn1 nn.BatchNorm2d(mid_ch) self.dwconv nn.Conv2d(mid_ch, mid_ch, 3, stridestride, padding1, groupsmid_ch, biasFalse) self.bn2 nn.BatchNorm2d(mid_ch) self.gconv2 nn.Conv2d(mid_ch, out_ch if stride 2 else out_ch, 1, groupsgroups, biasFalse) self.bn3 nn.BatchNorm2d(out_ch if stride 2 else out_ch) self.relu nn.ReLU(inplaceTrue) if stride 2: # 下采样旁支用平均池化对齐尺寸 self.shortcut nn.AvgPool2d(3, stride2, padding1) def forward(self, x): if self.stride 1: x1, x2 x.chunk(2, dim1) # 通道对半切 out self.relu(self.bn1(self.gconv1(x2))) out channel_shuffle(out, self.groups) out self.relu(self.bn2(self.dwconv(out))) out self.bn3(self.gconv2(out)) return self.relu(torch.cat([x1, out], dim1)) else: out self.relu(self.bn1(self.gconv1(x))) out channel_shuffle(out, self.groups) out self.relu(self.bn2(self.dwconv(out))) out self.bn3(self.gconv2(out)) return self.relu(torch.cat([self.shortcut(x), out], dim1))这段代码里groups3是 ShuffleNet 论文的常用取值通道数必须能被 3 整除否则view会直接报错。mid_ch out_ch // 2是下采样单元的关键主分支只输出一半通道另一半留给旁支 concat。biasFalse是因为后面接了 BN卷积自带偏置是冗余的这是轻量网络的标准写法。channel_shuffle里的transpose(1,2).contiguous()不能省transpose 后内存不连续直接 view 会抛错。2.3 缺陷分类头怎么接才不浪费前面的特征ShuffleNet 主干输出的是 (N, C, 7, 7) 这种特征图缺陷分类是整图单标签任务需要先全局平均池化压成 (N, C)再接一个全连接层映射到类别数。这里有个容易翻车的地方电力缺陷类别往往不均衡正常样本远多于缺陷样本直接接 FC 会让模型偏向多数类。常见做法是在分类头前加一层 Dropoutp0.5损失函数用带权重的交叉熵权重按类别频率的倒数设置。如果缺陷类别只有五六种FC 输出维度就是 6如果要做正常 多类缺陷的联合判别维度就是 1 缺陷类数。全局平均池化比直接 flatten 接大 FC 更稳参数量也小得多符合 ShuffleNet 的轻量定位。class ShuffleNetClassifier(nn.Module): def __init__(self, num_classes6, groups3): super().__init__() # 主干省略假设最终输出通道为 1024 self.stem nn.Sequential( nn.Conv2d(3, 24, 3, stride2, padding1, biasFalse), nn.BatchNorm2d(24), nn.ReLU(inplaceTrue), nn.MaxPool2d(3, stride2, padding1) ) self.stage nn.Sequential( ShuffleUnit(24, 240, stride2, groupsgroups), ShuffleUnit(240, 240, stride1, groupsgroups), ShuffleUnit(240, 480, stride2, groupsgroups), ShuffleUnit(480, 1024, stride1, groupsgroups), ) self.pool nn.AdaptiveAvgPool2d(1) # 任意输入尺寸都能压成 1x1 self.drop nn.Dropout(0.5) self.fc nn.Linear(1024, num_classes) def forward(self, x): x self.stem(x) x self.stage(x) x self.pool(x).flatten(1) return self.fc(self.drop(x))AdaptiveAvgPool2d(1)的好处是输入图像尺寸不必固定成 224巡检图裁剪成 256 或 320 都能跑现场数据长宽比不统一时这点很实用。Dropout(0.5)放在池化之后、FC 之前是分类网络的常规位置。类别数num_classes要和你数据目录里的子文件夹数量严格一致多一个少一个都会在训练时报维度错。3. 不含数据集图片时目录和标注该怎么组织3.1 按 ImageFolder 规范摆数据一行代码加载资料不含图片所以第一步是自己把巡检图整理成 PyTorch 能直接读的结构。最省事的是torchvision.datasets.ImageFolder它要求每个类别一个子文件夹文件夹名就是标签。目录长这样dataset/ train/ normal/ # 正常绝缘子、金具 broken/ # 破损 rust/ # 锈蚀 foreign/ # 异物 missing_pin/ # 销钉缺失 val/ normal/ broken/ ...训练集和验证集各自复制一份同样的类别目录结构。ImageFolder 会按文件夹名的字母序生成类别索引所以训练和验证的文件夹命名必须完全一致否则索引对不上验证精度会莫名其妙地低。from torchvision import datasets, transforms from torch.utils.data import DataLoader train_tf transforms.Compose([ transforms.Resize((256, 256)), # 统一尺寸 transforms.RandomHorizontalFlip(), # 巡检图左右翻转合理 transforms.RandomRotation(15), # 小角度旋转增强 transforms.ColorJitter(0.2, 0.2, 0.2), # 应对不同光照 transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # ImageNet 统计量 ]) val_tf transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.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) print(train_ds.classes) # 打印类别顺序确认和预期一致Normalize用的均值方差是 ImageNet 的统计量因为主干通常加载 ImageNet 预训练权重输入分布要对齐。RandomHorizontalFlip对巡检图是安全的杆塔左右翻转不改变缺陷语义但RandomVerticalFlip要慎用绝缘子上下翻转可能造出现实中不存在的姿态。num_workers在 Windows 上如果报错就设成 0这是 DataLoader 多进程的经典坑。3.2 类别不均衡时用带权损失别硬训巡检数据里正常样本可能占八成以上直接训会让模型学会全猜正常也能有 80% 准确率但缺陷一个都抓不到。解决办法是给交叉熵加类别权重权重取该类样本数的倒数再归一化。import numpy as np import torch # 统计每个类别的样本数 counts [len(os.listdir(fdataset/train/{c})) for c in train_ds.classes] counts np.array(counts, dtypenp.float32) weights 1.0 / counts weights weights / weights.sum() * len(counts) # 归一化均值约为 1 weights torch.tensor(weights, dtypetorch.float32) criterion nn.CrossEntropyLoss(weightweights)weights的均值归一化到 1 是为了不整体放大或缩小损失尺度影响学习率效果。如果某个类别样本极少比如只有十几张光靠权重可能还是不够常见做法是对这类样本做额外增强或者用重采样让每个 batch 里各类别比例接近。权重不是越大越好极端权重会让模型对少数类过拟合验证集上反而波动大。3.3 训练脚本的主循环和关键参数把模型、数据、损失、优化器串起来的主循环参数设置直接决定能不能收敛。import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model ShuffleNetClassifier(num_classeslen(train_ds.classes)).to(device) criterion nn.CrossEntropyLoss(weightweights.to(device)) optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max50) for epoch in range(50): 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 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs).argmax(1) correct (pred labels).sum().item() total labels.size(0) print(fepoch {epoch}, val_acc {correct/total:.4f})lr0.01配 SGD 是 ShuffleNet 从零训练的常用起点如果加载了预训练权重主干学习率可以调小到 0.001分类头保持 0.01。weight_decay1e-4抑制过拟合。CosineAnnealingLR让学习率按余弦曲线衰减比阶梯衰减更平滑T_max设成总 epoch 数。验证时一定要model.eval()加torch.no_grad()否则 BN 会更新统计量、显存也会爆。4. 训练和部署里最容易翻车的几个点4.1 现象loss 一直不降准确率卡在多数类比例原因通常是学习率太大导致梯度爆炸或者数据归一化没做、输入像素值在 0-255 范围。ShuffleNet 的 BN 层对输入尺度敏感没归一化时第一层卷积输出直接饱和。解决确认transforms.Normalize在ToTensor之后ToTensor 会把像素压到 0-1学习率从 0.01 往下试到 0.001观察前几个 epoch 的 loss 曲线。如果 loss 是 nan检查权重里有没有除零。4.2 现象验证准确率远高于训练准确率这看着像好事实际多半是验证集和训练集有重叠图片或者验证集太小、类别分布和训练集不一致。巡检图如果按时间连续采集相邻帧高度相似随机划分会让同一杆塔的图同时进训练和验证。解决按杆塔编号或采集批次划分训练验证保证同一基杆塔的图只出现在一边。验证集每个类别至少留几十张否则准确率抖动没有参考意义。4.3 现象通道混洗那行报 view size 不匹配原因是输入通道数不能被groups整除。比如某层输出 100 通道groups3100/3 不是整数view(N, 3, 100//3, H, W)就会对不上。解决设计每层通道数时保证是 groups 的倍数ShuffleNet 原版通道数都是精心选过的。自己改结构时把通道数往 3 的倍数上靠或者把 groups 改成能整除的 2 或 4。4.4 现象GPU 显存够但训练中途 OOM常见于验证阶段忘了torch.no_grad()或者 DataLoader 的num_workers太大、每个 worker 都缓存了一份数据。另外AdaptiveAvgPool2d之前如果特征图还很大中间激活占的显存会超预期。解决验证循环加with torch.no_grad()num_workers从 4 降到 2 试batch_size从 32 降到 16。用torch.cuda.empty_cache()清理缓存只能缓解根治还是减小 batch 或输入尺寸。4.5 现象模型在测试集上把锈蚀全判成正常这是典型的类别不均衡后遗症模型对少数类欠拟合。除了带权损失还要看数据里锈蚀样本是不是太少、增强是不是不够。解决对锈蚀类做针对性增强亮度扰动、局部遮挡或者用重采样让每个 batch 里锈蚀占比不低于 20%。验证时不要只看总体准确率要打印每个类别的召回率少数类召回率才是缺陷分类的真正指标。5. 用混淆矩阵和单图推理验证模型到底能不能用训练跑完总体准确率是个很粗的指标。缺陷分类真正要看的是每个类别的召回率和误判方向混淆矩阵能一次看清。下面这段代码在验证集上算混淆矩阵并支持拿单张巡检图做推理。from sklearn.metrics import confusion_matrix, classification_report import matplotlib.pyplot as plt import seaborn as sns model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) preds model(imgs).argmax(1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) print(classification_report(all_labels, all_preds, target_namestrain_ds.classes, digits4)) sns.heatmap(cm, annotTrue, fmtd, xticklabelstrain_ds.classes, yticklabelstrain_ds.classes, cmapBlues) plt.xlabel(predicted); plt.ylabel(true) plt.savefig(confusion_matrix.png, dpi150, bbox_inchestight)classification_report会给出每个类别的 precision、recall、f1重点看缺陷类的 recall。如果锈蚀的 recall 只有 0.5说明一半锈蚀被漏了得回去补数据或调权重。混淆矩阵的对角线是判对的非对角线能看出模型把哪两类搞混——比如破损和异物混淆说明这两类的视觉特征在模型眼里太像可能需要更细的标注或更强的数据增强。单图推理是上线前必做的验证拿几张没参与训练的现场图走一遍完整预处理加前向from PIL import Image def predict_one(img_path, model, class_names, device): tf transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img Image.open(img_path).convert(RGB) x tf(img).unsqueeze(0).to(device) # 加 batch 维度 model.eval() with torch.no_grad(): prob torch.softmax(model(x), dim1)[0] idx prob.argmax().item() return class_names[idx], prob[idx].item() label, conf predict_one(test/rust_01.jpg, model, train_ds.classes, device) print(f预测: {label}, 置信度: {conf:.4f})unsqueeze(0)补上 batch 维度因为模型 forward 期望 4 维输入。softmax把 logits 转成概率置信度低于 0.6 的样本建议人工复核这在巡检业务里比强行给标签更稳妥。预处理必须和验证集完全一致Resize 尺寸、Normalize 参数差一点推理结果就可能偏。我自己的习惯是每次改完结构或数据增强先跑一遍混淆矩阵盯着缺陷类的 recall 看而不是盯着总体 acc。总体 acc 高但缺陷漏检多这个模型在现场就是不可用的。另外把predict_one包成一个批量脚本对测试目录里所有图跑一遍把低置信度的图单独挑出来这批图往往就是模型边界所在也是下一轮补数据的重点。希望帮到你。本文还有配套的精品资源点击获取