用WGAN-GP生成故障轴承振动信号:解决小样本问题的可行路径
简介面向故障诊断与深度学习研究者的GAN-1D项目资源应用WGANWasserstein生成对抗网络生成故障轴承振动信号主要用于解决实际工程中轴承故障样本稀缺、类别不平衡等问题为数据增强和智能诊断提供实验基础。压缩包内共19个文件包含Python训练与测试脚本train.py、test.py、预处理后的.mat数据集、可用于继续训练的TensorFlow检查点文件.ckpt以及9张不同状态信号生成效果对比图整体大小约22.03MB结构清晰便于快速复现。资源提供完整的WGAN训练流程支持自定义学习率、训练轮数和采样频率并内置9类信号数据可选读者可基于此调整网络结构或损失函数适配自身故障诊断场景。目前已有3387人学习下载适合具备Python与TensorFlow基础、希望运用生成对抗网络扩充故障样本的研究者参考使用。1. 用 WGAN 生成故障轴承振动信号小样本问题的一条可行出路做故障诊断的工程师基本都撞过同一堵墙正常工况的振动数据要多少有多少但故障状态的数据——尤其是早期故障、罕见工况下的样本——少得可怜。拿这些样本去训练 CNN 或 LSTM模型要么过拟合要么对故障特征根本不敏感。与其到处找数据或者靠人工造样本不如用生成式模型直接补数据。这份资源做的就是这件事用 WGANWasserstein GAN的改进版本 WGAN-GP在一维振动信号上做生成把故障轴承的时域波形“造”出来。它解决的是故障诊断里最实际的痛点——样本不平衡适合正在做故障诊断、设备健康管理、PHM 相关课题的工程师和学生前提是你对 Python 和 PyTorch 有基本操作能力不需要懂太深的数学也能跟着把流程跑通。2. 为什么是 WGAN 而不是原始 GAN一维信号生成的选型逻辑2.1 原始 GAN 在振动信号上的两个典型失效场景先看原始 GAN 的结构生成器把随机噪声映射成假样本判别器判断输入是真实振动信号还是生成信号两者对抗训练。理论上这个机制能收敛但实际跑一维振动信号时问题非常具体。第一是JS 散度在低维流形上的不连续性。振动信号本质上落在高维空间里的低维流形上真实分布和生成分布的支撑集几乎不重叠JS 散度在这些位置是常数梯度传不回去生成器收不到有效反馈。第二是判别器太容易饱和。信号数据不像图像那样有丰富的纹理梯度判别器可能在几百步内就达到接近零的分类误差之后梯度消失生成器永远学不到故障冲击的周期性特征。这两个问题在图像生成里还能靠调参硬扛但在振动信号这种一维时序数据上几乎是致命的。我见过不少项目用原始 GAN 生成振动波形训练到后面生成器输出的全是正弦波甚至连故障频率都淹没在噪声里这就是典型的 mode collapse 和梯度消失叠加的结果。2.2 Wasserstein 距离改了什么从“判断真假”到“度量差距”WGAN 的核心改动是把判别器换成Critic评论家输出不再是一个“真/假”概率而是一个实数用来估计真实分布和生成分布之间的 Wasserstein 距离。这个距离的好处是即使两个分布的支撑集不重叠它依然能给出一个平滑且有意义的梯度信号。换句话说原始 GAN 告诉你“你错了”WGAN 告诉你“你离目标还有多远、往哪个方向走”。这个改动在代码层面非常小。原始 GAN 的判别器最后一层是 SigmoidWGAN 的 Critic 去掉 Sigmoid 直接输出一个标量。训练时Critic 的损失是真实样本输出均值减去生成样本输出均值生成器的损失是生成样本输出均值的负值。具体到 PyTorch 里改动只有几行但收敛行为的差异是质变。2.3 WGAN-GP 的梯度惩罚把权重裁剪换掉WGAN 原始论文用权重裁剪weight clipping来保证 Lipschitz 约束把 Critic 的权重限制在一个小范围内。这个做法有个副作用权重全被推到边界值模型容量损失严重训练后期容易出现振荡。WGAN-GPGradient Penalty换了一种思路——在真实分布和生成分布之间的采样点上强制梯度范数接近 1用软约束替代硬裁剪。梯度惩罚的实现有几个细节值得注意惩罚项要加在 Critic 的损失上采样点是在真实样本和生成样本连线上随机插值得到的惩罚系数一般取 10训练时 Critic 每更新 5 次生成器更新 1 次。这套配置是实践中验证过的稳定组合后续章节的代码也按这个来。提示如果你用的是 WGAN 原始版本务必检查是否做了权重裁剪如果用了 WGAN-GP检查 Critic 权重是否被错误地裁剪了。两者混用会直接破坏梯度惩罚的意义。3. 数据准备与预处理把轴承振动信号转换成模型能吃的格式3.1 数据来源与故障类型的组织这份资源没有自带原始数据需要你自己准备轴承振动信号。常见来源是公开的轴承故障数据集或者实验室台架采集的数据。不管哪种来源都需要按故障类型组织目录结构。我一般这样组织data/ normal/ normal_0001.csv normal_0002.csv inner_race/ inner_race_0001.csv outer_race/ ball/每个 CSV 文件是一段连续的振动信号采样频率可能不同但每段信号内部是等间隔采样的。如果数据结构是像某些平台那样一个文件里包含多列比如同时包含驱动端和风扇端的数据需要先做通道提取只保留你要用的那一列。故障类型建议至少准备两种内圈故障和外圈故障。这两种故障的振动特征差异明显生成模型学到的特征更有区分度。球故障滚动体故障由于特征较弱生成难度大不建议作为第一个试验对象。3.2 滑窗采样固定长度切分与重叠率设定原始信号是长序列不能直接喂给模型。需要按固定长度做滑窗切分。窗口长度的选择直接影响生成效果。常见做法是取一个包含多圈旋转信号的窗口长度——比如转速 1800 转/分钟即 30 转/秒采样率 12kHz 时每转对应 400 个采样点取 1024 个点大约覆盖 2.5 个旋转周期足以包含 2 到 3 次故障冲击。滑窗切分代码如下import numpy as np import pandas as pd def sliding_window(data, window_size1024, step512): 将一维信号切分为固定长度窗口。 参数: data: 一维 numpy 数组或 pandas Series window_size: 窗口长度建议 1024 或 2048 step: 滑窗步长建议 window_size // 2 返回: 二维数组形状为 (样本数, window_size) if isinstance(data, pd.Series): data data.values samples [] for start in range(0, len(data) - window_size 1, step): samples.append(data[start:start window_size]) return np.array(samples)这里有两处参数值得注意。window_size决定了生成信号的时间跨度太小比如 256时故障冲击可能只有一个脉冲模型学到的“故障特征”对位置极其敏感太大比如 4096时模型参数量需要增加收敛变慢。step控制重叠率50% 重叠step window_size // 2是数据增强的常见做法能有效扩大样本量。切分完成后对每个窗口做归一化。振动信号的幅值范围不固定直接用原始幅值训练会让生成器输出范围飘移。最简单的做法是每个窗口独立做零均值单位方差归一化def normalize_windows(samples, eps1e-8): 窗口级归一化每个窗口减均值除标准差。 mean samples.mean(axis1, keepdimsTrue) std samples.std(axis1, keepdimsTrue) return (samples - mean) / (std eps)为什么要窗口级归一化而不是全局归一化因为振动信号在不同时段的有效值差异很大——机器启动阶段和稳态阶段的幅值能差好几倍。全局归一化会把启动阶段的信号压成一条直线等于把有用信息抹掉了。3.3 训练集与测试集的划分不要源自私心划分数据集时有个容易踩的坑同一个长信号切出来的窗口之间存在高度相关性因为滑窗有重叠如果这些窗口同时出现在训练集和验证集中模型评估的“泛化能力”是虚高的。我一般这样处理先按连续段切分把一整段信号分为训练段和测试段然后在训练段内做滑窗测试段最后作为独立的真实样本评估集。# 假设 segments 是一个列表每个元素是一段完整信号 train_segments segments[:int(len(segments) * 0.8)] test_segments segments[int(len(segments) * 0.8):] train_windows np.concatenate( [sliding_window(seg, 1024, 512) for seg in train_segments] )这样划分后训练窗口和测试窗口之间没有重叠覆盖的原始时间点评估结果才能反映真实生成质量。4. 搭建 WGAN-1D 模型从生成器到训练循环的完整实现4.1 生成器结构转置卷积还是一维上采样生成器的输入是低维噪声常见是 100 维输出是 1024 点的一维信号。从 100 维映射到 1024 维有两种常见方案全连接层直接展开或者用转置卷积ConvTranspose1d逐步上采样。全连接展开的优点是结构简单、前期收敛快缺点是对局部特征的建模能力弱生成信号容易出现高频毛刺转置卷积的优点是能学习到多尺度的局部结构但需要精心设计各层通道数。我建议用转置卷积方案因为它对振动信号这种包含周期性脉冲的模式更友好。一个参考结构如下import torch import torch.nn as nn class Generator(nn.Module): WGAN-GP 生成器从 128 维噪声生成 1024 点一维信号。 def __init__(self, noise_dim128, base_channels64): super(Generator, self).__init__() self.noise_dim noise_dim # 先把噪声投影到 128 * 8 的向量再用转置卷积逐步上采样 self.fc nn.Linear(noise_dim, base_channels * 8 * 8, biasFalse) self.deconv nn.Sequential( # 输入形状: (batch, 64*8, 8) nn.ConvTranspose1d(base_channels * 8, base_channels * 4, kernel_size16, stride2, padding7, output_padding1), nn.BatchNorm1d(base_channels * 4), nn.ReLU(True), # (batch, 64*4, 16) nn.ConvTranspose1d(base_channels * 4, base_channels * 2, kernel_size16, stride2, padding7, output_padding1), nn.BatchNorm1d(base_channels * 2), nn.ReLU(True), # (batch, 64*2, 32) nn.ConvTranspose1d(base_channels * 2, base_channels, kernel_size16, stride2, padding7, output_padding1), nn.BatchNorm1d(base_channels), nn.ReLU(True), # (batch, 64, 64) nn.ConvTranspose1d(base_channels, 1, kernel_size16, stride2, padding7, output_padding1), # (batch, 1, 128) nn.Tanh() ) def forward(self, z): x self.fc(z) x x.view(-1, self.noise_dim if False else 512, 8) # 注意调整维度 x self.deconv(x) return x.squeeze(1)这里有个小坑fc的输出维度和deconv的输入形状要跟通道数严格对齐。base_channels64时fc输出 512 个特征reshape 成 64*8 个通道、每通道 8 个点。如果你改大base_channels必须同步调整fc的out_features和reshape的维度。这地方不报错但训练会乱套是典型的隐蔽问题。最后一层用Tanh做激活把输出限制在 [-1, 1]对应第 3 章归一化后的数据范围。如果数据归一化时用的是别的范围比如 [0, 1]这里的激活函数也要改。4.2 Critic判别器结构带梯度惩罚的容量控制Critic 的结构和生成器镜像对称输入 1024 点信号输出一个标量。用一维卷积逐步下采样最后做全局平均池化再接全连接层输出标量。关键是不要用 BatchNorm——梯度惩罚要求在样本维度上独立计算梯度范数BatchNorm 会引入样本间的统计依赖破坏这个约束。我一般用 InstanceNorm 或者干脆不用归一化。class Critic(nn.Module): WGAN-GP 的 Critic输出一个实数评估分布距离。 def __init__(self, base_channels64): super(Critic, self).__init__() self.conv nn.Sequential( nn.Conv1d(1, base_channels, kernel_size8, stride2, padding3), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv1d(base_channels, base_channels * 2, kernel_size8, stride2, padding3), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv1d(base_channels * 2, base_channels * 4, kernel_size8, stride2, padding3), nn.LeakyReLU(0.2, inplaceTrue), nn.Conv1d(base_channels * 4, base_channels * 8, kernel_size8, stride2, padding3), nn.LeakyReLU(0.2, inplaceTrue), ) self.fc nn.Linear(base_channels * 8 * 8, 1) def forward(self, x): x x.unsqueeze(1) # (batch, 1, 1024) x self.conv(x) x x.view(x.size(0), -1) return self.fc(x)x.unsqueeze(1)是把输入从(batch, 1024)变成(batch, 1, 1024)适配 Conv1d 的通道维。LeakyReLU 的负斜率 0.2 是经验值太小会让负区间的梯度几乎消失太大则逼近线性我一般固定在 0.2 不动。4.3 训练循环梯度惩罚的计算与两步交替WGAN-GP 的训练循环有三件事和普通 GAN 不同Critic 的梯度要裁剪等等是梯度惩罚不是权重裁剪、Critic 要训练好多次生成器才训一次、损失函数不加 log。完整训练循环如下def train_wgan_gp(generator, critic, train_loader, epochs200, lr1e-4, lambda_gp10, n_critic5): WGAN-GP 训练主循环。 参数: generator: 生成器模型 critic: Critic 模型 train_loader: 真实信号的数据加载器 epochs: 训练轮数 lr: 学习率WGAN-GP 通常用 1e-4比普通 GAN 低 lambda_gp: 梯度惩罚系数固定为 10 是稳定经验值 n_critic: Critic 每训练几次生成器训练一次 g_optim torch.optim.Adam(generator.parameters(), lrlr, betas(0.5, 0.9)) c_optim torch.optim.Adam(critic.parameters(), lrlr, betas(0.5, 0.9)) for epoch in range(epochs): for i, real_samples in enumerate(train_loader): batch_size real_samples.size(0) real_samples real_samples.float() # 训练 Critic n_critic 次 for _ in range(n_critic): c_optim.zero_grad() z torch.randn(batch_size, generator.noise_dim) fake_samples generator(z).detach() # 生成器梯度不需要回传 c_loss_real critic(real_samples).mean() c_loss_fake critic(fake_samples).mean() c_loss c_loss_fake - c_loss_real # 梯度惩罚 alpha torch.rand(batch_size, 1) interp (alpha * real_samples (1 - alpha) * fake_samples).requires_grad_(True) interp_outs critic(interp) grads torch.autograd.grad( outputsinterp_outs, inputsinterp, grad_outputstorch.ones_like(interp_outs), create_graphTrue, retain_graphTrue )[0] grad_norm grads.view(batch_size, -1).norm(2, dim1) gp ((grad_norm - 1) ** 2).mean() c_loss_total c_loss lambda_gp * gp c_loss_total.backward() c_optim.step() # 训练生成器 1 次 g_optim.zero_grad() z torch.randn(batch_size, generator.noise_dim) fake_samples generator(z) g_loss -critic(fake_samples).mean() g_loss.backward() g_optim.step() if (epoch 1) % 10 0: print(fEpoch {epoch1}/{epochs} | C Loss: {c_loss.item():.4f} | G Loss: {g_loss.item():.4f})逐块拆开说明。c_loss c_loss_fake - c_loss_real对应的是 Wasserstein 距离的估计值Critic 想让真实样本得分更高、生成样本得分更低所以这个损失越小越好。fake_samples.detach()是关键——Critic 反向传播时不能让梯度流到生成器否则生成器的参数会被 Critic 的更新连带改动。梯度惩罚部分是最容易写错的位置。alpha torch.rand(batch_size, 1)生成每个样本独立的插值系数interp是真实样本和生成样本的随机插值点——这些点理论上在真实分布与生成分布的“中间地带”对它们施加梯度范数约束能保证整个空间的 Lipschitz 连续。torch.autograd.grad里务必加上create_graphTrue因为后面c_loss_total.backward()需要对这个梯度再求一次梯度即梯度惩罚本身要参与 Critic 的反向传播。关于超参lr1e-4是 WGAN-GP 论文和实践中验证过的稳定值比 GAN 常用的 1e-3 低一个量级因为 Wasserstein 损失的梯度量纲和交叉熵不同学习率太大会直接震荡。lambda_gp10是原论文推荐的我测试过 5 和 20效果不如 10 稳定。n_critic5保证 Critic 有足够的训练量来估计分布距离如果你发现 Critic 的损失掉得很快而生成器跟不上可以改成 10。5. 训练遇坑排查四个高频问题的现象与解决路径5.1 Critic 损失持续下降但生成信号全是噪声这是最典型的现象训练日志显示 Critic 损失一路走低但每 50 轮保存一次生成的信号听感或频谱上全是白噪声没有明显的周期性脉冲。原因是生成器进入了退化状态——生成器输出的方差极小集中在某个点附近Critic 很容易区分真假所以损失一直在降但生成器没有学到任何故障特征。排查路径先看生成信号的均值方差。如果标准差接近 0.01 甚至更低基本是生成器退化。解决方法是降低学习率到 5e-5同时把n_critic增加到 10让 Critic 更快地跟上生成器的变化如果还不行检查生成器最后一层的Tanh输出范围和输入信号的实际范围是否一致——如果输入信号归一化后范围是[-1.5, 1.5]而你用了Tanh模型永远学不到超出范围的部分。5.2 生成信号存在周期性但故障频率不对生成信号看得出有明显的周期性波动但频谱峰值对应的频率和真实故障特征频率对不上。最常见的原因是滑窗切分时窗口长度小于一个旋转周期。比如转速 3000 转/分钟50 转/秒、采样率 12kHz 时一个周期 240 采样点如果窗口只有 256 点模型看到的波形里故障冲击的出现位置每次都不同它学到的“周期性”是窗口内的伪周期。解决方法是把窗口长度扩大到至少包含 4 个旋转周期本例即 960 点以上或者在做滑窗之前先用转速信号做角域重采样把时域信号变成角度域信号。如果不想改窗口长度降低转速波动的影响也可以——但这是治标不治本故障特征频率还是乱的。5.3 梯度惩罚值异常大或变成 NaN训练开始时gp项的数值在几十到几百之间是正常的因为初始模型输出不满足 Lipschitz 约束但正常应该在 100 轮内逐渐下降到 10 以下。如果gp一直大于 50 或者直接变成 NaN大概率是Critic 的输入数据范围问题。梯度惩罚要求插值点的梯度范数接近 1但如果你输入信号的幅值范围过大比如归一化前直接喂原始信号插值点上的梯度会被放大很多倍惩罚项剧烈震荡。另外检查torch.autograd.grad里的grad_outputs是否设置正确它必须和interp_outs形状一致。torch.ones_like(interp_outs)表示对输出向量的每个分量都计算梯度如果你忘了这个参数Pytorch 会报错但不会提示原因指向的方向很奇怪。5.4 不同故障类别的生成信号无法区分模型对正常信号和故障信号都能生成但生成的内圈故障信号和外圈故障信号放在频谱图上几乎一样。这不是模型能力问题是类别数据量不均衡导致的。假设内圈故障样本有 1000 个窗口外圈故障只有 200 个Critic 会优先拟合样本量大的那个分布生成器自然也更偏向生成内圈故障特征。解决思路有三种一是对外圈故障做更多滑窗采样减小step二是对少样本类别做信号增广加噪声、时间偏移、幅值扰动三是先单独训练两个类别的模型再混合训练做迁移。我一般优先用第一种因为实现成本最低且问题往往就出在窗口数量差一个数量级。注意如果你同时生成多个故障类别一定要在保存模型和评估时按类别分别保存生成样本不要混在一个文件夹里。否则后续做故障分类模型训练时类别标签对不上等于白做。6. 进阶技巧用频域指标量化生成质量调整超参数的落点训练完模型后光看损失曲线不能说明生成质量。我习惯用三步验证时域波形目视、频域特征对比、时频图观察冲击特征。第一步是频域对比。把真实故障信号和生成信号分别做 FFT得到频谱后比对主要的故障特征频率内圈故障特征频率、外圈故障特征频率、转频及其谐波处的幅值比例。用代码实现很简单import numpy as np import torch def spectral_similarity(real, fake, fs12000): 计算真实信号与生成信号的频谱相似度基于功率谱的余弦相似度。 real_p np.abs(np.fft.rfft(real))**2 fake_p np.abs(np.fft.rfft(fake))**2 real_p real_p / (np.linalg.norm(real_p) 1e-8) fake_p fake_p / (np.linalg.norm(fake_p) 1e-8) return float(np.dot(real_p, fake_p))这个数值通常在 0.7 到 0.9 之间说明频谱能量分布较为接近低于 0.5 说明生成信号失去了大部分故障特征。注意这里用的功率谱是对整段信号做的如果故障冲击是瞬态的比如内圈早期故障还应该配合短时傅里叶变换看时频图看冲击出现的周期是否对齐。第二步是诊断应用验证——把生成样本混入真实样本中训练一个简单的故障分类器。比如用 50% 真实 50% 生成样本训练一个两层 CNN然后单独用真实样本做测试。如果分类准确率比纯真实样本训练时下降不超过 3%说明生成样本基本可用。这套流程能同时验证生成样本的信息量和多样性。最后是超参数微调的落点。lr的敏感度最高建议先固定 1e-4不收敛就降 1e-5收敛太慢就升到 5e-4 但警惕振荡。n_critic在数据量偏小时可以提高——样本少、Critic 容易过拟合时n_critic10能稳定不少。window_size决定了模型看到的信号跨度通常不要在训练中频繁调整但可以从 1024 换到 2048 观察频谱覆盖范围是否更完整。还有一个容易忽略的参数是batch_size振动信号数据往往较大但 batch 太大比如 256会导致梯度惩罚计算开销暴涨训练速度显著变慢我一般取 64。整套流程跑完之后我习惯做一件事每 50 轮保存一次生成器的权重和噪声向量保证固定种子训练结束后把每个 check point 的生成样本做成一个频谱动画观察模型从噪声到故障特征的收敛过程。这个方法帮我定位过好几次 mode collapse 的发生节点。从那以后我每次训练生成模型都强制走一遍这套验证流程——先频谱相似度再分类器评估最后看收敛动画三关过了才敢把数据交给下游模型。希望帮到你。本文还有配套的精品资源点击获取