WGAN生成轴承故障振动信号：从数据不足到样本扩充的实战指南

简介：这份资源面向从事轴承故障诊断、信号处理与深度学习研究的工程师及学生，提供一套基于WGAN生成一维故障轴承振动信号的完整Python实现，用于缓解故障样本稀缺、类别不平衡等问题，适合具备一定TensorFlow基础的中高级读者复现与二次开发。压缩包共19个文件，约22.03MB，包含2个py训练与测试脚本、2个mat数据文件、9张png生成效果图，以及checkpoint、meta、index、data等模型权重与训练状态文件，另附md与readme说明文档，便于直接运行与结果对照。资源要求Python 3.5以上、TensorFlow GPU及numpy、scipy等环境，支持通过命令行调整学习率、训练轮数、采样频率，并可在9类信号中切换训练数据。目前已有3386人学习下载，读者可借此掌握WGAN在一维振动信号生成中的网络结构、训练流程与评估方式，快速搭建数据增强实验，为故障诊断模型提供扩充样本与排错参考。

1. 用WGAN生成故障轴承振动信号：从“数据不够”到“样本管够”的落地路径

设备故障诊断最头疼的问题从来不是算法不够先进，而是故障样本太少。一台正常运转的轴承，你可能连续采集几个月都碰不到一次内圈剥落或外圈裂纹，但深度学习模型偏偏需要大量故障样本来训练。更麻烦的是，不同工况下的振动信号分布差异极大，实验室台架上采集的故障数据搬到产线上往往直接翻车。GAN（生成对抗网络）的出现给了一条新路：既然真实故障样本难获取，那就让网络自己“造”出以假乱真的振动信号。而WGAN（Wasserstein GAN）相比原始GAN，在训练稳定性和生成质量上有明显优势，尤其适合振动信号这种一维时序数据。这篇内容面向已经了解轴承故障诊断基本流程、想用WGAN扩充故障样本集的工程师，从原理选型讲到代码复现，再到参数调优和踩坑记录，争取让你看完就能在自己的数据集上跑起来。

2. 为什么选WGAN而不是原始GAN：振动信号生成的选型逻辑

2.1 原始GAN在振动信号上的三个硬伤

原始GAN用JS散度衡量生成分布和真实分布的差异，这个选择在图像生成上勉强能用，但搬到一维振动信号上问题就被放大了。第一个硬伤是梯度消失：当判别器太强时，生成器几乎拿不到有效梯度，loss曲线看着在动，但生成的信号始终是噪声。第二个硬伤是模式崩溃：生成器发现只要反复输出某一段“万能信号”就能骗过判别器，于是所有生成样本长得一模一样，多样性直接归零。第三个硬伤是训练震荡：判别器和生成器的loss此起彼伏，你根本不知道什么时候该停。

振动信号本身的特点让这些问题更严重。轴承故障信号是典型的一维时序数据，采样率高、频带窄、故障特征往往藏在冲击成分的间隔和幅值调制里。原始GAN的判别器很容易通过捕捉局部噪声模式来区分真假，导致生成器学不到真正的故障冲击特征。

WGAN的核心改动是用Wasserstein距离替代JS散度。Wasserstein距离的好处是即使两个分布没有重叠，它也能提供有意义的梯度。具体实现上，WGAN做了三件事：去掉判别器最后一层的Sigmoid，把判别器改叫Critic；把loss改成真实样本得分减去生成样本得分；每次更新Critic后把权重裁剪到[-c, c]区间。这三步下来，训练稳定性提升非常明显。

2.2 WGAN-GP的改进与振动信号适配

权重裁剪有个副作用：它会把Critic的参数强行限制在一个小范围内，导致Critic学不到复杂的判别函数。WGAN-GP（Gradient Penalty）用梯度惩罚替代权重裁剪，让Critic的梯度范数接近1，既保证了Lipschitz约束，又不会过度限制模型容量。

对于振动信号生成，我一般会做以下适配：

生成器结构上，我习惯用三层一维转置卷积，每层后面接BatchNorm和LeakyReLU，最后用Tanh把输出压到[-1, 1]。Critic用四层一维卷积，每层接LeakyReLU，最后全局平均池化后输出一个标量。这个结构在轴承振动信号上表现稳定，生成样本的时域波形和包络谱都能和真实故障信号对上。

2.3 数据预处理：把振动信号切成WGAN能吃的形状

原始振动信号是长时序，直接喂给WGAN效果很差。常见做法是滑动窗口切分，每个窗口包含至少一个完整的故障冲击周期。以轴承故障为例，如果外圈故障特征频率是100Hz，采样率12kHz，那一个冲击周期对应120个采样点，窗口长度至少取1024才能覆盖多个周期。

import numpy as np from sklearn.preprocessing import MinMaxScaler def segment_signal(signal, window_size=1024, step=512): """ 滑动窗口切分振动信号 signal: 一维numpy数组，原始振动信号 window_size: 每个样本的采样点数 step: 滑动步长，通常取window_size的一半 返回: (n_samples, window_size)的数组 """ segments = [] for start in range(0, len(signal) - window_size + 1, step): segment = signal[start:start + window_size] segments.append(segment) return np.array(segments) def normalize_signal(segments): """ 归一化到[-1, 1]，适配Tanh输出 注意：按全局最大最小值归一化，不要按每个样本单独归一化 否则会破坏样本间的幅值相对关系 """ scaler = MinMaxScaler(feature_range=(-1, 1)) n_samples, window_size = segments.shape flat = segments.reshape(-1, 1) scaled = scaler.fit_transform(flat) return scaled.reshape(n_samples, window_size), scaler # 使用示例 raw_signal = np.load('bearing_vibration.npy') # 假设是单通道长信号 segments = segment_signal(raw_signal, window_size=1024, step=512) normalized, scaler = normalize_signal(segments) print(f"切分后样本数: {normalized.shape[0]}, 每样本长度: {normalized.shape[1]}")

这里有个容易忽略的点：归一化必须按全局做，不能每个样本单独归一化。如果每个样本各自缩放到[-1, 1]，那幅值信息就丢了，而轴承故障的严重程度恰恰体现在冲击幅值上。另外窗口长度建议取1024或2048，太短覆盖不了多个冲击周期，太长则单样本计算量太大。

3. 搭一个能跑的WGAN-GP：生成器和Critic的代码实现

3.1 生成器网络：从噪声到振动波形的映射

生成器的任务是把一个低维噪声向量映射成1024点的一维振动信号。我一般用全连接层先把噪声升维到合适长度，再堆叠一维转置卷积逐步上采样。

import torch import torch.nn as nn class Generator(nn.Module): def __init__(self, noise_dim=100, output_length=1024): super(Generator, self).__init__() self.noise_dim = noise_dim self.output_length = output_length # 全连接层：噪声 -> 初始特征图 self.fc = nn.Sequential( nn.Linear(noise_dim, 256 * 8), nn.BatchNorm1d(256 * 8), nn.ReLU(True) ) # 一维转置卷积堆叠 self.conv_blocks = nn.Sequential( # 输入: (batch, 256, 8) nn.ConvTranspose1d(256, 128, kernel_size=4, stride=2, padding=1), nn.BatchNorm1d(128), nn.ReLU(True), # 输出: (batch, 128, 16) nn.ConvTranspose1d(128, 64, kernel_size=4, stride=2, padding=1), nn.BatchNorm1d(64), nn.ReLU(True), # 输出: (batch, 64, 32) nn.ConvTranspose1d(64, 32, kernel_size=4, stride=2, padding=1), nn.BatchNorm1d(32), nn.ReLU(True), # 输出: (batch, 32, 64) nn.ConvTranspose1d(32, 16, kernel_size=4, stride=2, padding=1), nn.BatchNorm1d(16), nn.ReLU(True), # 输出: (batch, 16, 128) nn.ConvTranspose1d(16, 1, kernel_size=4, stride=2, padding=1), # 输出: (batch, 1, 256) ) # 自适应池化到目标长度 self.adaptive_pool = nn.AdaptiveAvgPool1d(output_length) self.tanh = nn.Tanh() def forward(self, z): x = self.fc(z) x = x.view(x.size(0), 256, 8) x = self.conv_blocks(x) x = self.adaptive_pool(x) return self.tanh(x)

生成器的关键参数是noise_dim和output_length。noise_dim一般取100或128，太小生成多样性不足，太大训练收敛慢。output_length要和预处理时的窗口长度一致。转置卷积的kernel_size和stride决定了上采样倍数，这里用kernel_size=4、stride=2的组合，每层长度翻倍。最后用AdaptiveAvgPool1d强制对齐到目标长度，避免手工计算每层尺寸。

3.2 Critic网络：用一维卷积提取振动特征

Critic不用Sigmoid，输出是实数分数。结构上我用一维卷积逐层降采样，最后全局平均池化。

class Critic(nn.Module): def __init__(self, input_length=1024): super(Critic, self).__init__() self.conv_blocks = nn.Sequential( # 输入: (batch, 1, 1024) nn.Conv1d(1, 32, kernel_size=5, stride=2, padding=2), nn.LeakyReLU(0.2, inplace=True), # 输出: (batch, 32, 512) nn.Conv1d(32, 64, kernel_size=5, stride=2, padding=2), nn.LeakyReLU(0.2, inplace=True), # 输出: (batch, 64, 256) nn.Conv1d(64, 128, kernel_size=5, stride=2, padding=2), nn.LeakyReLU(0.2, inplace=True), # 输出: (batch, 128, 128) nn.Conv1d(128, 256, kernel_size=5, stride=2, padding=2), nn.LeakyReLU(0.2, inplace=True), # 输出: (batch, 256, 64) ) self.global_pool = nn.AdaptiveAvgPool1d(1) self.fc = nn.Linear(256, 1) def forward(self, x): x = self.conv_blocks(x) x = self.global_pool(x).squeeze(-1) return self.fc(x)

Critic里没有BatchNorm，这是WGAN-GP的常见做法，因为BatchNorm会引入样本间的依赖，破坏梯度惩罚的独立性。LeakyReLU的负斜率取0.2，比默认的0.01更适合振动信号，能让负半轴的梯度更好地回传。

3.3 梯度惩罚与训练循环

WGAN-GP的核心在梯度惩罚项。每次迭代时，在真实样本和生成样本之间随机插值，计算Critic对插值样本的梯度范数，惩罚其偏离1的程度。

def gradient_penalty(critic, real, fake, device, lambda_gp=10): """ 计算WGAN-GP的梯度惩罚项 real: 真实样本 (batch, 1, length) fake: 生成样本 (batch, 1, length) lambda_gp: 惩罚系数，常用10 """ batch_size = real.size(0) # 随机插值系数 alpha = torch.rand(batch_size, 1, 1, device=device) # 插值样本 interpolated = alpha * real + (1 - alpha) * fake interpolated.requires_grad_(True) # Critic对插值样本的评分 score = critic(interpolated) # 计算梯度 grad = torch.autograd.grad( outputs=score, inputs=interpolated, grad_outputs=torch.ones_like(score), create_graph=True, retain_graph=True )[0] # 梯度范数 grad_norm = grad.view(batch_size, -1).norm(2, dim=1) # 惩罚项：梯度范数偏离1的平方 penalty = ((grad_norm - 1) ** 2).mean() return lambda_gp * penalty def train_wgan_gp(generator, critic, dataloader, epochs=2000, n_critic=5, lr=1e-4, device='cuda'): """ WGAN-GP训练循环 n_critic: 每更新一次生成器，Critic更新n_critic次 """ g_optimizer = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.9)) c_optimizer = torch.optim.Adam(critic.parameters(), lr=lr, betas=(0.5, 0.9)) generator.to(device) critic.to(device) for epoch in range(epochs): for i, (real_signal,) in enumerate(dataloader): real_signal = real_signal.to(device) batch_size = real_signal.size(0) # ---- 训练Critic ---- for _ in range(n_critic): z = torch.randn(batch_size, 100, device=device) fake_signal = generator(z).detach() c_real = critic(real_signal).mean() c_fake = critic(fake_signal).mean() gp = gradient_penalty(critic, real_signal, fake_signal, device) c_loss = c_fake - c_real + gp c_optimizer.zero_grad() c_loss.backward() c_optimizer.step() # ---- 训练生成器 ---- z = torch.randn(batch_size, 100, device=device) fake_signal = generator(z) g_loss = -critic(fake_signal).mean() g_optimizer.zero_grad() g_loss.backward() g_optimizer.step() if (epoch + 1) % 100 == 0: print(f"Epoch [{epoch+1}/{epochs}] " f"C_loss: {c_loss.item():.4f} " f"G_loss: {g_loss.item():.4f} " f"W_distance: {(c_real - c_fake).item():.4f}") return generator, critic

训练循环里有几个参数需要关注。n_critic=5表示每训练一次生成器，Critic训练五次，这是WGAN论文推荐的值。学习率用1e-4，Adam的betas设为(0.5, 0.9)，这是WGAN-GP的标配。W_distance是真实样本得分减去生成样本得分，这个值应该随着训练逐渐减小，如果它一直震荡或增大，说明训练出了问题。

4. 生成质量评估与故障诊断验证：怎么证明“假数据”真有用

4.1 时域和频域的双重评估指标

生成样本不能只看loss，必须从时域和频域两个维度验证。时域看波形形态和统计特征，频域看包络谱里的故障特征频率是否对得上。

import numpy as np from scipy.fft import fft, fftfreq from scipy.stats import kurtosis, skew def evaluate_generated_signals(real_signals, fake_signals, fs=12000): """ 评估生成信号质量 real_signals: 真实信号 (n, length) fake_signals: 生成信号 (n, length) fs: 采样率 """ metrics = {} # 时域统计特征对比 for name, signals in [('real', real_signals), ('fake', fake_signals)]: metrics[f'{name}_mean'] = np.mean(np.abs(signals)) metrics[f'{name}_std'] = np.std(signals) metrics[f'{name}_kurtosis'] = np.mean(kurtosis(signals, axis=1)) metrics[f'{name}_skewness'] = np.mean(skew(signals, axis=1)) # 包络谱对比 def envelope_spectrum(signal): # Hilbert变换求包络 from scipy.signal import hilbert analytic = hilbert(signal) envelope = np.abs(analytic) # 去均值 envelope = envelope - np.mean(envelope) # FFT n = len(envelope) spectrum = np.abs(fft(envelope))[:n//2] * 2 / n freqs = fftfreq(n, 1/fs)[:n//2] return freqs, spectrum # 取平均包络谱 real_env = np.mean([envelope_spectrum(s)[1] for s in real_signals[:50]], axis=0) fake_env = np.mean([envelope_spectrum(s)[1] for s in fake_signals[:50]], axis=0) freqs = envelope_spectrum(real_signals[0])[0] # 找包络谱峰值对应的频率 real_peak_freq = freqs[np.argmax(real_env)] fake_peak_freq = freqs[np.argmax(fake_env)] metrics['real_peak_freq'] = real_peak_freq metrics['fake_peak_freq'] = fake_peak_freq metrics['freq_error'] = abs(real_peak_freq - fake_peak_freq) return metrics # 使用示例 real_samples = normalized[:200] # 取200个真实样本 z = torch.randn(200, 100).to(device) with torch.no_grad(): fake_samples = generator(z).cpu().numpy().squeeze(1) metrics = evaluate_generated_signals(real_samples, fake_samples) for k, v in metrics.items(): print(f"{k}: {v:.4f}")

评估时重点看三个指标：kurtosis（峭度）反映冲击成分的强弱，故障轴承信号的峭度通常大于3；包络谱峰值频率应该对应故障特征频率；生成样本的std不能太小，否则说明模式崩溃了。如果生成信号的峭度远低于真实信号，说明生成器没学到冲击特征，需要调整网络结构或增加训练轮次。

4.2 用生成样本扩充训练集做故障分类

生成样本的最终目的是提升故障诊断模型的性能。我一般会做对比实验：只用真实样本训练分类器 vs 真实样本+生成样本训练分类器，看测试集准确率的变化。

from sklearn.svm import SVC from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score, classification_report def fault_classification_experiment(real_data, real_labels, generator, device, n_generate_per_class=200): """ 对比实验：真实样本 vs 真实+生成样本 real_data: 真实信号 (n, length) real_labels: 标签 (n,) """ # 提取时域统计特征作为分类器输入 def extract_features(signals): features = [] for s in signals: feat = [ np.mean(np.abs(s)), np.std(s), kurtosis(s), skew(s), np.max(np.abs(s)), np.sqrt(np.mean(s**2)) ] features.append(feat) return np.array(features) # 只用真实样本 X_real = extract_features(real_data) X_train, X_test, y_train, y_test = train_test_split( X_real, real_labels, test_size=0.3, random_state=42, stratify=real_labels ) clf_real = SVC(kernel='rbf', C=10, gamma='scale') clf_real.fit(X_train, y_train) acc_real = accuracy_score(y_test, clf_real.predict(X_test)) # 真实+生成样本 generated_signals = [] generated_labels = [] for class_id in np.unique(real_labels): z = torch.randn(n_generate_per_class, 100).to(device) with torch.no_grad(): fake = generator(z).cpu().numpy().squeeze(1) generated_signals.append(fake) generated_labels.extend([class_id] * n_generate_per_class) generated_signals = np.vstack(generated_signals) generated_labels = np.array(generated_labels) X_gen = extract_features(generated_signals) X_combined = np.vstack([X_real, X_gen]) y_combined = np.concatenate([real_labels, generated_labels]) X_train_c, X_test_c, y_train_c, y_test_c = train_test_split( X_combined, y_combined, test_size=0.3, random_state=42, stratify=y_combined ) clf_combined = SVC(kernel='rbf', C=10, gamma='scale') clf_combined.fit(X_train_c, y_train_c) acc_combined = accuracy_score(y_test_c, clf_combined.predict(X_test_c)) print(f"仅真实样本准确率: {acc_real:.4f}") print(f"真实+生成样本准确率: {acc_combined:.4f}") print(f"提升: {(acc_combined - acc_real) * 100:.2f}%") return acc_real, acc_combined

这个实验里，生成样本的数量要控制好。如果生成样本远多于真实样本，分类器会被生成样本主导，反而可能降低性能。我一般让生成样本和真实样本数量相当，或者生成样本略少一些。另外分类器用SVM就够了，不用上深度学习，因为我们要验证的是生成样本的质量，不是分类器的能力。

5. 避坑与排查：WGAN生成振动信号时最容易翻车的五个地方

5.1 生成信号全是噪声，loss看着正常但样本没法用

现象：训练几千轮后，生成的信号波形杂乱无章，和噪声没区别，但Critic loss和Generator loss都在缓慢下降，看不出异常。

原因：最常见的原因是Critic太强，生成器拿不到有效梯度。在WGAN-GP里，如果n_critic设得太大（比如10以上），或者Critic的学习率比生成器高，Critic很快就能完美区分真假样本，生成器的梯度就消失了。

解决：把n_critic降到3~5，Critic和生成器用相同的学习率。另外检查梯度惩罚系数lambda_gp，如果太小（比如1），Lipschitz约束不够，Critic会变得太强；如果太大（比如100），Critic被限制得太死，学不到东西。我一般从10开始调，观察W_distance的变化，如果它很快降到0附近并保持不变，说明Critic太强了。

5.2 生成样本多样性不足，所有样本长得差不多

现象：生成1000个样本，画出来发现波形几乎一样，只是幅值有微小差异。计算样本间的欧氏距离，发现远小于真实样本间的距离。

原因：模式崩溃。生成器发现只要输出某一段“万能信号”就能骗过Critic，于是放弃了学习完整的故障分布。在振动信号上，这个问题特别容易出现在故障特征频率比较单一的情况下。

解决：首先检查噪声维度，noise_dim=100是底线，如果降到32或16，多样性肯定不够。其次可以在生成器里加Dropout层，训练时随机丢弃一些神经元，强迫生成器不依赖固定模式。另外可以尝试minibatch discrimination，让Critic同时看一批样本，惩罚过于相似的生成结果。我自己的经验是，在生成器的全连接层后加一个Dropout(0.3)，模式崩溃的概率会明显降低。

5.3 包络谱峰值对不上故障特征频率

现象：生成信号的时域波形看着还行，但包络谱里的峰值频率和真实故障特征频率偏差超过10%。

原因：生成器没有学到故障冲击的周期特性。一维转置卷积的感受野有限，如果故障冲击周期比较长，生成器可能只学到了局部波形，没学到全局的周期结构。

解决：增大生成器的感受野。具体做法是增加转置卷积的层数，或者用空洞卷积（dilated convolution）扩大感受野。另一个办法是在生成器里加入自注意力机制，让网络能捕捉长距离依赖。我一般会先检查窗口长度是否覆盖了至少3个冲击周期，如果窗口太短，生成器根本看不到完整的周期结构，包络谱肯定对不上。

5.4 训练后期W_distance突然崩掉

现象：训练前1000轮W_distance稳步下降，然后突然跳到很大的值，生成样本质量急剧恶化。

原因：梯度惩罚项计算时，插值样本的梯度可能爆炸。如果真实样本和生成样本的幅值范围差异很大，插值后的样本可能落在极端区域，导致梯度范数异常。

解决：确保真实样本和生成样本都归一化到相同的范围（比如[-1, 1]）。另外可以在梯度惩罚计算时加一个小的epsilon，防止除零。如果崩掉后恢复不了，就降低学习率，从1e-4降到5e-5，重新加载崩掉之前的模型继续训练。我习惯每500轮保存一次模型，崩了就从最近的检查点恢复，比从头训练省时间。

5.5 生成样本用于分类后准确率反而下降

现象：用真实+生成样本训练分类器，测试准确率比只用真实样本还低。

原因：生成样本的质量参差不齐，有些样本虽然看着像故障信号，但特征分布和真实故障有偏差，相当于引入了噪声标签。另外如果生成样本数量远多于真实样本，分类器会被生成样本主导。

解决：不要把所有生成样本都塞进训练集。我一般会先用包络谱峰值频率筛选一遍，只保留峰值频率和真实故障特征频率偏差小于5%的生成样本。另外控制生成样本和真实样本的比例在1:1到1:2之间，不要超过2:1。如果准确率还是下降，说明生成样本质量不够，需要回到训练阶段继续调WGAN。

6. 从生成到落地：把WGAN变成故障诊断流水线里的一个模块

6.1 条件WGAN：按故障类型定向生成

无条件WGAN生成的是混合故障样本，你没法控制它生成的是内圈故障还是外圈故障。实际产线上，你往往需要针对特定故障类型补充样本。条件WGAN（Conditional WGAN）把故障标签作为额外输入，生成器和Critic都接收标签信息。

class ConditionalGenerator(nn.Module): def __init__(self, noise_dim=100, n_classes=4, output_length=1024): super(ConditionalGenerator, self).__init__() # 标签嵌入 self.label_embedding = nn.Embedding(n_classes, 50) # 噪声+标签嵌入拼接后输入 self.fc = nn.Sequential( nn.Linear(noise_dim + 50, 256 * 8), nn.BatchNorm1d(256 * 8), nn.ReLU(True) ) # 后续转置卷积层与无条件生成器相同 self.conv_blocks = nn.Sequential( nn.ConvTranspose1d(256, 128, 4, 2, 1), nn.BatchNorm1d(128), nn.ReLU(True), nn.ConvTranspose1d(128, 64, 4, 2, 1), nn.BatchNorm1d(64), nn.ReLU(True), nn.ConvTranspose1d(64, 32, 4, 2, 1), nn.BatchNorm1d(32), nn.ReLU(True), nn.ConvTranspose1d(32, 16, 4, 2, 1), nn.BatchNorm1d(16), nn.ReLU(True), nn.ConvTranspose1d(16, 1, 4, 2, 1), ) self.adaptive_pool = nn.AdaptiveAvgPool1d(output_length) self.tanh = nn.Tanh() def forward(self, z, labels): label_emb = self.label_embedding(labels) x = torch.cat([z, label_emb], dim=1) x = self.fc(x) x = x.view(x.size(0), 256, 8) x = self.conv_blocks(x) x = self.adaptive_pool(x) return self.tanh(x)

条件WGAN的训练循环和无条件版本基本一致，只是每次前向传播时要同时传入噪声和标签。Critic也要接收标签，通常做法是把标签嵌入后和信号特征拼接。条件WGAN的好处是你可以指定生成“内圈故障”或“外圈故障”，生成样本的标签是确定的，不需要额外标注。

6.2 在线生成：训练好的生成器怎么嵌入诊断系统

训练好的生成器可以导出为TorchScript或ONNX，嵌入到诊断系统的推理流程里。当系统检测到某类故障样本不足时，实时调用生成器补充样本。

# 导出为TorchScript generator.eval() example_noise = torch.randn(1, 100) example_label = torch.tensor([0]) traced_generator = torch.jit.trace(generator, (example_noise, example_label)) traced_generator.save('wgan_generator.pt') # 推理时加载 loaded_generator = torch.jit.load('wgan_generator.pt') loaded_generator.eval() def generate_fault_samples(fault_type, n_samples=100): """ 按故障类型生成样本 fault_type: 0-正常, 1-内圈故障, 2-外圈故障, 3-滚动体故障 """ z = torch.randn(n_samples, 100) labels = torch.full((n_samples,), fault_type, dtype=torch.long) with torch.no_grad(): samples = loaded_generator(z, labels) return samples.numpy().squeeze(1)

在线生成时要注意推理速度。生成器本身计算量不大，1000个样本在GPU上不到1秒。但如果诊断系统部署在边缘设备上，可能需要把生成器量化成FP16或INT8。我一般会在导出前做一次量化感知训练，精度损失控制在1%以内。

6.3 一个我踩过的坑：生成样本的幅值分布偏移

最后说一个我实际踩过的坑。训练WGAN时，我用MinMaxScaler把信号归一化到[-1, 1]，生成器输出也用Tanh压到[-1, 1]。看起来没问题，但生成样本的幅值分布明显偏向0附近，大冲击成分的幅值比真实样本低。原因是Tanh在接近±1时梯度很小，生成器倾向于输出中间值来避免梯度消失。

解决办法是在生成器最后一层用线性输出，然后在loss里加一个幅值约束项，惩罚生成样本和真实样本的幅值分布差异。或者更简单：在归一化时不要用MinMaxScaler，改用StandardScaler，让数据分布接近标准正态，Tanh的输出范围刚好匹配。这个改动让生成样本的峭度从2.1提升到了3.8，和真实故障信号的峭度基本一致。

做故障诊断这行，数据永远是最贵的。WGAN给了一条低成本扩充样本的路，但这条路不是一键生成的，参数要调、结构要改、生成样本要筛。我现在的习惯是每拿到一个新工况的数据，先花半天时间把WGAN跑通，生成一批样本，然后用包络谱筛一遍，最后混入训练集做对比实验。这套流程跑下来，通常能把小样本故障诊断的准确率提升5到15个百分点。希望帮到你。

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