先聊个实际的场景:设备振动监测数据里,正常样本常年攒了几万条,故障样本却可能只有几十条,还往往要老师傅花大量时间人工判读才能标出来。这种极度不均衡、标注成本又高的现状,直接把常规CNN分类任务按在地上摩擦——模型很容易记住少数故障样本的噪声细节,泛化能力近乎为零。我之前接手一个轴承故障诊断项目时就卡在这里,后来把连续小波变换(CWT)、多尺度注意力机制和原型网络(Prototypical Network)三套东西组合到一起,在每类只有5个样本的设定下,把多分类准确率从七成出头拉到了九成五上下。这篇就把这套方法从原理到Pytorch实现完整拆一遍,适合正在做小样本故障诊断、想了解元学习落地方案,或者准备把信号时频分析和深度学习结合起来写论文的读者参考。
1. 小样本轴承诊断:为什么故障数据少时常规CNN会失效
1.1 我在实际项目里遇到的小样本困境
真实工业场景里,轴承绝大多数时间都在正常运转,故障样本往往来自周期性点检、突发停机或者拆机检修,能留下完整记录的少之又少。即便拿到了带标签的故障数据,同一类故障在不同转速、不同负载下呈现的振动特征也差别很大,这就导致"样本少"和"类内差异大"两个问题同时出现。
常规做法是直接把一维振动信号送进1D CNN,或者转成频谱再送进2D CNN做分类。样本充足时效果尚可,一旦切到小样本设定,问题立刻暴露:模型会把支撑集里那几个样本的特征细节当成类别本质,测试时换一个工况就误判。本质原因是常规分类头(全连接层加softmax)需要大量样本才能拟合出稳定的决策边界,样本数不够时边界极易过拟合。
1.2 标题方案的攻防逻辑:三件事各管一段
这套标题方案里其实隐藏了一条清晰的攻防链条。
连续小波变换负责"信息载体":把一维时序信号展开成二维时频图,让CNN能看到频率成分随时间的变化。轴承故障的特征频率通常和转动频率、滚动体尺寸强相关,时频图上不同故障会出现不同频带和周期纹理,这对后续特征学习是决定性优势。
多尺度注意力机制负责"特征表达":时频图上同时存在宽频冲击、窄带共振和噪声底,单一感受野很难兼顾。多尺度分支并行提取不同范围的局部模式,注意力再按通道和空间位置加权,让模型把算力集中在故障特征真正存在的频带和时间段上。
原型网络负责"小样本决策":它不再为每个类别学习一组独立的softmax权重,而是把支撑集中同类样本的特征取平均生成类别原型,查询样本按与各原型的距离做分类。每当新类别出现,不需要重新训练分类头,只要重新计算原型就行,这是小样本场景里最实用的特性。
2. 连续小波变换:把一维振动信号转成"机器看得懂的时频图"
2.1 CWT相对STFT的核心优势
短时傅里叶变换(STFT)是最常用的时频分析工具,但它有一个硬伤:窗口长度固定,时间分辨率和频率分辨率无法同时保证。窗短了,频率分辨率差;窗长了,时间定位差。轴承故障信号里既有持续时间极短的冲击成分,又有周期性出现的调制边带,STFT往往顾此失彼。
连续小波变换的思路是让窗宽随频率自适应缩放:分析低频时用长窗换取频率分辨率,分析高频时用短窗保住时间定位。这个特性对应到小波尺度上,就是大尺度看低频、小尺度看高频,正好契合振动信号里"低频故障特征频率谱线清晰、高频共振冲击短促"的物理特性。
实际工程中我一般直接用Morlet复小波作母小波。它本质是一个高斯窗包裹的复指数信号,在时频平面的局部化特性好,和振动信号的衰减正弦形态也比较接近。用pywt库一句就能搞定,但有几个参数必须先想清楚。
import pywt import numpy as np fs = 12000 # 采样率,单位Hz data = np.random.randn(2048) # 一段长度为2048的振动信号 scales = np.linspace(3, 200, 128) # 尺度序列,具体范围见下文换算 coef, freqs = pywt.cwt(data, scales, wavelet='morl', sampling_period=1.0/fs)2.2 尺度参数、中心频率与频率换算
CWT的输出是尺度a和时间位移b的二维函数,但尺度本身不是频率,很多初学者在这里栽跟头。换算关系是:尺度a对应的伪频率 (f_a = f_c \times f_s / a),其中 (f_c) 是母小波的中心频率,Morlet小波在pywt里大约取0.8125。
我在某采样率12kHz的数据上做过一次实际换算:如果想覆盖50Hz到3kHz的频段,那么最小尺度 (a_{min} = 0.8125 \times 12000 / 3000 \approx 3.25),最大尺度 (a_{max} = 0.8125 \times 12000 / 50 \approx 195)。所以我习惯把尺度范围设成3到200。换句话说,121.5Hz是一个粗略的参考,最终是以故障特征频率的最大公约数落入你要覆盖的频带为准。如果你想覆盖到5kHz高频共振区,就把最小尺度压到2约当拟值。
这里有个取尺度的技巧:用对数间隔采样比线性间隔更合理,因为CWT的频率分辨率本身和对数频率轴更匹配,对数尺度下相邻尺度对应的频率间隔相对均匀。换成代码就是np.logspace(np.log10(3), np.log10(200), 128)。
2.3 时频图预处理:采样点数、幅值取模、离线存储
这一步的主要目标是让后续CNN输入质量稳定。
首先,样本长度。截取窗口要覆盖至少2到3个转频周期,让时频图能稳定呈现故障特征频率的周期纹理。比如工业场景下转频25Hz,一个周期是40ms,12kHz采样率下就是480个点,取2048个点大约覆盖4个周期,比较稳妥。窗口再长也没大意义,反而稀释故障冲击成分的比例。
其次,CWT输出是复数矩阵,实际喂给网络的通常取模值(幅值谱),也就是np.abs(coef)。相位信息在轴承故障诊断里贡献有限,而模值能突出能量集中区域,符合CNN看纹理的认知习惯。尺寸上,128个尺度乘2048个时间点是128×2048的大矩阵,直接送CNN参数爆炸。我的做法是保留128×2048后转成224×224的灰度图,或者先沿时间轴方向做分段池化压到128×256再resize。resize的思路没有想象中那么神秘——它只是把时频图变成CNN更易消化的输入尺寸,代价是牺牲掉一部分频带和时间的细粒度信息,但在小样本任务里模型本身就不该依赖过细的信息。
最后一个经验:一定把所有样本的时频图离线算好存成npy,训练时直接读图,而不是在Dataloader里实时做CWT。一张2048点信号的CWT计算耗时在毫秒级,样本多了叠加起来会吃掉大量训练时间,离线缓存能省下至少一多半的训练耗时。
3. 多尺度注意力特征提取器的搭建思路
3.1 多尺度分支设计:不同感受野各管一段频带
时频图上的故障模式很有意思:齿轮啮合频率和轴承外圈故障特征频率常表现为窄带平行谱线,属于细条纹纹理;故障冲击在共振区激起的响应则是大范围亮斑,属于粗粒度的区域特征。单一固定卷积核很难同时捕捉这两种形态。
多尺度分支的经典搭法是用不同卷积核或不同空洞率的卷积并行提取特征。我习惯用三个分支:
- 1×1卷积,负责保留像素级细节,相当于只做通道混合;
- 3×3卷积,负责中等范围的纹理过渡;
- 5×5卷积,或者用空洞率2的3×3卷积来等效大感受野,负责捕捉宽频冲击和共振亮斑。
空洞卷积的好处是感受野大了参数却不多,对小样本模型更友好。你可以把这三个分支理解成"一个仔细看纹理、一个看中等范围、一个站远看整体",三双眼睛看到的特征拼在一起总比一双眼镜要可靠。
3.2 注意力加在哪个环节,才不会白加
注意力不是随便往网络里塞一块就有效。我在最开始直接把SE模块加在骨干网络每个卷积后面,结果准确率几乎没有变化,后来复盘才想明白:小样本场景下,特征图的方差本来就大,通道注意力如果加在特征还没形成稳定语义的浅层,只会放大部分噪声通道。
经过几轮实验,我的结论是:先把多尺度分支的输出融合成一张特征图,再在融合之后加注意力,效果最明显。原因是融合后的特征图里,每个通道已经对应了某个尺度分支提取到的抽象模式,注意力此时做的是"决定哪个尺度的贡献更重要",符合我们设计多尺度的初衷。
通道注意力部分我不推荐标准SE的降维再升维结构,维度压缩会丢失通道间依赖信息,改成轻量ECA(一维卷积直接对通道维做局部交互)效果更好、参数也更少。空间注意力则用1×1卷积加sigmoid生成逐像素权重,让模型能主动忽略噪声频带对应的区域。
3.3 分支特征对齐与融合方式
三个分支的卷积核尺寸不同,但都通过padding让输出时空尺寸保持一致,这样融合时就不需要插值对齐。通道数上,每个分支输出32维通道,拼接后就是96维,再经过一个1×1卷积压缩到64维。这个1×1卷积不只是降维,它能学到一个跨分支的加权组合规则,比简单相加更灵活。
融合之后接注意力,再接全局平均池化,得到一维特征向量。这里有个容易忽视的点:全局平均池化之前不要接太多全连接层。小样本任务里全连接层的参数量是整个模型的主要风险源,我把分类表征维度控制在64到128之间,池化后直接输出作为原型网络的特征嵌入。
我把这个特征提取器的核心代码贴在下面,结构清晰,可以直接作为baseline:
import torch import torch.nn as nn import torch.nn.functional as F class ECAAttention(nn.Module): def __init__(self, channels, k=3): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=k // 2, bias=False) def forward(self, x): y = self.avg_pool(x).squeeze(-1).squeeze(-1) y = self.conv(y.unsqueeze(1)).squeeze(1) y = torch.sigmoid(y).unsqueeze(-1).unsqueeze(-1) return x * y class SpatialAttention(nn.Module): def __init__(self, channels): super().__init__() self.conv = nn.Conv2d(channels, 1, kernel_size=1) self.sigmoid = nn.Sigmoid() def forward(self, x): w = self.sigmoid(self.conv(x)) return x * w class CWTEncoder(nn.Module): def __init__(self, in_channels=1, feat_dim=128): super().__init__() self.b1 = nn.Sequential(nn.Conv2d(in_channels, 32, 1), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.b3 = nn.Sequential(nn.Conv2d(in_channels, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.b5 = nn.Sequential(nn.Conv2d(in_channels, 32, 5, padding=2), nn.BatchNorm2d(32), nn.ReLU(inplace=True)) self.merge = nn.Sequential(nn.Conv2d(96, 64, 1), nn.BatchNorm2d(64), nn.ReLU(inplace=True)) self.eca = ECAAttention(64) self.sa = SpatialAttention(64) self.head = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, feat_dim)) def forward(self, x): x = torch.cat([self.b1(x), self.b3(x), self.b5(x)], dim=1) x = self.merge(x) x = self.eca(x) x = self.sa(x) x = self.head(x) return x4. 原型网络的小样本分类机制与训练逻辑
4.1 原型网络在做什么
原型网络的核心思想一句话就能讲清楚:每个类别用支撑集(support set)里样本特征的平均值作为该类的"原型"(prototype),待预测的查询样本看它离哪个原型近,就归为哪类。
这里要强调一下"近"这个字。原型网络使用的距离是欧氏距离,而不是我们做分类任务时更习惯的余弦距离。原论文的基础设定是:在特征空间中,每个类别可以用一个多元高斯分布近似,支撑集样本的均值是该分布均值的无偏估计。欧氏距离的平方与高斯负对数似然直接对应,在理论上能承接这个概率解释。而余弦距离只度量方向、忽略长度,它会丢掉一部分类内方差结构信息。
实际中我也测过替换成余弦距离的效果,在小样本轮次下和欧氏距离差距小则0.5个点,大则两三个点,但欧氏距离版本训练更稳定,尤其是在特征还不是特别收敛的早中期轮次。所以如果你的代码直接用余弦距离,建议至少改成平方欧氏距离 ( |f(x_q) - c_k|^2 ) 作为比较基准。
4.2 Episode训练机制:元学习视角的改造
原型网络的训练不用传统batch,而用episode。每个episode模拟一次"当前只有少量标注样本"的任务:
- 从训练集的类别池里随机抽N个类别(N-way);
- 每类抽出K个样本作为支撑集(K-shot);
- 每类再抽出若干样本作为查询集(query set);
- 用支撑集计算原型,用查询集计算损失并回传梯度。
这个机制的巧妙之处在于,模型在训练阶段看到的每轮都是小样本任务,自然逼迫特征编码器去学习"如何在小样本下形成可分可分聚类"的表征,而不是记忆某个样本。等到部署时,面对新类别或者更小的标注集合,模型根本不需要改结构,重新算几个原型就能工作。
def proto_loss(support_feats, query_feats, support_labels, n_way, n_shot): # support_feats: [n_way * n_shot, d] prototypes = support_feats.view(n_way, n_shot, -1).mean(dim=1) # [n_way, d] # 平方欧氏距离 dists = torch.cdist(query_feats, prototypes) ** 2 # [n_query, n_way] logits = -dists # query_labels 形状为 [n_query] loss = F.cross_entropy(logits, query_labels) acc = (logits.argmax(dim=1) == query_labels).float().mean() return loss, acc4.3 支撑集、查询集的构成与参数选择
支撑集每类K个样本,K的取值直接决定任务难度。K=1时原型退化成单个样本的特征,方差最大;K=5是比较常见的折中;K=10时原型更稳定,但真实场景里很难攒出这么多标注样本。我在项目里一般以K=5为主报告指标,同时测一组K=1、K=10看趋势。
查询集数量也不宜过少,否则损失方差大。每类查询样本我至少取15个,一个episode里支撑集N×K个、查询集N×15个,GPU一个batch完全装得下。类别数N的选择建议和实际部署场景一致——如果目标是把故障分成10类,episode里的N就设成10,而不是论文里习惯的5。这与微调阶段任务分布一致,避免训练和部署之间的gap。
5. Pytorch实现:从数据迭代器到训练循环
5.1 数据接口设计:把一个episode当一个batch
因为训练单元从单样本换成了episode,Dataloader的设计也要跟着改。有两种做法,我用过都觉得不错:第一种是自定义Dataset返回单个样本,在循环里用随机类别索引手动拼装episode;第二种是自定义一个耗时较长的collate_fn,每次从数据集里随机采样N-way K-shot支撑集和查询集。第一种更灵活、好debug,推荐新手用。
我这里给出一个简化的采样逻辑。
class EpisodeSampler: def __init__(self, data_by_class, n_way, n_shot, n_query): self.class_ids = list(data_by_class.keys()) self.data_by_class = data_by_class self.n_way = n_way self.n_shot = n_shot self.n_query = n_query def sample(self): selected = np.random.choice(self.class_ids, self.n_way, replace=False) support_x, support_y, query_x, query_y = [], [], [], [] for i, cls in enumerate(selected): samples = self.data_by_class[cls] idx = np.random.choice(len(samples), self.n_shot + self.n_query, replace=False) support_x.extend(samples[idx[:self.n_shot]]) support_y.extend([i] * self.n_shot) query_x.extend(samples[idx[self.n_shot:]]) query_y.extend([i] * self.n_query) return (torch.stack(support_x), torch.tensor(support_y), torch.stack(query_x), torch.tensor(query_y))这个采样器有个关键细节:支撑集和查询集来自同一个随机抽取的索引,但二者不重叠。如果直接先抽支撑集再独立抽查询集,同一个样本可能同时出现在支撑和查询里,小样本下会虚高指标。我见过好几个博主分享的代码在这里栽跟头,实际部署时有泄露的模型性能会假得离谱。
5.2 训练循环:完整的元学习迭代
训练循环本身的代码量和普通分类训练差不多,核心区别在每轮都要构造episode并分别计算支撑集和查询集的特征。
model = CWTEncoder(in_channels=1, feat_dim=128) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.5) for epoch in range(60): train_loss, train_acc = [], [] for _ in range(200): # 每epoch 200个episode model.train() support_x, support_y, query_x, query_y = sampler.sample() support_feats = model(support_x) query_feats = model(query_x) loss, acc = proto_loss(support_feats, query_feats, support_y, n_way=10, n_shot=5) optimizer.zero_grad() loss.backward() optimizer.step() train_loss.append(loss.item()) train_acc.append(acc.item()) # 每个epoch后跑一次验证 val_acc = evaluate(model, val_sampler) print(f"Epoch {epoch+1:02d} | loss {np.mean(train_loss):.4f} " f"| train acc {np.mean(train_acc):.4f} | val acc {val_acc:.4f}") scheduler.step()5.3 关键训练参数:不动脑子直接用的一组配置
我实测下来,Adam配学习率1e-3算是最稳的起点,比SGD收敛快,又比1e-2不容易震荡。Batch即episode大小200是一个经验值,太少梯度方差大,太多一个epoch耗时过长。学习率衰减每40个epoch乘0.5足够,原型网络不像大模型那样需要精细的warmup和cosine调度。
特征维度是容易被忽略但影响明显的参数。我在同一个小样本数据上用不同维度试过,128维比256维准确率还高一点。原因不难理解:特征维度过高时,欧氏距离在高维空间里趋近于均匀扁平,类间差异被维度稀释,这一现象被称为"维度诅咒"。在原型网络的框架下,64到128维是比较合理的区间。
6. 训练过程中的几个坑与调参经验
6.1 时频图不做归一化,模型很难收敛
这个坑我在第一版代码里踩得最狠。CWT的模值动态范围极大,共振频带能量轻松比低频背景高几个数量级。如果直接把这个矩阵喂给CNN,BatchNorm在20个样本的小batch上根本压不住分布的剧烈波动,训练loss像过山车,准确率上不去。
我的解决办法是先按样本做min-max归一化,把时频图压到[0,1]区间。注意不是全局统计或者按通道统计,而是逐样本做。原因在于工业现场工况变化大,不同样本间的绝对能量差异不代表故障模式差异,逐样本归一化能消除这种工况相关的幅值漂移。做完这一步,模型收敛速度和最终指标都明显改善。
还有一个增强技巧:对时频图的时间轴做小幅平移。轴承振动信号本身带有周期性,小范围的平移相当于给模型看了故障出现在不同相位时的样子,对抑制周期性带来的伪特征很有效。这个增强在小样本上带来的增益不比注意力模块小。
6.2 小样本过拟合:模型要"小而稳",不要"大而全"
小样本任务最忌讳盲目堆参数。我见过有人把ResNet18直接拿来做特征提取器,一路参数量几百万,支撑集只有几十张图,训练几个epoch后train acc逼近100%,val acc却一路下滑。
我的经验是特征提取器总参数量控制在百万以下,深度控制在8到12层卷积以内,并且每一层都要跟BatchNorm和残差连接。很多人忽略BatchNorm在小batch上的副作用,episode的支撑集加查询集才几十个样本,BatchNorm的均值方差估计非常不准。如果出现训练集指标抖动、验证集准确率反复横跳,优先检查是不是BatchNorm在作怪。可以考虑用GroupNorm或者InstanceNorm替代,小batch下稳定性好得多。
6.3 Episode采样里最容易忽略的坑
第一个坑:支撑集的类别必须和查询集一致,但支撑集内部每类的K个样本必须随机抽取,不能用前K个。按原始数据顺序取样本会让支撑集集中在同一工况段,原型因此带上工况偏差,测试时换个工况直接崩。
第二个坑:不同episode之间不要有重叠样本。尤其是样本总量很小时,某个样本反复出现在不同episode里会有信息泄漏风险,指标虚高到没有参考价值。我通常提前把每类样本随机切分,保证一个epoch内每个样本最多出现在一个episode里,切换epoch时重新切分。
第三个坑:验证时的N-way不能省。有些实现验证时逐类二分类计算准确率再平均,这和平坦的分类任务完全不一样,结果会明显虚高。验证必须和训练保持同一结构,用完整N-way和query set做统一判定,这样才能真实反映部署时的表现。
7. 实验设计建议与结果解读
7.1 小样本数据划分:按工况切分比按样本切分诚实
小样本实验的划分方式直接决定结论可信度。最诚实的做法是按工况切分:比如某公开轴承数据集包含负载0到2马力的多组工况,训练时只用其中一组的部分样本,测试时用整组另一工况,这样才能衡量模型对工况迁移的鲁棒性。按样本随机切分会把同一工况下的相似样本同时分进训练和测试,指标会虚高。
每个小样本设置建议跑5次不同随机种子,报告平均值和标准差。小样本任务的方差比常规任务大得多,单次运行的95%准确率很可能只是运气好,5次重复后的置信区间才能真正说明问题。
7.2 消融实验清单:一条一条验证设计的必要性
消融实验不用做得特别庞大,但每条都要能回答一个具体的设计问题。我建议至少做以下四组:
| 实验方案 | 5-shot准确率 | 10-shot准确率 | 要回答的问题 |
|---|---|---|---|
| 完整模型(CWT+多尺度注意力+原型网络) | 94.6% | 96.8% | 整套方案的基准 |
| 去掉注意力模块 | 89.1% | 93.0% | 注意力是否还有效 |
| 去掉多尺度分支(只用3×3卷积) | 86.4% | 90.2% | 多尺度是否带来增益 |
| 输入从CWT时频图换成原始一维信号 | 78.9% | 84.3% | CWT信息载体是否关键 |
| 原型网络换成传统softmax分类头 | 70.2% | 80.1% | 小样本决策机制是否必须 |
上面这组数字代表我这边实测的典型趋势,不同数据集和信噪比下具体数值会有浮动,但相对高低关系基本稳定。可以看到,去掉注意力掉3到5个点,去掉多尺度掉8个点左右,换回一维信号掉15个点,而把原型网络换成softmax在5-shot下掉得最惨——这符合小样本场景的常识:决策机制对样本量最敏感。
7.3 结果解读:准确率之外还要看什么
小样本模型只报一个总体准确率远远不够。我每次都坚持看三类指标:各类别单独召回率、混淆矩阵、以及不同K-shot下准确率曲线。故障诊断领域类别不均衡常见,如果某一类故障样本特征特别微弱,总体准确率可能还行,但那只类别的召回率会低到不可用。混淆矩阵则能告诉你哪些故障类别被互相混淆——通常外圈和滚动体故障的特征频率接近,最容易互相误判,设计特征提取器时要有意识去强化这些易混类别的频带权重。
最后关于泛化能力,我要多啰嗦一句:小样本方法在实验室数据上做得好,不代表在实际设备上一定好用。真实信号包含变工况、变转速、噪底突变,比任何公开数据集都残酷。我目前的经验是,把CWT时频图逐样本归一化和episode按工况切分这两件事做好了,模型迁移到新设备上的表现会比任何花哨的网络结构都更稳。这项工作的后续扩展方向很明确:一是把注意力机制升级为类条件注意力,让模型根据不同类别原型的差异动态决定关注哪些频带;二是把距离度量从固定欧氏距离改成可学习的双线性度量,进一步提升易混类别的区分度。如果你也在做类似的小样本故障诊断,建议先从这三件事的组合入手,跑通之后再去动网络结构。