企业官网建设流程全解析

去年帮某制造企业做设备健康管理项目，对方提了个很现实的需求：现场维修师傅看不懂频谱图，能不能直接告诉他轴承是外圈坏了还是内圈坏了？我们试了一圈方案，最后落地的是——用 Pytorch 把原始振动信号转成小波时频图，再交给 SwinTransformer 做轴承故障诊断。这篇文章就把整套方案的思路、代码和踩坑过程拆开讲，适合正在研究故障诊断方向的研究生，也适合想把深度学习模型真正用起来的设备工程师。

1. 为什么轴承故障诊断要绕道"时频图"再喂给Transformer

1.1 一维CNN做轴承诊断的三个局限

很多论文的起点是把原始振动信号直接扔给一维卷积神经网络，这种做法在实验室里跑通很容易，但放到工程现场往往不太经得起推敲。

第一是感受野问题。一维卷积的感受野是沿着时间轴线性扩展的，轴承故障冲击往往非常短暂，内圈故障转一圈只出现几次窄脉冲，中间夹杂着大量平稳振动。如果卷积核不够大、网络层数不够深，模型很容易被占多数的"无故障段"带偏。想覆盖完整的故障冲击周期，就得堆层数或者加大卷积核，参数和计算量都不小。

第二是工况敏感。同一套模型在转速 1800rpm 下调出来，换到 1200rpm 的现场电机上，分布立刻变样。一维信号里混着转速、载荷、环境噪声的信息，模型很难把"故障模式"和"工况差异"分开。

第三是解释性差。一维CNN最后接全连接层输出类别，你很难说清楚模型到底学到了什么。现场故障如果判断错，排查起来很痛苦，相当于一个黑盒给了你一个你不信任的答案。

1.2 小波时频图的优势：从"有什么频率"到"频率何时出现"

傅里叶变换能告诉我们振动信号里有哪些频率成分，但丢掉了一个关键信息——这些频率在什么时间出现。轴承故障产生的冲击是非平稳的、瞬态的，单看频谱时，某些短时冲击的能量会被平均到整个时间轴上，很容易被其他背景频率淹没。

短时傅里叶变换（STFT）用固定窗口切信号，窗口长了时间分辨率差，窗口短了频率分辨率差，这个矛盾是固定的。而连续小波变换通过缩放小波基函数，在不同频段自动调节时间-频率分辨率：低频段频率分辨率高，高频段时间分辨率高。恰好匹配轴承故障信号的特点——故障冲击是瞬态高频，同时又被旋转频率调制，需要同时看清楚"什么时候发生冲击"和"冲击分布在哪个频带"。

把CWT系数的模值画成热力图，就是小波时频图。不同故障类别在时频图上的纹理差异非常明显：正常轴承能量集中在低频连续带上；外圈故障会出现周期性的垂直冲击线；内圈故障的冲击线幅度随载荷区周期起伏；滚动体故障则表现出冲击线位置逐渐漂移的特征。这些视觉差异，正是图像分类模型能区分的依据。

1.3 SwinTransformer为什么适合图像化的振动信号

时频图虽然是"人造图像"，但它和自然图像有一个共通点：信息集中在局部纹理上，而不是某个像素的绝对位置。ViT 的全局自注意力复杂度是 O(N²)，时频图 patch 化后序列长度不短，训练数据不够时极容易欠拟合。

SwinTransformer 把自注意力限制在窗口内部，复杂度从图像尺寸的平方降为线性；又通过移位窗口让不同窗口之间的信息可以跨窗口流动。这种"先局部、再全局"的思路，正好匹配时频图上的多尺度特征：故障冲击线是最小尺度的高频细节，调制边带是中尺度结构，整体能量分布是大尺度背景。Swin 的层级结构把 C、2C、4C、8C 的特征逐级拼出来，四类特征同时被利用。

工程上还有一个优势：Swin 在 ImageNet 上提供了预训练权重，用预训练模型微调时频图任务，比从零训练一维CNN稳定得多。这一点在小样本故障诊断场景里几乎决定成败。

2. 环境准备与轴承振动数据预处理

2.1 Pytorch环境与依赖库版本

方案定了，先搭环境。我用的版本组合比较常规，Python 3.10、Pytorch 2.1、CUDA 11.8。如果你买的新显卡是 Ada 架构，CUDA 版本建议 12.x，差别不大，关键是 Pytorch 和 CUDA 的匹配关系别搞错，否则 import torch 会直接报错。

除了 torch 和 torchvision，还需要几个信号处理和图像处理库：pywavelets 做小波变换，matplotlib 保存图片，timm 加载 SwinTransformer 预训练模型，opencv-python 做图像预处理，scikit-learn 用来算指标和混淆矩阵。安装命令长期不变：

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 pip install timm pywavelets matplotlib opencv-python scikit-learn pandas tqdm

有一个版本细节值得提醒：timm 版本不要用太老的，早期版本的 Swin 模型命名不规范，后来统一成了swin_tiny_patch4_window7_224这种写法。建议 timm 不低于 0.9.0，我遇到过老版本读取预训练权重时 key 对不上的情况，排查起来很浪费时间。

2.2 轴承数据获取与切片策略

实验用的是某公开滚动轴承故障数据集，包含正常（Normal）、外圈故障（OR）、内圈故障（IR）、滚动体故障（B）四类状态，每种故障还分了不同的损伤直径。有人会直接把原始 mat 文件里的长信号全部读进来训练，这不对——样本数量太少，模型学不到鲁棒特征，必须做切片。

采样率是 12kHz，电机转速约 1730rpm 到 1790rpm，换算下来每个故障冲击周期大概 400 个采样点。我采用的切片策略是每段 1024 点、步长 512 个点滑窗，相当于每段覆盖两三个冲击周期，保留足够的瞬态细节。每类样本各切出几千段，全数据集在 1 万到 2 万张图之间。

这里有一个全局大坑：轴承振动信号是高度自相关的，同一段连续信号里相邻两个切片几乎一模一样。划分训练集和测试集时，绝对不能再把所有切片混在一起后随机打乱，必须按原始信号文件的先后顺序划分。同一个文件里的所有切片只能全部进入训练集，或者全部进入验证集、测试集。否则相邻切片会同时出现在训练集和测试集里，测试准确率虚高到 99%，拿到现场数据立刻打回原形。

2.3 生成小波时频图的完整代码

CWT 的原理一句话讲：选定一个小波基函数，让它与信号在不同尺度上做内积，尺度与频率存在反比关系——尺度小对应高频，尺度大对应低频。PyWavelets 的pywt.cwt直接封装了整个流程。

import numpy as np import pywt import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt sampling_rate = 12000 scales = np.arange(1, 128) wavelet = 'morl' image_size = 224 def signal_to_time_frequency_image(signal: np.ndarray, save_path: str): # 连续小波变换，返回系数矩阵和对应的伪频率 coef, freqs = pywt.cwt(signal, scales, wavelet, sampling_period=1.0 / sampling_rate) # 取模，做幅度归一化 mag = np.abs(coef) mag = mag / (mag.max() + 1e-9) # 直接将幅度矩阵保存为伪彩色图片 plt.imsave(save_path, mag, cmap='jet') # 返回归一化后的矩阵，方便后续直接加载 return mag

有几个细节必须说清楚。

第一，scales从 1 到 128 是经验值。采样率 12kHz 下，这个尺度范围对应的伪频率能够覆盖从几十赫兹到几千赫兹的频带，轴承故障特征频率多数落在这个区间。如果你的数据采样率不同，要重新算尺度范围，不能照搬他人代码。

第二，cmap用jet还是viridis都可以，但整个数据集必须统一。如果训练图用 jet、测试图用 viridis，模型学到的是颜色映射差异而不是故障纹理差异，这在调试时非常隐蔽。

第三，plt.imsave保存 JPG 是有损压缩的，如果追求极致可复现性，改用 PNG 或者直接保存归一化后的 npy 矩阵。工程上为了省磁盘我还是用的 JPG，但固定了quality=95，效果可以接受。

第四，所有图片统一 resize 到 224x224，这一步在保存图片时就做完，不要在训练时反复做不同尺寸的缩放，否则不同 batch 之间的分布会有细微漂移。

from PIL import Image import os def resize_time_frequency_image(src_path: str, dst_path: str): img = Image.open(src_path).convert('RGB') img = img.resize((image_size, image_size), Image.BILINEAR) img.save(dst_path)

3. SwinTransformer模型搭建细节

3.1 它到底在做什么：窗口自注意力与移位窗口

SwinTransformer 最核心的机制是窗口自注意力。输入特征图被切成一个个不重叠的窗口，例如 7x7 的格子，每个窗口内部做标准自注意力。这个设计让计算复杂度从整张图的 O(H²W²) 降为 O(HW·w²)，其中 w 是窗口边长。图像变大时，复杂度近似线性增长，这是它比 ViT 更适合实际训练的原因之一。

但窗口切分有个副作用：不同窗口之间完全没信息交换，右边界上的像素永远看不到左边界上的上下文。Swin 的解法是移位窗口——在下一层把窗口整体错位移动几个 patch，使上一层处于不同窗口的像素在下一层被分到同一个窗口里。通过这种交替的窗口模式，信息可以在多层之间不断扩散，最终达到全局感受野的效果。

窗口自注意力里还加了一个可学习的相对位置偏置，每个注意力头都有一组独立的 bias。这对时频图任务特别有价值：模型可以对"冲击线出现在时频图哪个相对位置、相隔多远"更加敏感，不必每个像素都硬学。故障诊断里冲击的周期性正是通过这种相对位置关系体现出来的。

Swin 的层级结构也不能忽略。Patch Merging 不断合并相邻 patch，特征图从 1/4 分辨率一路到 1/32 分辨率，通道数逐级翻倍。时频图上的冲击线属于小尺度纹理，调制边带属于中尺度结构，整体能量分布属于大尺度背景。这种金字塔式的特征表达，和 CWT 天然的多分辨率特性非常契合。

3.2 两种接入方式：timm一行代码与官方源码

接入 SwinTransformer 有两种常见方式。

第一种是直接使用 timm：

import timm import torch.nn as nn model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=True, num_classes=4 )

这一行代码会做三件事：加载 Swin-Tiny 结构、读取 ImageNet 预训练权重、把最后的分类头替换成 4 类输出。多数场景下够用了，不需要自己动手改造模型内部结构。

第二种是拉取官方仓库的 models 文件放进项目里，自己写前向传播。这种方式适合你要改 patch size、窗口大小、或者想做权重初始化等自定义操作的时候。缺点是代码量明显变大，还需要自己处理相对位置编码的初始化。

我的建议是先用 timm 跑通完整流程，确认实验结果没问题了再考虑要不要换成官方源码。不要在一开始就把精力花在模型实现细节上，模型结构和振动信号的关系弄清楚才是本篇的重点。

timm 里 Swin 模型的num_classes参数已经帮你替换了分类头，非常方便。如果你需要更细的控制，也可以这样改：

in_features = model.head.in_features model.head = nn.Linear(in_features, 4)

故障诊断通常类别数很少，4 类、10 类都常见，改一行就够了。

3.3 针对时频图场景的模型配置

Swin-Tiny 是参数量和效果比较均衡的选择，约 28M 参数。如果数据量只有几千张图，可以降低drop_path_rate防止过拟合，这是 Swin 的内置正则化手段，类似 Dropout 但作用于路径：

model = timm.create_model( 'swin_tiny_patch4_window7_224', pretrained=True, num_classes=4, drop_path_rate=0.2 )

预训练权重的处理也值得说。时频图和自然图像差异确实存在，但 ImageNet 预训练模型学习的底层纹理特征——边缘、角点、局部对比度——仍然可以迁移。我的实测是：加载预训练权重微调，比从零训练准确率高 2 到 3 个百分点，收敛速度也快很多。前提是数据量不能太小，否则预训练权重反而会带来偏置。

如果数据量非常小（每类只有几百张），可以采用分层解冻策略：先冻结 backbone 只训练分类头 5 个 epoch，再解冻最后两层，最后才全量微调。学习率从 1e-4 起步，每个阶段除以 10。这个方法在故障诊断这种典型小样本场景里非常管用。

4. 训练流程、超参数与实验结果

4.1 Dataset与DataLoader的实现

数据准备好后，第一步是把图片路径和标签组织成 Pytorch 的 Dataset。用 ImageFolder 是最快的做法，目录结构如下：

data/ train/ normal/xxx.jpg inner/xxx.jpg outer/xxx.jpg ball/xxx.jpg val/ normal/xxx.jpg ...

然后用torchvision.datasets.ImageFolder直接加载，省去自己写 Dataset 的麻烦。但如果你需要控制文件的划分方式（比如按原始文件分组），建议还是自定义 Dataset：

from torch.utils.data import Dataset from PIL import Image class FaultDataset(Dataset): def __init__(self, image_paths, labels, transform=None): self.image_paths = image_paths self.labels = labels self.transform = transform def __len__(self): return len(self.image_paths) def __getitem__(self, idx): image = Image.open(self.image_paths[idx]).convert('RGB') if self.transform: image = self.transform(image) return image, self.labels[idx]

有一个训练细节值得特别注意：时频图和平常的猫狗图片不一样，它的横轴是时间、纵轴是频率，物理语义非常强。随机水平翻转和随机垂直翻转都不能随便用。水平翻转会把时间反向，一个外圈故障的冲击序列顺序就颠倒了；垂直翻转会把高频变低频，模型学到的频率位置关系被破坏。这会严重损伤模型对故障模式的判断能力。

我在实验中用的数据增强组合是：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.8, 1.0)), transforms.RandomHorizontalFlip(p=0.0), # 默认关掉 transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

RandomResizedCrop相当于做轻微的尺度扰动，模拟工况变化；ColorJitter只调亮度和对比度，不改变颜色通道的波长语义，相对安全。Normalize 的均值和标准差用 ImageNet 的标准值，timm 预训练模型默认要求这个输入分布。

4.2 超参数配置与选择逻辑

训练阶段超参数的选择直接影响最终效果。下表是我在多个数据集上调出来的稳定配置：

为什么用 AdamW 而不是 SGD？Swin 原论文在 ImageNet 上就是用 AdamW 配余弦退火，这套组合在大模型上表现稳定，尤其是解耦权重衰减之后，权重衰减不会干扰 Adam 的动量估计。如果你用 SGD，只能收敛到差不多的水平，但需要调更多参数。

4.3 训练主循环代码

训练主循环可以直接参考这个骨架：

import torch import torch.nn as nn from torch.cuda.amp import autocast, GradScaler device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) criterion = nn.CrossEntropyLoss(label_smoothing=0.1) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50) scaler = GradScaler() best_acc = 0.0 for epoch in range(50): model.train() total_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss += loss.item() scheduler.step() # 验证 acc = validate(model, val_loader) if acc > best_acc: best_acc = acc torch.save({'model_state': model.state_dict(), 'config': config}, 'best_model.pth')

validate函数就是标准的无梯度推理，计算准确率和各类别 F1，这里不展开写。混合精度用autocast包住前向传播，可以减少显存占用，并且训练速度在 3090 上大概提升三到四成。注意验证阶段也需要用autocast，否则 fp16 训练的模型在 fp32 推理下结果会有微小差异。

4.4 实验对比：Swin-T vs ResNet50 vs 一维CNN

我在同一套时频图数据上做了三组对比实验，训练集约 12000 张、测试集约 4000 张，结果如下：

一维CNN的 92.6% 并非不能接受，但混淆矩阵显示它在外圈故障和滚动体故障之间容易混。时频图的意义在于给模型提供了明确的频带和纹理信息；Swin 的窗口注意力进一步强化了对局部冲击序列的建模，所以 98.7% 的准确率不是模型单纯堆参数堆出来的。

测试集上还观察到一个有意思的现象：Swin 在正常类和内圈故障类上的表现最稳，几乎零混淆。滚动体故障是最难分类的一类，原因是滚动体故障冲击位置随转动而移动，时频图上没有固定周期的冲击线，这种"游走型"纹理对模型的时序理解能力要求很高。

4.5 进一步的可视化：混淆矩阵与特征分布

只看准确率不够，我习惯把混淆矩阵和 t-SNE 特征分布一起画出来。混淆矩阵能告诉你错误集中在哪一对类别之间；t-SNE 能直观显示最后一个全连接层之前的特征向量是否聚类清晰。

Swin 的特征分布通常是四团明显的簇，簇与簇之间的边界很干净。而一维CNN的特征分布就散得多，尤其是滚动体故障的特征点会往外圈故障的簇里渗透。这也是 SSN 在故障诊断里真正值钱的地方——它把原始信号的复杂模式压缩成了紧致而可分的表征。

5. 实战中遇到的那些坑，每一个都可能让你的模型白跑

5.1 小波参数不一致：一个让实验不可复现的隐形杀手

我踩过的第一个坑是"小波参数不一致"。最初跑实验时，小波基函数、尺度范围、采样率这些参数散落在各个 notebook 里，后来换了环境重新跑数据，发现结果完全对不上。排查到最后，原因是第一次用morl小波、尺度 1-128；第二次用的是cmor小波、尺度 1-256，生成出来的时频图像素分布完全不同。

解决方案是统一用 config 文件管理所有预处理参数，参数写到同一个字典里，每次实验前打印一遍：

config = { 'wavelet': 'morl', 'scales': list(range(1, 128)), 'sampling_rate': 12000, 'sample_length': 1024, 'step': 512, 'image_size': 224, 'cmap': 'jet', 'seed': 42, }

同时固定所有随机种子，包括 Python、NumPy、Pytorch 的 random seed，否则即便数据一样，训练结果也可能有几毛钱的波动。小波参数和随机种子这种问题，表面看只是"结果对不上"，实际会让整个实验过程失去可信度。

5.2 数据泄漏：准确率 99% 的假象是怎么来的

这个问题必须单独拿出来讲。如果你把切片后的所有图片全部混在一起，然后随机按照 8:2 划分训练集和测试集，你会得到一个非常漂亮的准确率——我见过有人做到 99.5%，一度以为是模型太强。把训练好的模型拿去预测现场新采集的数据，准确率跌到 85% 以下，这才发现是数据泄漏。

原因是同一个原始文件里相邻切片的重叠率太高。步长 512、窗口 1024，意味着相邻切片有 50% 的内容完全一致。这些"相似图片"如果一部分进了训练集、一部分进了测试集，模型相当于直接记住了图片本身，而不是理解故障模式。

正确做法在前面提过：按原始信号文件分组划分数据集。每个传感器记录文件的所有切片只能属于一个集合。更严格的方案是采用 leave-one-file-out 的交叉验证，拿一个完整文件做测试集、其余做训练集，反复验证。这样得到的指标才是真实泛化能力。我也建议把这一点写进任何故障诊断实验的 code review 清单。

5.3 SwinTransformer 在小样本上的过拟合

Swin-Tiny 有 28M 参数，而故障图片通常只有几千到一两万张。从零训练时，训练集准确率很快到 99%，验证集却卡在 90% 上下，这就是典型的过拟合。

解决过拟合的组合拳有三个。第一，加载 ImageNet 预训练权重，这是最有效的——边缘、纹理等底层特征可以迁移到频率纹理上。第二，开启 Swin 的drop_path_rate，我试过 0.2 到 0.3 之间效果比较稳，太高会造成欠拟合。第三，加标签平滑label_smoothing=0.1，相当于告诉模型不要把训练集的硬标签学得太绝对。

还可以使用 MixUp 增强：

def mixup_data(x, y, alpha=0.2): lam = np.random.beta(alpha, alpha) index = torch.randperm(x.size(0)).to(x.device) mixed_x = lam * x + (1 - lam) * x[index] y_a, y_b = y, y[index] return mixed_x, y_a, y_b, lam

MixUp 在时频图任务上表现不错，因为它相当于把两类故障的时频纹理按比例叠加，模型被迫学习更泛化的边界。但不要叠加太强，alpha 超过 0.4 之后验证集准确率反而下降。

5.4 显存优化与推理速度考量

Swin-T 在 224x224 输入、batch size 32 的情况下，fp32 训练需要 10GB 左右显存。如果你的卡只有 8GB，会直接 OOM。三个解决办法。

一是开混合精度，显存占用几乎减半，训练速度还更快。二是减小 batch size 到 16，配合梯度累积：

accumulation_steps = 2 for i, (images, labels) in enumerate(train_loader): loss = loss / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

三是推理阶段可以转成 ONNX 或 TensorRT。轴承故障诊断的落地场景经常是边缘设备，Swin-T 单张推理在 GPU 上 8ms 不算慢，但在 CPU 上可能要 200ms 以上。如果部署环境没有 GPU，建议先蒸馏到一个小 CNN，或者换用更轻量的 Swin-Nano 变体。

5.5 类别不均衡与现场工况变化

公开数据集里正常类样本往往远多于故障类，现场采集时故障样本更难获得。类别不均衡会直接推高多数类的准确率，少数类被牺牲。处理办法有几层：首先在 loss 里加class_weight，或者直接用 Focal Loss；其次对少数类做过采样，让每个 batch 里各类数量尽量均衡；最后评估指标不能只看 accuracy，要看加权 F1 和混淆矩阵。

工况变化是另一个大问题。公开数据集大多是恒定转速、恒定载荷，现场设备不可能这么干净。模型在一种工况下训练、另一种工况下测试，准确率普遍下降 5 到 10 个百分点。解决方向是做 domain adaptation，或者对原始信号做幅值归一化和转速同步重采样，但这些属于进阶话题，超出这篇文章的篇幅，后续可以单独写一篇。

写在最后：一个当时被忽略的细节

项目做完一年多，这套方案还在现场跑着。回头整理时，我觉得最有价值的反而不是 98.7% 的准确率，而是整个流程中对物理意义和数据划分的坚持。很多人把故障诊断当成纯图像分类问题来做，大量精力花在调模型上，却忽略了时频图生成参数、数据泄漏、工况差异这些更底层的问题。模型选型再花哨，数据划分有问题，一切都是空中楼阁。

最后分享一个小技巧：保存模型权重时，一定同时保存一份预处理 config，包括小波基函数、尺度范围、切片长度、步长、图片尺寸、cmap 这些参数。这样三个月后重新推理时，你不需要翻几十个 notebook 去猜当时是怎么做数据的。这个小细节在工程落地时比调参重要得多。