☰
WGAN-GP实现轴承故障振动信号生成与数据增强实战
2026/10/10 14:30:36 网站建设 项目流程

简介:面向轴承故障诊断中的数据稀缺与类别不平衡问题,这个基于WGAN的生成对抗网络工程可对原始振动信号建模,输出合成故障样本,适合深度学习、故障诊断与数据增强方向的工程师、研究生复现及二次开发。压缩包共19个文件、体积约22MB,包含Python训练与测试脚本、MAT格式数据集、模型权重文件以及9张信号可视化结果图,目录简洁,运行入口清晰。项目基于TensorFlow GPU实现,训练脚本支持自定义学习率、迭代轮数和采样周期,内置9种信号可供选择,并附带checkpoint与数据增强文件,既降低了环境配置门槛,也便于直接加载模型进行推理或继续训练。目前已有3386人学习浏览,整体资料结构完整,既能快速生成故障振动信号,也为科研对比、论文实验提供了可复现的代码底座。

1. 用WGAN补足故障轴承振动样本:为什么常规GAN在这件事上翻车

故障轴承的振动信号是典型的非平稳、非高斯一维时序数据,诊断模型想要训练得扎实,缺的不是算法而是样本。实际产线里正常样本一抓一大把,故障样本却少得可怜,尤其是早期损伤、复合故障这类工况,能采集到的完整样本可能只有几十条,根本喂不饱深度学习模型。我最初试过用普通DCGAN直接生成一维振动信号,结果训练不到两三百个epoch,生成器就把所有样本都收敛到同一条波形上去,时域波形看起来像模像样,频域一分析全是噪声,这种样本拿去做诊断增强只会把模型带偏。后来换成WGAN(Wasserstein GAN)做一维故障信号生成,训练稳定性明显改善,生成样本的频谱包络和真实故障特征能对齐。这份资源的核心就是把WGAN完整落到轴承故障振动信号生成上,从数据切分、生成器判别器搭建到训练调参、效果验证都给了可运行的Python代码,适合正在做故障诊断数据增强、样本不平衡处理的从业者直接参考。

2. 从真实振动信号到训练样本:一维数据怎么喂给WGAN

2.1 为什么不用二维图而直接用一维原始信号

很多做故障诊断的同学第一反应是把振动信号做小波变换或短时傅里叶变成时频图,再用DCGAN去生成图片。这个路线不是不行,但有两个实际麻烦:一是时频图的分辨率、窗函数、变换参数都会直接影响生成样本质量,引入额外人为因素;二是生成出来的时频图还需要做逆变换回到一维信号,这个环节的误差会让后续诊断模型无所适从。GAN-1D这个项目直接用一维原始信号训练,生成器输出的就是和输入同长度的采样序列,省掉了来回变换的环节,诊断模型拿到的增强样本和真实样本在数据结构上完全一致。

一维信号建模有一个关键点需要先想清楚:振动信号不像图像有强烈的空间局部性,但相邻采样点之间确实存在短时相关性。处理这种关系的常见做法是堆叠一维卷积核,感受野逐层扩大,让生成器先学局部振荡模式,再逐步合成整体波形。项目里生成器的上采样策略也是顺着这条思路设计的,先在小尺度上生成细节,再向上采样到目标长度。

2.2 数据切片、归一化与批量读取的标准化流程

故障轴承振动信号原始文件通常是长序列连续采样,训练WGAN前必须先切成固定长度片段。我一般用1024个采样点作为单条样本长度,这个长度在8kHz到12kHz采样率下能覆盖几个完整的轴承转频周期,同时包含足够多的冲击特征。切分窗口设置50%重叠能扩充样本量,但要注意重叠率太高会让训练集和验证集之间的样本高度相关,我实际用下来30%到50%比较稳妥。

归一化是这里最容易踩坑的环节。WGAN的生成器输出层用的是tanh激活函数,输出范围是[-1, 1],因此输入的真实信号也必须归一化到同一个范围。有人直接用MinMaxScaler把数据压到[0, 1],训练时loss确实能降,但生成样本的幅值分布和真实信号对不上,诊断模型拿去用效果很差。正确做法是:

import numpy as np def normalize_signal(signal, eps=1e-8): """将振动信号归一化到 [-1, 1],保持相对幅值信息""" min_val = np.min(signal) max_val = np.max(signal) if max_val - min_val < eps: return np.zeros_like(signal) # 2 * (x - min) / (max - min) - 1 映射到 [-1, 1] normalized = 2.0 * (signal - min_val) / (max_val - min_val) - 1.0 return normalized def sliding_window_slice(signal, window_size=1024, stride=512): """滑窗切分,返回二维数组 [n_samples, window_size]""" samples = [] n = len(signal) for start in range(0, n - window_size + 1, stride): seg = signal[start:start + window_size] samples.append(normalize_signal(seg)) return np.array(samples, dtype=np.float32)

这里的stride参数控制重叠率,stride=512在1024窗口下就是50%重叠。归一化是按每一条切片独立做的,不是对整个长序列做全局MinMax,因为WGAN训练时逐batch输入,如果全局归一化,不同工况下幅值差异大的信号会被压平,生成器学不到真实的冲击幅值分布。

批量读取用PyTorch的TensorDataset和DataLoader就够用,shuffle=True是必须的,否则每个epoch样本顺序固定,生成器容易记住批次顺序产生周期性伪影。batch size我习惯设64,显存够的话可以加大到128,但注意batch size和critic更新次数之间存在联动关系,后面第4章会细说。

3. WGAN-GP模型搭建:一维信号生成器和判别器的关键设计

3.1 生成器:噪声向量如何一步步变成1024点振动序列

WGAN的生成器本质是一个上采样网络,输入是服从标准正态分布的隐向量,输出是1024维的振动信号序列。项目采用的是一维转置卷积(ConvTranspose1d)堆叠结构,每一层将序列长度翻倍,通道数逐层递减,直到最后一层输出单通道信号。结构上从100维噪声开始,通过线性层先扩成256通道的短序列,再经过3层转置卷积逐步上采样到1024点。

每一层转置卷积之间需要激活函数,这里用的不是ReLU而是LeakyReLU,负斜率设0.2。原因是ReLU会把负激活全部截断为0,导致生成信号出现削波,对振动信号的冲击特征影响很大。LeakyReLU保留少量负值信息,生成的波形更平滑、冲击形态更自然。

import torch import torch.nn as nn class Generator1D(nn.Module): """一维振动信号生成器:100维噪声 -> 1024采样点信号""" def __init__(self, latent_dim=100, signal_len=1024, channels=64): super().__init__() self.latent_dim = latent_dim # 先从 [batch, 100] 映射到 [batch, 256, 4] 的初始序列 self.init_fc = nn.Linear(latent_dim, channels * 4 * 4) # 三层转置卷积,长度 4 -> 16 -> 64 -> 256 -> 1024 更稳定 self.deconv1 = nn.ConvTranspose1d(channels, channels // 2, kernel_size=8, stride=4, padding=2) self.deconv2 = nn.ConvTranspose1d(channels // 2, channels // 4, kernel_size=8, stride=4, padding=2) self.deconv3 = nn.ConvTranspose1d(channels // 4, channels // 8, kernel_size=8, stride=4, padding=2) self.deconv4 = nn.ConvTranspose1d(channels // 8, 1, kernel_size=8, stride=4, padding=2) self.leaky = nn.LeakyReLU(0.2) self.tanh = nn.Tanh() self._init_weights() def _init_weights(self): # 转置卷积的零初始化技巧,避免初始阶段产生极端幅值 for module in [self.deconv1, self.deconv2, self.deconv3, self.deconv4]: nn.init.normal_(module.weight, mean=0.0, std=0.02) nn.init.zeros_(module.bias) def forward(self, z): x = self.init_fc(z).view(z.size(0), -1, 4) x = self.leaky(x) x = self.leaky(self.deconv1(x)) x = self.leaky(self.deconv2(x)) x = self.leaky(self.deconv3(x)) x = self.tanh(self.deconv4(x)) return x.squeeze(1) # [batch, 1024]

这里有个值得注意的细节是_init_weights中标准差取0.02而不是默认的1.0。常见做法是控制初始权重范围,标准差在0.02附近,避免生成器在训练早期输出幅值特别大的信号,判别器直接饱和导致梯度消失。最后一层用固定步长4、kernel size 8的组合,每次上采样4倍,四层是4到1024点的完整映射。

如果不想用ConvTranspose1d,也可以用Upsample + Conv1d的组合,上采样倍率更好控制,棋盘伪影更轻,但训练速度稍慢。项目里转置卷积方案是经过验证的,生成1024点序列不需要追求过大kernel,8以内足够了。

3.2 判别器(Critic):一维卷积结构输出的是分数而不是概率

WGAN的关键变化在判别器上。普通GAN的判别器输出一个sigmoid概率值,表示“真或假”的概率;WGAN的Critic输出一个没有经过sigmoid的实数值,表示输入样本的Wasserstein距离估计。这个设计直接解决了JS散度导致的梯度消失问题。项目里Critic用4层一维卷积叠加,每层使用步长2的普通卷积压缩序列长度,通道数逐层增加,最后通过线性层输出一个标量。

每层卷积后面跟的还是LeakyReLU(0.2),和生成器保持一致。这里强烈建议不要用BatchNorm,尤其是Critic中更不能用。BatchNorm会破坏WGAN-GP对每个样本独立计算梯度惩罚的要求,批量统计量会引入样本间的依赖,导致惩罚项失效。实际训练时会发现loss波动明显加剧,生成样本多样性下降。

class Critic1D(nn.Module): """WGAN判别器:输入1024点信号,输出Wasserstein距离估计分数""" def __init__(self, signal_len=1024, channels=64): super().__init__() self.conv1 = nn.Conv1d(1, channels, kernel_size=8, stride=2, padding=3) self.conv2 = nn.Conv1d(channels, channels * 2, kernel_size=8, stride=2, padding=3) self.conv3 = nn.Conv1d(channels * 2, channels * 4, kernel_size=8, stride=2, padding=3) self.conv4 = nn.Conv1d(channels * 4, channels * 8, kernel_size=8, stride=2, padding=3) self.leaky = nn.LeakyReLU(0.2) # 输入经过4层stride=2卷积,序列长度 1024 -> 512 -> 256 -> 128 -> 64 self.fc = nn.Linear(64 * channels * 8, 1) def forward(self, x): # x: [batch, 1024],需增加通道维度 -> [batch, 1, 1024] x = x.unsqueeze(1) x = self.leaky(self.conv1(x)) x = self.leaky(self.conv2(x)) x = self.leaky(self.conv3(x)) x = self.leaky(self.conv4(x)) x = x.view(x.size(0), -1) return self.fc(x) # 输出标量,无sigmoid

Critic不需要输出层归一化到特定范围,它输出的分数绝对值本身没有直接意义,有意义的是真实样本和生成样本之间分数的差值。训练的目标就是让这个差值尽可能反映两个分布之间的Wasserstein距离。这也是为什么WGAN的Loss曲线看起来不像普通GAN那样收敛到某个稳定值,而是在一定范围内持续波动,这是正常现象。

3.3 梯度惩罚项的实现细节

WGAN-GP是在WGAN基础上加了梯度惩罚项来满足Lipschitz约束,替代了原先的权重裁剪方案。权重裁剪会把参数强行限制在一个小范围内,容易导致模型表达能力下降,而且只能粗略近似约束效果。梯度惩罚的思路是对真实样本和生成样本之间的插值点求梯度,约束梯度的范数尽可能接近1。

def compute_gp(critic, real_data, fake_data, lambda_gp=10.0): """WGAN-GP梯度惩罚项:在真实与生成之间随机插值""" batch_size = real_data.size(0) # 随机采样插值权重 alpha,每个样本独立 alpha = torch.rand(batch_size, 1, device=real_data.device) alpha = alpha.expand_as(real_data) # 插值样本介于真实和生成之间 interpolated = alpha * real_data + (1.0 - alpha) * fake_data interpolated.requires_grad_(True) # Critic对插值样本求梯度 critic_interpolated = critic(interpolated) grads = torch.autograd.grad( outputs=critic_interpolated, inputs=interpolated, grad_outputs=torch.ones_like(critic_interpolated), create_graph=True, retain_graph=True )[0] grads = grads.view(batch_size, -1) grad_norm = grads.norm(2, dim=1) penalty = lambda_gp * ((grad_norm - 1.0) ** 2).mean() return penalty

这里lambda_gp=10.0是原论文推荐的默认值,实际项目里我试过1到20之间的不同取值,10确实是个稳定点,小于5时约束不够训练容易波动,大于20时生成样本偏保守、多样性明显下降。代码中create_graph=True和retain_graph=True是必须的,因为梯度惩罚的梯度还要反传到Critic参数上。这个细节如果漏掉,PyTorch会直接报错或用错计算图最终梯度算错。

4. 训练循环与参数调优:n_critic、学习率、损失平衡怎么设

4.1 WGAN-GP的标准训练循环拆解

WGAN的训练和普通GAN最大的区别在于:Critic要训练多次,生成器才训练一次。原论文设定为每训练5次Critic,再训练1次生成器,这个比例通常写作n_critic=5。原因在于Wasserstein距离的估计质量依赖Critic拟合的充分程度,如果Critic没训好就急着更新生成器,梯度方向是错的,训练反而震荡。在振动信号这种数据分布比较复杂的情况下,我先把n_critic设到5跑通,再试3和8对比稳定性。

训练循环里每个batch的具体流程是:先取真实信号样本,再从标准正态分布采样隐向量z,生成器产出fake信号;然后Critic分别对real和fake算分,计算Wasserstein距离的估计值(real分数的均值减fake分数的均值);在real和fake之间插值算梯度惩罚;三项组合成Critic的loss并反向传播更新参数;重复n_critic次后,冻结Critic梯度,用生成器损失反向传播更新生成器参数。

def train_wgan(generator, critic, dataloader, epochs=2000, n_critic=5, lr_g=1e-4, lr_c=1e-4, lambda_gp=10.0, device='cuda', save_interval=50): # 两个优化器独立设置,WGAN普遍使用RMSprop或Adam,学习率必须小 opt_g = torch.optim.Adam(generator.parameters(), lr=lr_g, betas=(0.5, 0.9)) opt_c = torch.optim.Adam(critic.parameters(), lr=lr_c, betas=(0.5, 0.9)) for epoch in range(epochs): for real_signals in dataloader: real_signals = real_signals.to(device) batch_size = real_signals.size(0) # ---------- 更新Critic n_critic次 ---------- critic_loss_epoch = 0.0 for _ in range(n_critic): # 采样隐向量,标准正态分布 z = torch.randn(batch_size, generator.latent_dim, device=device) fake_signals = generator(z) # 清空调优器梯度 opt_c.zero_grad() real_score = critic(real_signals) fake_score = critic(fake_signals) # Wasserstein距离估计:最大化 real_score - fake_score w_distance = real_score.mean() - fake_score.mean() gp = compute_gp(critic, real_signals, fake_signals.detach(), lambda_gp) critic_loss = -w_distance + gp critic_loss.backward() opt_c.step() critic_loss_epoch += critic_loss.item() # ---------- 更新生成器 ---------- z = torch.randn(batch_size, generator.latent_dim, device=device) fake_signals = generator(z) opt_g.zero_grad() fake_score = critic(fake_signals) # 生成器目标:让Critic对fake样本打出更高的分 generator_loss = -fake_score.mean() generator_loss.backward() opt_g.step() if (epoch + 1) % save_interval == 0: torch.save(generator.state_dict(), f'generator_epoch_{epoch+1}.pth')

4.2 关键参数的经验窗口和调整方向

WGAN-GP的训练参数在振动信号这个场景下有一个经验窗口。学习率这里用的是1e-4的Adam,之前试过普通GAN常用的2e-4和1e-3,训练前几百个epoch看着loss降得飞快,但生成样本的频谱会出现高频毛刺,这是优化步长过大导致生成器在局部震荡。learning rate降下来之后,所需epoch数大约翻倍,但生成质量明显提升。learning rate这个参数建议从1e-4起步,loss出现异常波动时先降到5e-5再看,不要一上来就用大学习率赌。

n_critic这个参数影响是两方面的:设得太小,Critic拟合不充分,生成器接收到的梯度噪声大,样本质量不稳定;设得太大,Critic强到梯度太干净,生成器反而容易停滞。在1024点一维信号这个任务上,n_critic=5稳,8偶尔更好但耗时增加60%。epoch数方面,在单张消费级显卡上2000轮左右能看到明显效果,1000轮内主要是粗结构成形,1500轮后冲击细节才开始出现。

beta值也是一个容易被忽略的设定。Adam默认的b1=0.9在这里会导致训练震荡,将b1调低到0.5是GAN训练中常见做法,b2保持0.9。原因是较低的一阶动量衰减能让优化器更快响应Critic输出的变化,避免陷入局部平坦区域。梯度惩罚系数lambda_gp=10在上一章介绍过,这里补充一句:如果数据分布特别稀疏,比如某些故障频率带的样本极少,可以试试把lambda_gp降到5,让生成器有更大探索空间。

5. 训练常见问题排查:模式坍缩、Loss波动、梯度爆炸

5.1 生成信号全部雷同(模式坍缩)

现象:不同批次生成出来的信号两两之间余弦相似度超过0.9,时域波形几乎完全重叠,频谱能量分布也集中在一个固定频带上。

原因:WGAN-GP对比普通GAN已经大幅缓解了模式坍缩,但生成器容量过大或学习率偏高时仍会发生。还有一个容易被忽略的诱因:训练数据里某一类故障样本占比过高,比如内圈故障有500条而外圈只有50条,生成器学会的分布会偏向多数类。

解决:先从数据入手,把各类别的样本数量对齐,按少样本类别数量做下采样或用滑窗重叠把少数类扩充到相近量级。然后调低生成器学习率到5e-5,同时把n_critic升到8,给Critic更多时间分辨真实样本中的细节差异。这个组合我试过多次,大多数情况3到5百个epoch内能恢复多样性。

5.2 Critic Loss持续走低但生成样本质量却在恶化

现象:打印出来的critic_loss从几十一路降到接近0,真实信号和生成信号的得分差也在缩小,但生成的信号时域波形变成一条近似正弦曲线,完全没有冲击特征。

原因:Wasserstein距离收敛只说明Critic对两个分布的评价趋于一致,无法保证生成分布真的覆盖真实分布。当multimodal分布时可能出现极小值陷阱,生成器找到一个能骗过Critic的狭窄区域就停住不动。这属于WGAN-GP的标准失效模式,不是代码错误。

解决:检查训练数据是否包含足够丰富的相位信息。轴承振动信号对相位非常敏感,同一故障模式下不同位置采集的信号相位差很大,如果数据处理时只保留整数周期的切片,相位多样性不足,生成器就没必要学复杂的冲击形态。把切分起点随机化,让每个样本在信号序列中的起始相位不完全一样,通常能解决。

5.3 训练中段出现NaN并伴随梯度异常放大

现象:epoch到某一轮附近,loss突然变成NaN,随后所有数值都跟着NaN。

原因:梯度惩罚计算中对梯度norm的平方在极端情况下会产生极大值,叠加Adam自适应学习率的更新后权重一步跳变。这类问题多发生在激活函数选择不当的情况下,比如Critic隐藏层用了ReLU,负半轴梯度全断,插值样本的反向传播梯度波动大。

解决:确认Critic和生成器隐藏层都使用的是LeakyReLU(0.2)。调低学习率至5e-5后重新训练。如果数据长度不是2的幂次,卷积下采样会在边界出现对齐问题,建议把信号pad或crop到1024这样的2次幂长度。这样每一层卷积stride=2都能整除。

5.4 训练能跑完但生成的样本只有轻微改进

现象:训练了上千个epoch,用包络谱对比生成样本和真实样本,故障特征频率的对齐度始终较差,冲击特征时有时无。

原因:特征对齐度差的直接原因是判别器对故障特征不敏感,本质是Critic的分辨能力不足。1024点信号中冲击成分占比往往很小,在时域上可能只占据几十个采样点的宽度,如果Critic的起始感受野过大,冲击细节在第一层卷积后就被抹掉了。

解决:把Critic第一层卷积核从8改小到4,保留更多局部信息。另一个有效方法是生成器的最后一层上采样倍数从4改成2两次,让网络有更多参数去精修冲击形态。还要确认RMSprop或者Adam的学习率过了1500个epoch后没有一直保持1e-4,我一般到1000轮后手动降一次到5e-5。

6. 生成效果验证:从时域波形、包络谱到特征分布的检查清单

WGAN训练完不能只看生成波形“像不像”,故障诊断场景要求生成样本的故障特征可追溯。我的验证流程分成四步,全部过完才敢把生成样本交给下游诊断模型。第一步是时域波形检查,滚动轴承故障信号最典型的形态是周期性的冲击衰减波,冲击间隔对应转频相关的故障特征频率。直接在生成样本里计算两个相邻冲击峰之间的采样点数,换算成频率,这个值应该和真实样本的故障特征频率对得上,误差控制在5%以内才合格。

第二步是包络谱对齐,这是诊断里最关键也最容易被忽略的环节。振动信号在低频段能量弱,直接用FFT观察特征频率不明显,标准做法是做Hilbert变换取包络,再对包络信号做FFT得到包络谱。计算生成样本和真实样本的包络谱,在故障特征频率及其二倍频、三倍频处应该都有明显峰值。批量验证时我写了一个简单脚本,循环比较生成样本和真实样本的前三阶特征频率峰值差。

from scipy.signal import hilbert import numpy as np def envelope_spectrum(signal, fs=12000): """计算包络谱,返回频率轴和幅值""" analytic = hilbert(signal) envelope = np.abs(analytic) # 去除直流分量 envelope = envelope - np.mean(envelope) spectrum = np.abs(np.fft.rfft(envelope)) freqs = np.fft.rfftfreq(len(signal), d=1.0/fs) return freqs, spectrum def check_feature_alignment(real_sig, fake_sig, fault_freq, fs=12000, tolerance=0.05): """检查生成样本在故障特征频率附近是否有对齐的峰值""" freqs, spec_real = envelope_spectrum(real_sig, fs) _, spec_fake = envelope_spectrum(fake_sig, fs) idx_fault = np.argmin(np.abs(freqs - fault_freq)) window = int(len(freqs) * 0.01) peak_real = np.max(spec_real[max(0, idx_fault-window): idx_fault+window+1]) peak_fake = np.max(spec_fake[max(0, idx_fault-window): idx_fault+window+1]) ratio = peak_fake / (peak_real + 1e-8) return ratio > 1 - tolerance

第三步和第四步可以用t-SNE或PCA把生成样本和真实样本的特征空间可视化。诊断模型如果有现成的中间层特征提取器,提取倒数第二层的特征做可视化最高效。生成样本和真实样本在图中应该交叠但不能完全混为一体,完全混在一起说明生成器只是照抄了训练样本,对下游泛化没有帮助。

我自己的教训是:验证这一步必须在每次训练完成之后强制完整走一遍,不要因为赶进度只做时域目测就急着用。有一次生成样本的包络谱基频恰好对齐了,但二倍频完全缺失,拿去训练诊断模型后故障严重程度误判率明显上升。从那以后,我每次训练完都会把四步验证脚本做成固定流程,生成数据进入下游模型之前先跑完包络谱对齐和t-SNE分布检查这两道关,宁可在验证上多花一小时,也不让劣质增强样本污染诊断模型。希望这套验证路径对你有帮助,如果照本文步骤训练后仍然出现生成样本和真实样本分布对不齐的情况,优先检查数据源信号是否包含完整的故障特征周期,这是大多数问题真正藏身的位置。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询