简介:面向少量样本故障诊断场景的PyTorch实现,结合注意力机制提升小样本下的特征提取与泛化能力,适合故障诊断方向的研究生、算法工程师及相关开发者参考学习。压缩包共18个文件,包含10个MAT格式数据文件与8个Python源码文件,包体大小10.89MB,MAT数据覆盖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和label,data形状一般是[样本数, 序列长度],label是该组样本对应的类别编号,读取时必须做降维和类型转换,否则后续CrossEntropyLoss会对不上维度。
import scipy.io as sio import numpy as np def read_mat(path): mat = sio.loadmat(path, squeeze_me=True, struct_as_record=False) data = np.asarray(mat["data"], dtype=np.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_len=1024, overlap=0.75): step = max(1, int(window_len * (1 - overlap))) if len(x) < window_len: return np.zeros((0, window_len), dtype=np.float32) n = (len(x) - window_len) // step + 1 idx = np.arange(window_len)[None, :] + step * np.arange(n)[:, None] return x[idx]切片后再分组,不要直接对原始长序列随机划分。step是每次滑动的步长,overlap=0.75时步长是窗口长度的四分之一,样本量能膨胀四倍左右。对小样本任务来说这是最简单的数据增强,但副作用是相邻样本几乎一样,训练集的多样性被高估。工程上我一般先切窗,再按原始记录 ID 分组,最后用组级划分切验证集,这样才能近似真实部署时“见到新数据”的条件。
| 参数 | 常见取值 | 对小样本的影响 |
|---|---|---|
| 窗口长度 | 1024 / 2048 | 太短丢周期信息,太长增加计算量且样本数变少 |
| 重叠率 | 0.5~0.75 | 提高样本数但带来高相关性,需配合组级划分 |
| 每类样本数 | 30~200 | 属于少样本区间,需要正则化和早停配合 |
2.3 分层划分与数据泄漏规避
类别不均衡在小样本里经常出现,比如正常样本多、故障样本少。随机划分可能让某一类在训练集里只剩个位数。用分层抽样保证训练集和验证集里每个类的比例一致:
from sklearn.model_selection import GroupShuffleSplit gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42) train_idx, val_idx = next(gss.split(data, labels, groups=record_ids)) x_train, x_val = data[train_idx], data[val_idx] y_train, y_val = labels[train_idx], labels[val_idx]groups=record_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, reduction=8): super().__init__() self.fc = nn.Sequential( nn.Linear(channels, channels // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels, bias=False), 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_size=7): super().__init__() self.conv = nn.Conv1d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): avg_map = torch.mean(x, dim=1, keepdim=True) max_map = torch.amax(x, dim=1, keepdim=True) out = torch.cat([avg_map, max_map], dim=1) # [B, 2, L] attn = self.sigmoid(self.conv(out)) # [B, 1, L] return x * attn这里把通道维压缩成两路:一路均值、一路最大值。最大值能捕捉冲击特征,均值能反映整体能量分布,这是故障诊断里常见的两个互补统计量。Conv1d的padding要手动算好,kernel_size=7时padding=3,避免序列长度在卷积后变短。组合起来就是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], dim=1)permute把L转到最后一维,才能用adaptive_avg_pool1d。adaptive_avg_pool1d(out, 1)等价于全局平均池化,在L维度上取平均。取平均的意义在于,双向 GRU 在不同时间步上看到的故障特征不同,末尾隐状态容易丢失早期出现的冲击特征;平均池化相当于把所有时间步的证据汇总,是一种天然的集成效果。
| 池化方式 | 代码 | 对小样本的隐患 | 我的取舍 |
|---|---|---|---|
| 取最后时间步 | out[:, -1, :] | 信息只看尾部,早期冲击特征丢失 | 不推荐 |
| Max Pool | adaptive_max_pool1d | 容易放大单点噪声 | 作为补充特征 |
| GAP | adaptive_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_size=3, padding=1), nn.BatchNorm1d(16), nn.ReLU(inplace=True), ) self.bigru = nn.GRU(16, 32, batch_first=True, bidirectional=True) 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 维。bidirectional=True时输出特征数是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), keepdim=True) var = x.var(dim=(0, 2), unbiased=False, keepdim=True) return F.batch_norm( x, mean.squeeze(), var.squeeze(), bn.weight, bn.bias, training=False, momentum=0.0, eps=bn.eps, )注意mean.squeeze()是把[1, C, 1]变回[C],F.batch_norm要求统计量是一维张量。training=False表示不更新 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, smoothing=0.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, dim=1) return (-one_hot * log_probs).sum(dim=1).mean()target.unsqueeze(1)是必须的,scatter_在dim=1方向需要目标索引和矩阵同维。smoothing=0.1是经验值,让每个负类分到0.1 / 9左右的梯度,模型不会被单一类别带偏。如果平滑值调到 0.3 以上,损失下降会变慢,且测试精度可能不升反降,因为模型学到的类别边界太模糊。
4.3 EarlyStopping 的 patience 与 min_delta
少样本训练里验证损失往往不是平滑下降的,而是带明显噪声。patience设得太小容易在第一个小波动时就中断训练,设太大又失去早停意义。我一般配合min_delta来过滤噪声。
class EarlyStopping: def __init__(self, patience=15, min_delta=1e-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_delta=1e-4表示验证损失必须比历史最优低至少 0.0001 才认为是真正的提升,这样小幅波动不会触发 checkpoint 覆盖。patience=15表示连续 15 个 epoch 没有明显进展就停止。少样本下建议 checkpoint 保存最优模型而不是最后一个 epoch,否则最终模型往往已经不是验证集上最好的那个状态。
4.4 AdamP_amsgrad 与 oneD_Meta_ACON 的选型建议
AdamP_amsgrad.py本质上是把 AdamP 的权值解耦思想引入 Adam,并叠加 AMSGrad 的长期梯度缓存。我使用时的默认配置是这样的:
optimizer = AdamP_amsgrad.AdamP_amsgrad( model.parameters(), lr=3e-4, betas=(0.9, 0.999), weight_decay=1e-2, amsgrad=True, )weight_decay在 AdamP 里是解耦的,和AdamW类似,对少样本更友好。学习率不要设到1e-3以上,因为样本少时梯度方差大,过高的学习率会让注意力模块的通道权重震荡。amsgrad=True让每个参数维度记录历史最大二阶矩,避免罕见的大梯度冲掉已学到的特征。
oneD_Meta_ACON.py实现的是自适应激活函数,根据输入动态切换 ReLU 和线性激活的比例。少样本下注意力模块之后的特征分布很敏感,固定激活函数可能不够灵活。下面是适配一维信号的最小实现:
class MetaAcon1d(nn.Module): def __init__(self, channels, reduction=4): super().__init__() self.avg_pool = nn.AdaptiveAvgPool1d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channels // reduction, 1, bias=False), ) 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 或 32,num_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, keepdim=True) / (G.sum(dim=-1, keepdim=True) + 1e-6) cam = torch.relu((alpha * A).sum(dim=0)).detach().cpu().numpy() h1.remove() h2.remove() return camalpha的计算是 Grad-CAM++ 与 Grad-CAM 的主要区别,它对小目标的梯度做了加权,不会因为某个强脉冲占据主导而忽略分布在整个周期里的微弱故障特征。+1e-6防止梯度全零时除零。拿到cam后,把它插值到原始信号长度,看高激活区域是否对应波形上的冲击段。如果注意力集中在每一个周期的固定相位附近,说明模型学到的不是故障频率而是规则噪声,这时需要回查数据预处理是否有 50Hz 工频干扰没滤干净。
本文还有配套的精品资源,点击获取