简介面向少量样本故障诊断场景的PyTorch实现结合注意力机制提升小样本下的特征提取与泛化能力适合故障诊断方向的研究生、算法工程师及相关开发者参考学习。压缩包共18个文件包含10个MAT格式数据文件与8个Python源码文件包体大小10.89MBMAT数据覆盖0-9号样本Python脚本涵盖模型训练、早停、标签平滑、AdaBN等完整流程。资源集成1D信号注意力机制、AMSGradP优化器、1D-Meta-ACON激活函数与1D-Grad-CAM可视化方法并特别加入GAP全局平均池化来增强BiGRU/BiLSTM输出处理属于一套相对完整的小样本故障诊断改进方案。代码注释清晰、模块化程度高便于按需替换数据集或网络组件可用于复现实验、对比消融研究或快速搭建自己的故障诊断基线。目前已有358人学习使用对正在探索注意力机制与少样本诊断结合的读者有直接参考价值。1. 从一次注意力模块误放开始的小样本故障诊断实验做故障诊断的人很容易被“注意力机制”四个字带偏以为把 CBAM 插到哪个位置都能提点尤其是一维振动信号。实际上一维信号的时间位置语义和二维图像的空间语义差别很大把二维注意力原样搬过来多数时候只是增加参数少样本下还会更早过拟合。这个仓库的做法是把注意力、池化、归一化和优化器都针对“小样本”重新调过0.mat到9.mat是按故障类别归档的样本oneD_CS_attention.py是改造后的一维通道-空间注意力GAP 接在 BiGRU 之后是我在实验里收益最大的一处小改动加上 AdaBN、标签平滑和 EarlyStopping 让模型在小批量训练时更稳。适合正在做轴承、齿轮箱等旋转机械故障诊断手头每类样本只有几十到几百个片段、又不想直接上大模型的人。读这份代码前至少要能跑通 PyTorch 训练循环并且对一维信号做过分帧切片。2. 小样本数据组织从 0.mat 到训练/验证集的划分少样本诊断里测试集怎么切有时比模型结构更影响结论。如果只是把所有样本随机分成 8:2相邻滑窗的强相关性会让验证集虚高最后的部署效果反而对不上实验指标。这一章先把.mat文件的读取和划分规则说清楚后面训练脚本才不是黑盒。2.1 .mat 文件的读取路径与字段约定从项目结构看每个.mat文件对应一类样本datasave.py负责把原始振动波形切片后落盘。常见的保存字段是data和labeldata形状一般是[样本数, 序列长度]label是该组样本对应的类别编号读取时必须做降维和类型转换否则后续CrossEntropyLoss会对不上维度。import scipy.io as sio import numpy as np def read_mat(path): mat sio.loadmat(path, squeeze_meTrue, struct_as_recordFalse) data np.asarray(mat[data], dtypenp.float32) label int(np.asarray(mat[label]).reshape(-1)[0]) # 有的 .mat 里 data 是 [N, 1, L]需要先压成 [N, L] if data.ndim 3: data data[:, 0, :] return data, labelloadmat默认会把一维标签读成[1, 1]或[N, 1]我用reshape(-1)[0]强制取标量。data在保存时也可能被写成[N, 1, L]这种三层结构在 PyTorch 里不好直接和Conv1d对接先压维。读取后建议立即打印一遍data.shape和label的分布少样本最怕数据读错而训练流程不报错这类错误会潜伏到验证阶段才暴露。2.2 小样本拆分前的滑窗与重叠率控制一维振动信号通常是一条几十万点的连续记录直接拿整段做样本既浪费又容易欠拟合。标准做法是滑窗切片窗口长度选 1024 或 2048重叠率控制在 0.5 到 0.75 之间。重叠率越高切出的样本越多但相邻窗口的相关系数也越高这会直接污染验证集。def sliding_window(x, window_len1024, overlap0.75): step max(1, int(window_len * (1 - overlap))) if len(x) window_len: return np.zeros((0, window_len), dtypenp.float32) n (len(x) - window_len) // step 1 idx np.arange(window_len)[None, :] step * np.arange(n)[:, None] return x[idx]切片后再分组不要直接对原始长序列随机划分。step是每次滑动的步长overlap0.75时步长是窗口长度的四分之一样本量能膨胀四倍左右。对小样本任务来说这是最简单的数据增强但副作用是相邻样本几乎一样训练集的多样性被高估。工程上我一般先切窗再按原始记录 ID 分组最后用组级划分切验证集这样才能近似真实部署时“见到新数据”的条件。参数常见取值对小样本的影响窗口长度1024 / 2048太短丢周期信息太长增加计算量且样本数变少重叠率0.5~0.75提高样本数但带来高相关性需配合组级划分每类样本数30~200属于少样本区间需要正则化和早停配合2.3 分层划分与数据泄漏规避类别不均衡在小样本里经常出现比如正常样本多、故障样本少。随机划分可能让某一类在训练集里只剩个位数。用分层抽样保证训练集和验证集里每个类的比例一致from sklearn.model_selection import GroupShuffleSplit gss GroupShuffleSplit(n_splits1, test_size0.2, random_state42) train_idx, val_idx next(gss.split(data, labels, groupsrecord_ids)) x_train, x_val data[train_idx], data[val_idx] y_train, y_val labels[train_idx], labels[val_idx]groupsrecord_ids是每个样本对应的原始记录编号滑窗前给每个长序列分配唯一 ID切出的窗口继承这个 ID。GroupShuffleSplit会保证同一条记录的窗口全部进同一个集合避免窗口级别的信息泄漏。这一步看似简单但对最终指标的可靠度影响最大。随机切分时验证集可能是乐观的部署到新的工况数据上精度会明显下跌。3. 一维注意力机制实现oneD_CS_attention 与时序池化这一章是仓库的核心。oneD_CS_attention.py实现了一维信号的通道注意力和时序注意力和图像里的 CBAM 用法不同输入是[B, C, L]其中L是一维序列长度。很多现成代码直接用Conv2d的注意力模块放在一维信号上会直接报维度错或者因为transpose搞乱通道顺序导致训练不收敛。3.1 通道注意力与空间注意力的一维化改造通道注意力的思路是对每个通道做全局池化再用全连接层学出一组通道权重。一维场景下池化是在L维度上进行squeeze时要特别小心维度顺序。import torch import torch.nn as nn class ChannelAttention1d(nn.Module): def __init__(self, channels, reduction8): super().__init__() self.fc nn.Sequential( nn.Linear(channels, channels // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, channels, biasFalse), nn.Sigmoid(), ) def forward(self, x): # x: [B, C, L] y x.mean(dim-1) # [B, C] attn self.fc(y).unsqueeze(-1) # [B, C, 1] return x * attnmean(dim-1)是把整个时间维度压成一个标量等价于对一维信号做全局平均池化。注意unsqueeze(-1)必须在最后这样attn才能和x在C维上广播相乘。reduction控制中间层的压缩比例少样本场景我一般取 8太小会减少参数量但表达能力下降太大会让注意力模块本身过拟合。空间注意力在一维场景里其实是“时序注意力”在L维上生成一个掩码突出或抑制不同时间位置的特征。实现上用一维卷积代替二维卷积class TemporalAttention1d(nn.Module): def __init__(self, kernel_size7): super().__init__() self.conv nn.Conv1d(2, 1, kernel_size, paddingkernel_size // 2, biasFalse) self.sigmoid nn.Sigmoid() def forward(self, x): avg_map torch.mean(x, dim1, keepdimTrue) max_map torch.amax(x, dim1, keepdimTrue) out torch.cat([avg_map, max_map], dim1) # [B, 2, L] attn self.sigmoid(self.conv(out)) # [B, 1, L] return x * attn这里把通道维压缩成两路一路均值、一路最大值。最大值能捕捉冲击特征均值能反映整体能量分布这是故障诊断里常见的两个互补统计量。Conv1d的padding要手动算好kernel_size7时padding3避免序列长度在卷积后变短。组合起来就是oneD_CS_attention先通道后时序与 CBAM 的次序一致。3.2 BiGRU/BiLSTM 后接 GAP 为什么对小样本更稳源码贡献里特别提到很多模型在 BiGRU/BiLSTM 后面直接取最后一个时间步的隐状态而不是做 GAP。我在复现时对比过这个选择对少样本的影响比想象中大得多。import torch.nn.functional as F # 假设 bigru 输出 [B, L, H]H 是双向拼接后的维度 out, _ self.bigru(x) # [B, L, H] out out.permute(0, 2, 1) # [B, H, L] gap_out F.adaptive_avg_pool1d(out, 1).squeeze(-1) # [B, H] max_out F.adaptive_max_pool1d(out, 1).squeeze(-1) # [B, H] x torch.cat([gap_out, max_out], dim1)permute把L转到最后一维才能用adaptive_avg_pool1d。adaptive_avg_pool1d(out, 1)等价于全局平均池化在L维度上取平均。取平均的意义在于双向 GRU 在不同时间步上看到的故障特征不同末尾隐状态容易丢失早期出现的冲击特征平均池化相当于把所有时间步的证据汇总是一种天然的集成效果。池化方式代码对小样本的隐患我的取舍取最后时间步out[:, -1, :]信息只看尾部早期冲击特征丢失不推荐Max Pooladaptive_max_pool1d容易放大单点噪声作为补充特征GAPadaptive_avg_pool1d弱化短时强冲击首选稳定性最高用squeeze(-1)而不是squeeze()因为batch_size恰好为 1 时squeeze()会把 batch 维也去掉后续全连接层直接报错。把 GAP 和 Max Pool 拼起来可以同时保留平均能量和峰值特征少样本下比单用其中一种更稳。3.3 把注意力模块嵌入骨干网络的位置注意力模块放在哪里是个容易被忽略的问题。常见做法是把oneD_CS_attention接在 BiGRU 之后、GAP 之前或者接在 GAP 之后。我的经验是放在 BiGRU 之后、GAP 之前更合理因为这时张量还是[B, C, L]通道注意力和时序注意力可以同时作用于时间维。如果先 GAP 再塞注意力时间信息已经被压缩掉时序注意力模块就失去意义了。class FaultModel(nn.Module): def __init__(self, in_channels, seq_len, num_classes): super().__init__() self.cnn nn.Sequential( nn.Conv1d(in_channels, 16, kernel_size3, padding1), nn.BatchNorm1d(16), nn.ReLU(inplaceTrue), ) self.bigru nn.GRU(16, 32, batch_firstTrue, bidirectionalTrue) self.attn oneD_CS_attention(64) self.classifier nn.Linear(64, num_classes) def forward(self, x): x self.cnn(x) # [B, 16, L] x, _ self.bigru(x.transpose(1, 2)) # [B, L, 64] x self.attn(x.transpose(1, 2)) # [B, 64, L] x F.adaptive_avg_pool1d(x, 1).squeeze(-1) return self.classifier(x)bigru的输入要转成[B, L, C]因为 GRU 默认把序列维放在第 1 维。bidirectionalTrue时输出特征数是hidden_size * 2这里64。注意力模块的channels参数必须等于 BiGRU 输出的64否则全连接层维度对不上。少样本下 BiGRU 的隐层不要设太大32 已经能覆盖大多数轴承故障特征再大就依赖更多训练数据。4. 少样本训练稳定化AdaBN、标签平滑、早停与优化器变体小样本训练最大的问题是模型容易记住训练集上的局部噪声。单纯加正则化可能还不够这个仓库把几个组件单独封装成脚本组合起来效果更明显。这一章逐个讲清楚每个组件的适用场景和参数。4.1 AdaBN 重新估计 BatchNorm 统计量批量归一化在训练时用当前 batch 的均值和方差在测试时用滑动平均得到的 running stats。如果训练数据和测试数据来自不同工况比如负载、转速变了running stats 就会偏移诊断精度下跌。AdaBN 的核心是验证时临时用当前输入 batch 的统计量替代 running stats。import torch.nn.functional as F def adabn_forward(x, bn: nn.BatchNorm1d): mean x.mean(dim(0, 2), keepdimTrue) var x.var(dim(0, 2), unbiasedFalse, keepdimTrue) return F.batch_norm( x, mean.squeeze(), var.squeeze(), bn.weight, bn.bias, trainingFalse, momentum0.0, epsbn.eps, )注意mean.squeeze()是把[1, C, 1]变回[C]F.batch_norm要求统计量是一维张量。trainingFalse表示不更新 running stats只是借用当前 batch 的统计量做归一化。这个做法对少量样本风险在于单个 batch 的均值和方差不一定可靠一般在 16 以上才推荐使用只有 4 个样本时会引入额外噪声。4.2 标签平滑解决硬标签过拟合故障诊断里标签通常是 one-hot模型会为了把真实类的 logit 拉到很大而变得过度自信。少量样本下这种过度自信会直接转化为验证集精度抖动。标签平滑就是把 one-hot 目标从 1 和 0 变成1 - smoothing和smoothing / (num_classes - 1)。def smooth_cross_entropy(logits, target, num_classes, smoothing0.1): n logits.size(0) one_hot torch.full((n, num_classes), smoothing / (num_classes - 1)) one_hot one_hot.to(logits.device) one_hot.scatter_(1, target.unsqueeze(1), 1.0 - smoothing) log_probs F.log_softmax(logits, dim1) return (-one_hot * log_probs).sum(dim1).mean()target.unsqueeze(1)是必须的scatter_在dim1方向需要目标索引和矩阵同维。smoothing0.1是经验值让每个负类分到0.1 / 9左右的梯度模型不会被单一类别带偏。如果平滑值调到 0.3 以上损失下降会变慢且测试精度可能不升反降因为模型学到的类别边界太模糊。4.3 EarlyStopping 的 patience 与 min_delta少样本训练里验证损失往往不是平滑下降的而是带明显噪声。patience设得太小容易在第一个小波动时就中断训练设太大又失去早停意义。我一般配合min_delta来过滤噪声。class EarlyStopping: def __init__(self, patience15, min_delta1e-4): self.patience patience self.min_delta min_delta self.counter 0 self.best_score None def step(self, val_loss, model, save_path): if self.best_score is None or val_loss self.best_score - self.min_delta: self.best_score val_loss self.counter 0 torch.save(model.state_dict(), save_path) return False self.counter 1 return self.counter self.patiencemin_delta1e-4表示验证损失必须比历史最优低至少 0.0001 才认为是真正的提升这样小幅波动不会触发 checkpoint 覆盖。patience15表示连续 15 个 epoch 没有明显进展就停止。少样本下建议 checkpoint 保存最优模型而不是最后一个 epoch否则最终模型往往已经不是验证集上最好的那个状态。4.4 AdamP_amsgrad 与 oneD_Meta_ACON 的选型建议AdamP_amsgrad.py本质上是把 AdamP 的权值解耦思想引入 Adam并叠加 AMSGrad 的长期梯度缓存。我使用时的默认配置是这样的optimizer AdamP_amsgrad.AdamP_amsgrad( model.parameters(), lr3e-4, betas(0.9, 0.999), weight_decay1e-2, amsgradTrue, )weight_decay在 AdamP 里是解耦的和AdamW类似对少样本更友好。学习率不要设到1e-3以上因为样本少时梯度方差大过高的学习率会让注意力模块的通道权重震荡。amsgradTrue让每个参数维度记录历史最大二阶矩避免罕见的大梯度冲掉已学到的特征。oneD_Meta_ACON.py实现的是自适应激活函数根据输入动态切换 ReLU 和线性激活的比例。少样本下注意力模块之后的特征分布很敏感固定激活函数可能不够灵活。下面是适配一维信号的最小实现class MetaAcon1d(nn.Module): def __init__(self, channels, reduction4): super().__init__() self.avg_pool nn.AdaptiveAvgPool1d(1) self.fc nn.Sequential( nn.Linear(channels, channels // reduction, biasFalse), nn.ReLU(inplaceTrue), nn.Linear(channels // reduction, 1, biasFalse), ) self.sigmoid nn.Sigmoid() def forward(self, x): b, c, _ x.size() beta self.sigmoid(self.fc(self.avg_pool(x).view(b, c))) beta beta.view(b, 1, 1) return (1 - beta) * F.relu(x) beta * xbeta是从全局通道统计里学出来的一个标量每个 batch 动态变化。当beta接近 0 时激活退化为 ReLU接近 1 时退化为线性函数。这个模块建议放在注意力之后而不是卷积之后因为注意力模块输出的特征需要更精细的缩放。少样本下reduction取 4 或 8 即可太小的全连接层学不到通道差异。5. model_train.py 实战日志解读与 1D-Grad-CAM 验证这一章直接讲怎么把模型训练起来并判断是否学对了。model_train.py把前面的组件串起来核心流程是配置数据路径、定义模型、初始化优化器和调度器、循环训练并在验证集上做早停判断。我在跑这类代码时会额外加一个torch.save的完整 checkpoint包括优化器状态和 epoch 数方便从断点继续。for epoch in range(epochs): model.train() train_loss 0.0 for xb, yb in train_loader: xb, yb xb.cuda(), yb.cuda() optimizer.zero_grad() out model(xb) loss smooth_cross_entropy(out, yb, num_classes) loss.backward() optimizer.step() train_loss loss.item() model.eval() val_loss validate(model, val_loader) if early_stopping.step(val_loss, model, save_path): break注意训练阶段要model.train()验证阶段要model.eval()因为BatchNorm1d在两种模式下的行为不同。验证函数里用torch.no_grad()包裹否则每个 batch 都会多算一遍反向图内存占用直接翻倍。小样本下 batch size 通常取 16 或 32num_workers设成 0 或 2过高的num_workers在 Windows 环境下容易触发 DataLoader 卡死。日志里出现三种异常形态要特别注意训练损失不断下降但验证损失在第 10 个 epoch 后回升说明过拟合应该增强标签平滑或增大weight_decay两个损失同时居高不下说明学习率过大或模型容量不够把学习率降到1e-4再试验证损失曲线抖动剧烈但整体不下降往往不是模型问题而是数据划分没有按记录 ID 分组窗口泄漏导致指标失真。最后用 1D-Grad-CAM 验证注意力到底聚焦在哪里。1D-Grad-CAM 比普通 Grad-CAM 对多目标和小目标更敏感适合故障冲击特征较弱的场景def grad_cam_plusplus_1d(model, x, target_layer, target_class): act, grad {}, {} def fwd_hook(m, i, o): act[value] o def bwd_hook(m, gi, go): grad[value] go[0] h1 target_layer.register_forward_hook(fwd_hook) h2 target_layer.register_backward_hook(bwd_hook) x.requires_grad_() out model(x) model.zero_grad() out[0, target_class].backward() A act[value].squeeze(0) # [C, L] G grad[value].squeeze(0) # [C, L] alpha G.sum(dim-1, keepdimTrue) / (G.sum(dim-1, keepdimTrue) 1e-6) cam torch.relu((alpha * A).sum(dim0)).detach().cpu().numpy() h1.remove() h2.remove() return camalpha的计算是 Grad-CAM 与 Grad-CAM 的主要区别它对小目标的梯度做了加权不会因为某个强脉冲占据主导而忽略分布在整个周期里的微弱故障特征。1e-6防止梯度全零时除零。拿到cam后把它插值到原始信号长度看高激活区域是否对应波形上的冲击段。如果注意力集中在每一个周期的固定相位附近说明模型学到的不是故障频率而是规则噪声这时需要回查数据预处理是否有 50Hz 工频干扰没滤干净。本文还有配套的精品资源点击获取