
简介面向少量样本故障诊断场景的深度学习资源包基于PyTorch框架实现整合一维信号注意力机制与多种训练优化技巧适合研究生、科研人员及工业AI开发者进行复现和二次开发。代码覆盖1D信号注意力机制、AMSGradP优化器、1D-Meta-ACON自适应激活、GAP全局池化、1D-Grad-CAM可视化以及AdaBN域自适应等关键模块并针对BiGRU/BiLSTM后接GAP操作的有效性给出实证说明便于理解小样本建模中的典型提升手段。资源共18个文件包括10个mat格式的数据集与8个Python脚本整包约10.89MB数据文件可直接用于训练脚本则按模型、优化器、训练工具等维度划分结构清晰、取用方便。已有358人学习下载配套了早停、标签平滑、数据保存与继续训练等实用流程也可以替换自己的mat数据快速验证故障识别效果。整体兼顾算法原理与工程落地为后续改进或论文复现提供了良好起点。1. 少量样本故障诊断的难点与注意力方向的选型如果你在产线上遇到一类故障只采集到几十条有效样本常规深度分类网络基本会瞬间过拟合验证集准确率停在 60% 上下还不停震荡。这类任务有一个统一的称呼少量样本故障诊断。它的核心矛盾不是模型不够深而是故障特征在正常样本的强噪声背景下太微弱样本稍少一点模型就会去拟合噪声峰而不是故障模式。基于注意力机制的少量样本故障诊断 pytorch 方案本质上是把注意力当成一种特征筛选器让网络在极少量样本下集中关注与故障相关的频带或通道再配合元学习式的训练策略完成分类。本文面向做过深度学习、但对少样本任务还不太熟的工程师覆盖从数据组织、模型构建到训练参数调整的完整落地路径。2. 数据组织振动信号切片、归一化与少样本采样器2.1 为什么先处理数据而不是直接堆模型少量样本任务对数据划分的要求比普通分类严格得多。很多人直接在原始振动序列上滑窗切出几万条样本再随机划分训练测试集。这在样本量充足时问题不大但少量样本场景下同一个文件切出来的相邻窗口高度相似随机划分会让验证集里出现训练样本的“近亲”评估出来的准确率高出实际水平 10 到 20 个百分点。常见做法是以故障类型为最小单位划分域。比如公开轴承数据集 CWRU每个故障位置和损伤程度对应一段连续采集信号我会把这一段信号内切出的所有窗口全部划到同一个集合里不让它跨集合出现。否则模型记忆的是某一段连续信号的背景噪声而不是故障本身。2.1.1 切片的三个参数怎么定窗口长度决定频率分辨率。振动信号采样率通常在 12kHz 到 48kHz 之间1024 点窗口对应约 12Hz 的频率分辨率对轴承故障特征频率已经够用。窗口太短注意力机制能看到的时间上下文不足窗口太长样本占比被正常信号稀释故障样本更难切。步长控制样本数量。少量样本场景下每个文件切成几百条可以接受步长设为窗口长度的四分之一即可。下表是我常用的几组配置基本覆盖常见振动信号场景窗口长度步长单文件样本数适用任务轴承高速运转故障特征频率高512128约几百条保持架、外圈故障齿轮箱或低速重载设备2048512约一两百条齿面磨损、断齿实验室公开数据集带标注1024256数百条跨工况泛化验证2.1.2 归一化方法对跨工况故障诊断的影响我用 z-score 归一化而不是 min-max。原因很简单不同工况下振动幅值差异很大min-max 会被偶尔出现的冲击峰主导把正常的微小幅值压缩到零附近。z-score 对幅值尺度不敏感保留的是信号形状信息。需要重点注意的是统计量只在训练域计算然后直接应用到测试域。如果在每个 episode 内部重新计算均值和方差会把测试域的分布信息泄给模型跨工况评估就失去意义。数据加载阶段把这个逻辑固定下来。2.2 写一个少样本采样器way-shot-query 的代码少样本训练时不能按 batch 随机打乱。常见做法是采用 episode 采样每轮随机挑出 way 个故障类别每个类别抽 shot 条样本作为支持集再抽 query 条作为查询集。模型只在这个小型分类任务上计算损失并回传梯度。import numpy as np def sample_episode(labels, way5, shot1, query10, seedNone): if seed is not None: np.random.seed(seed) unique_cls np.unique(labels) chosen np.random.choice(unique_cls, way, replaceFalse) support_idx, query_idx, cls_map [], [], {} for new_id, cls in enumerate(chosen): idx np.where(labels cls)[0] np.random.shuffle(idx) cls_map[cls] new_id support_idx.extend(idx[:shot]) query_idx.extend(idx[shot:shot query]) return (np.array(support_idx), np.array(query_idx), cls_map)这段代码的核心是把分类目标从原始故障标签重映射到 0 到 way-1。调用时传入全局样本标签数组返回支持集和查询集的样本索引。注意 query 的数量不要超过该类别可用样本数否则会选空。2.3 对比随机采样与 episode 采样的差别随机采样是最常见的误用方式。标签为 0 到 9 的 10 类故障每个 batch 随机抽 32 条样本模型在梯度更新时看到的是全局类别分布这种方式训练出来的模型在支持样本只有 1 到 5 条时几乎没有泛化能力。episode 采样让每次梯度更新都模拟“少量样本分类”的真实测试条件模型必须在 5 个类别、每类 1 到 5 条样本的支持下做出判断。这种训练与测试分布的一致性比任何复杂模型结构都重要。在少量样本故障诊断的代码实现中我建议把采样器单独封装而不是和数据集写在一起方便后续切换不同 way-shot 配置。3. 基于注意力机制的特征提取器 PyTorch 实现3.1 通道注意力在少量样本下的三个可用位置注意力机制在故障诊断里常见放置位置有三个通道维度、空间/时间维度、以及跨模态维度。少量样本场景下通道注意力优先级最高。通道注意力做的是对卷积输出的每个通道打分保留对故障敏感的频带信息抑制噪声通道。信号经过第一层卷积后不同通道对应不同频率带通道注意力直接作用于这些频带。时间维度注意力对信号长度维做加权能增强瞬态冲击段但在样本量少时容易出现注意力分数饱和表现为所有时间位置权重接近相同。多头自注意力原理上最强但计算量随序列长度二次增长1024 点输入光 QK 矩阵就接近百万量级少量样本下很容易把注意力集中在个别噪声峰上。基于这个对比我一般选择把通道注意力嵌入残差块里这也是实践中最稳妥的结构注意力类型参数增量对振动信号的适用性少量样本下风险SE 通道注意力很小强但全连接降维丢失通道相关性降维比例需要调CBAM中等强通道与空间并行空间注意力容易饱和ECA 通道注意力极小强无降维保持通道依赖基本无多头自注意力大适合长序列但计算重过拟合和注意力崩溃3.2 可复现的 ECA 模块与残差编码器ECA 注意力去掉了 SE 里的两个全连接层改用一维卷积捕捉相邻通道的依赖参数只有 kernel_size 大小对少量样本训练非常友好。PyTorch 实现只需要二十行左右import torch from torch import nn class ECABlock(nn.Module): def __init__(self, channels, k_size5): super().__init__() self.avg_pool nn.AdaptiveAvgPool1d(1) self.conv nn.Conv1d(1, 1, kernel_sizek_size, paddingk_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): # x shape: [B, C, L]B 为 batchC 为通道数L 为信号长度 y self.avg_pool(x) # [B, C, 1]每通道全局均值得一个标量 y y.squeeze(-1).unsqueeze(1) # [B, 1, C]把通道维放到序列维便于 1D 卷积 y self.conv(y) # [B, 1, C]让相邻 C 个通道互相加权 y y.squeeze(1).unsqueeze(-1) # [B, C, 1] y torch.sigmoid(y) # 权重范围 [0, 1] return x * y参数里唯一需要调整的是 k_size它表示每个通道和左右邻居的交互范围。默认 5 在 256 通道的模型上表现稳定通道数较少时改成 3 更稳。整个模块没有 bias也没有 BN减少了不少需要训练的参数量。编码器部分使用带残差的一维卷积堆叠。下采样过程中不断增大通道数ECA 放在每个残差块的加法之前class ResidualConvBlock(nn.Module): def __init__(self, in_ch, out_ch, stride1): super().__init__() self.conv1 nn.Conv1d(in_ch, out_ch, 7, stride, 3, biasFalse) self.bn1 nn.BatchNorm1d(out_ch) self.relu nn.ReLU(inplaceTrue) self.conv2 nn.Conv1d(out_ch, out_ch, 7, 1, 3, biasFalse) self.bn2 nn.BatchNorm1d(out_ch) self.eca ECABlock(out_ch) self.shortcut nn.Sequential() if stride ! 1 or in_ch ! out_ch: self.shortcut nn.Sequential( nn.Conv1d(in_ch, out_ch, 1, stride, biasFalse), nn.BatchNorm1d(out_ch) ) def forward(self, x): out self.relu(self.bn1(self.conv1(x))) out self.bn2(self.conv2(out)) out self.eca(out) # 注意力加在残差加法之前 out self.shortcut(x) return self.relu(out) class ProtoNetEncoder(nn.Module): def __init__(self, in_len1024, feat_dim128): super().__init__() self.stem nn.Sequential( nn.Conv1d(1, 64, 15, 2, 7, biasFalse), nn.BatchNorm1d(64), nn.ReLU(inplaceTrue) ) self.layer1 ResidualConvBlock(64, 128, stride2) self.layer2 ResidualConvBlock(128, 128, stride2) self.layer3 ResidualConvBlock(128, 256, stride2) self.avg_pool nn.AdaptiveAvgPool1d(1) self.fc nn.Linear(256, feat_dim) def forward(self, x): x self.stem(x) x self.layer1(x) x self.layer2(x) x self.layer3(x) x self.avg_pool(x).squeeze(-1) return self.fc(x) # 输出 128 维嵌入向量编码器输出的是故障信号的嵌入向量不做直接类别输出。嵌入向量的维度选 128 比较合适太小不足以区分细微故障模式太大在样本少时维度灾难明显。整个模型参数量大约在 80 万左右一张消费级显卡足够训练。3.3 少样本头部原型网络距离度量特征提取器之后接的是原型网络头。把支持集中每个类别的向量取平均得到该类的原型向量然后计算查询样本到所有原型的欧氏距离。这个设计在少量样本任务里几乎是标配原因是它不对类别边界做显式建模完全靠距离度量完成分类。def compute_prototypes(support_feat, way, shot): # support_feat 已按类别顺序排列前面假设每次 episode 内每类连续放置 support_feat support_feat.view(way, shot, -1) return support_feat.mean(dim1) # 每个类别所有支持向量的均值 def classify_with_prototypes(query_feat, prototypes, temperature10.0): # 欧氏距离越小越接近原型取负后乘 temperature 得到 logits dist torch.cdist(query_feat, prototypes, p2) logits -dist * temperature return logits, dist这里的 temperature 参数值得细说。欧氏距离的数值范围与嵌入向量的模长强相关直接进 softmax 会显得过度平滑。乘以一个大于 1 的温度后距离差异被放大logits 拉开分类决策更锐利。温度过大会导致梯度消失过小则模型几乎不更新。实际使用中 8 到 12 是一个可靠区间具体表现要看嵌入向量的模场量级。4. 训练参数设定的关键细节与故障诊断避坑4.1 episode 训练主循环端到端更新注意力有了编码器和原型计算函数训练循环可以写得非常短。下面这段代码是完整可跑的主循环骨架import torch import torch.nn.functional as F device torch.device(cuda if torch.cuda.is_available() else cpu) encoder ProtoNetEncoder().to(device) optimizer torch.optim.Adam(encoder.parameters(), lr1e-3) steps_per_epoch 200 way, shot, query 5, 5, 15 temperature 10.0 for epoch in range(60): encoder.train() total_loss 0.0 for step in range(steps_per_epoch): sup_idx, que_idx, cls_map sample_episode( train_labels, wayway, shotshot, queryquery, seedepoch * 1000 step ) sup_x torch.tensor(train_x[sup_idx], dtypetorch.float32).to(device) que_x torch.tensor(train_x[que_idx], dtypetorch.float32).to(device) que_y torch.tensor([cls_map[int(train_labels[i])] for i in que_idx], dtypetorch.long).to(device) sup_feat encoder(sup_x) que_feat encoder(que_x) prototypes compute_prototypes(sup_feat, way, shot) logits, _ classify_with_prototypes(que_feat, prototypes, temperature) loss F.cross_entropy(logits, que_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() if (epoch 1) % 10 0: print(fepoch {epoch 1:02d}, loss {total_loss / steps_per_epoch:.4f})逻辑上需要注意两点。第一支持集和查询集都要经过编码器但只有查询集参与损失计算支持集只负责生成原型这个区分不能混淆。第二每次 episode 都重新随机采样同一个样本会在多轮中反复出现但和它同时出现的支撑集不同梯度方向也会变化这正是 episode 训练能缓解过拟合的机制。4.2 超参表way、shot、query、temperature 的经验值下面的经验值列表适合大部分机器振动场景。way 表示每次任务里的类别数一般和真实诊断场景的候选故障数一致不要为了追求难度去加大到 10 以上shot 是每类提供的支持样本数真实应用中是固定资源训练时尽量贴近部署条件。参数推荐值说明way5 或真实故障类别数类别太多会加大类间混淆shot1 到 5训练用 5评估时逐级降为 1query10 到 20query 太少会让梯度波动大temperature8 到 12过小梯度消失过大决策尖锐嵌入维度128比类别数大一个量级即可学习率1e-3 起步30 个 epoch 后降到 1e-4如果发表论文或做基准实验还建议把 way-shot 组合固定成几个标准配置比如 5-way-1-shot 和 5-way-5-shot便于和其他故障诊断代码比较。不要在一个实验里混用不同配置去对比准确率这样得出来的结论不可比。4.3 数据增强与过拟合防御的常见做法少量样本下数据增强是不可缺少的一环。对振动信号来说有三个增强操作是安全的叠加高斯噪声、幅值缩放、起始相位随机偏移。def augment_signal(x): # x 形状 [1, 1024] noise torch.randn_like(x) * 0.02 x x noise scale torch.empty(1).uniform_(0.9, 1.1).to(x.device) x x * scale return x增强时需要避免破坏相位信息的操作比如时间反转或时域裁剪。故障信号的特征频率和相位结构对诊断有实际意义破坏之后注意力机制学到的东西会失去可解释性。过拟合防御方面建议在 epoch 20 左右观察一次验证集表现连续 5 个 epoch 不提升就提前终止。上面代码里的 lr 直接用一个常数也可以但配合 cosine 衰减调度器会更稳尤其当支持样本只有 1 条时。5. 用注意力权重定位故障频段的验证技巧5.1 跨工况验证脚本与多 seed 评估少量样本诊断模型最大的风险是评估方差大。单次训练的准确率可能在 60% 到 90% 之间剧烈波动只有重复跑 5 到 10 次并统计均值方差才有分析价值。下面这段脚本在给定测试域上重复评估输出均值加减标准差def evaluate(encoder, test_x, test_labels, way5, shot5, query20, num_episodes50): encoder.eval() acc_list [] with torch.no_grad(): for i in range(num_episodes): sup_idx, que_idx, cls_map sample_episode( test_labels, wayway, shotshot, queryquery, seed100 i ) sup_x torch.tensor(test_x[sup_idx], dtypetorch.float32).to(device) que_x torch.tensor(test_x[que_idx], dtypetorch.float32).to(device) que_y [cls_map[int(test_labels[j])] for j in que_idx] sup_feat encoder(sup_x) que_feat encoder(que_x) prototypes compute_prototypes(sup_feat, way, shot) logits, _ classify_with_prototypes(que_feat, prototypes, temperature10.0) pred torch.argmax(logits, dim1).cpu().numpy() acc (pred np.array(que_y)).mean() acc_list.append(acc) return np.mean(acc_list), np.std(acc_list)评估时优先选测试域协议。比如用某一负载工况的数据作为测试域支持集只从测试域随机抽 few-shot 样本查询集来自其他负载工况这样能验证注意力机制在工况迁移下是否还稳定。5.2 从 ECA 权重里提取频率敏感通道训练好的注意力权重不是黑盒。ECA 模块的一维卷积输出经过 sigmoid 后每个通道都有一个 0 到 1 之间的重要性分数。把该分数应用到测试域信号上统计所有测试样本的均值就能得到模型最依赖的通道列表def extract_attention_weights(encoder, data_tensor): # data_tensor 形状 [N, 1, 1024] att_weights [] x data_tensor with torch.no_grad(): x encoder.stem(x) for layer in [encoder.layer1, encoder.layer2, encoder.layer3]: y layer.conv2(layer.bn2(layer.conv1(x))) # 简化调用实际应保留完整前向逻辑 eca layer.eca pooled eca.avg_pool(y) pooled pooled.squeeze(-1).unsqueeze(1) pooled eca.conv(pooled) w torch.sigmoid(pooled.squeeze(1).unsqueeze(-1)).squeeze(-1) att_weights.append(w.mean(dim0).cpu().numpy()) x layer(x) return att_weights拿到每个通道的权重后和模型第一层卷积的滤波器频率响应计算对应关系就能定位出到底是哪个频段对故障判别的贡献最大。这个验证技巧在生产项目中特别有价值如果模型关注的高权重频段覆盖了已知的轴承故障特征频率比如外圈故障特征频率及其倍频就可以确认模型确实学到了物理意义上的故障模式而不是靠偶然的数据偏差取胜。本文还有配套的精品资源点击获取