简介：面向机械故障诊断与预测领域的研究人员与PyTorch开发者，这份资源提供一套基于图神经网络（GNN）的完整Python实现框架。框架采用PyTorch与PyTorch Geometrics，数据预处理阶段内置KNNGraph、RadiusGraph、PathGraph三种图构造方法，可将多传感器时间序列转换为图结构，充分保留空间关联特征；诊断阶段则通过节点分类和图分类回归任务完成故障识别与趋势预测。压缩包共111个文件，包含56个Python源码、53个编译pyc文件及2张框架示意图，整体仅1.6MB，包体紧凑；源码中除训练工具与通用图构建模块外，还涵盖CWRU、SEU、PU等多个轴承数据集案例，便于对照实验和二次开发。目前已有711人学习下载，对希望快速上手GNN故障诊断、扩展空间特征建模思路的读者具有直接参考价值。
1. 用GNN重构故障诊断：当振动信号不再是一维数组
传统故障诊断把多传感器振动信号拼接成矩阵喂给CNN或LSTM，但这里有个被忽略的假设：传感器之间是独立的。实际上，安装在轴承座、齿轮箱、电机端盖上的传感器共享同一套机械传动链，振动会沿着轴系传播，它们之间的关系应该用图来描述，而不是用矩阵的列。GNN（图神经网络）的引入就是为了让模型显式感知这种空间拓扑。本文这套基于PyTorch和PyTorch Geometric的框架，把信号预处理、构图、图分类/节点分类封装成一条完整流水线，支持KNNGraph、RadiusGraph、PathGraph三种构图策略，并在CWRU、SEU、PU三个公开数据集上给出了可直接改用的训练脚本。如果你已经跑通CNN做故障诊断但对准确率不满意，或者想在论文里换一个更有解释性的基线，这套代码值得拆开看。
2. 三种图构造方法：为什么KNNGraph不是万能的
2.1 从时间序列到图：先想清楚边的语义
把一维振动信号切成窗口后，每个窗口是一个样本。如果做节点分类，每个窗口作为一个节点，边的含义是“这两个样本在特征空间里接近”；如果做图分类，每个窗口内部再切子段，子段作为节点，边的含义是“子段之间存在时序或空间耦合”。这套框架两种都支持，但构图方式决定了模型能学到什么。
选构图方法时先问自己一个问题：你希望GNN捕捉哪种结构？如果传感器布点之间存在明确的物理连接（比如同一个轴承座上的径向和轴向测点），PathGraph最合理；如果没有任何先验拓扑，KNNGraph通过特征距离自动找邻居；RadiusGraph则用固定阈值控制邻居数量，适合特征分布相对均匀的场景。
2.2 CWRUKnn.py：KNN构图的核心逻辑
from sklearn.neighbors import NearestNeighbors import torch def knn_graph(features, k=5): """ features: [N, D] N个节点, D维特征 k: 每个节点的邻居数量 """ neigh = NearestNeighbors(n_neighbors=k + 1, metric='cosine') neigh.fit(features) knn_indices = neigh.kneighbors(features, return_distance=False)[:, 1:] # 构造COO格式边索引 n = features.shape[0] src = [] dst = [] for i in range(n): for j in knn_indices[i]: src.append(i) dst.append(j) edge_index = torch.tensor([src, dst], dtype=torch.long) return edge_index逻辑说明：knn_graph先用scikit-learn的NearestNeighbors计算每个节点在特征空间中的k个最近邻居，返回的knn_indices形状是[N, k]，第i行是节点i的邻居编号列表。用[:, 1:]去掉自身是因为kneighbors返回的第一个最近邻永远是节点自己。最终输出PyTorch Geometric要求的COO格式edge_index，shape为[2, M]，第一行是源节点，第二行是目标节点。
参数调整时注意两点。第一，k值太小会导致图稀疏，消息传递不充分；k值太大会把不同故障类别的样本强行连边，造成特征混淆。CWRU数据集上k=5~8是经验区间，跨数据集时先看t-SNE可视化再定。第二，metric选cosine还是euclidean取决于特征是否做过标准化——如果特征已经z-score归一化，euclidean和cosine效果接近；如果没有，cosine对幅值不敏感，更适合振动信号这种幅值随工况变化的场景。
2.3 CWRURadius.py：固定阈值的构图策略
import torch def radius_graph(features, r=0.5): """ features: [N, D] r: 距离阈值, 超过该距离不连边 """ n = features.shape[0] src = [] dst = [] for i in range(n): for j in range(i + 1, n): dist = torch.norm(features[i] - features[j], p=2) if dist < r: src.extend([i, j]) dst.extend([j, i]) edge_index = torch.tensor([src, dst], dtype=torch.long) return edge_index逻辑说明：RadiusGraph遍历所有节点对，计算欧氏距离，小于阈值r就连一条无向边。嵌套循环的时间复杂度是O(N²)，N是节点数——当节点数超过5000时，这个函数会成为瓶颈，建议改用PyTorch Geometric内置的radius_graph（用C++实现，批量计算距离矩阵）。
r的选择比k更敏感。r太小会得到大量孤立点，GNN的消息传递在这些节点上失效；r太大会让图变成稠密团，和原始特征没有区别。我的做法是先算全样本距离分布，取所有距离的P50到P70分位数作为r的初值，再在验证集上微调。
CWRU数据集的采样频率是12kHz和48kHz两档，不同转速下的振动幅值差异很大。如果用RadiusGraph，建议先对每个窗口做幅值归一化再算距离，否则高转速样本会天然聚集——模型学到的是转速，不是故障。
2.4 CWRUPathGraph.py：当物理路径已知时怎么构图
PathGraph适用于测点位置明确的场景。比如CWRU数据集里，驱动端和风扇端各有一个加速度传感器，信号从故障源传到传感器经过的路径长度不同，PathGraph按物理距离或传播路径长度连边。
def path_graph(sensor_positions, max_path_len=2): """ sensor_positions: {sensor_name: [x, y, z]} max_path_len: 最长路径长度, 超过该距离不直接连边 """ import itertools sensors = list(sensor_positions.keys()) src, dst = [], [] for s1, s2 in itertools.combinations(sensors, 2): dist = euclidean(sensor_positions[s1], sensor_positions[s2]) if dist <= max_path_len: src.extend([sensors.index(s1), sensors.index(s2)]) dst.extend([sensors.index(s2), sensors.index(s1)]) return torch.tensor([src, dst], dtype=torch.long)逻辑说明：itertools.combinations生成所有传感器对，计算空间欧氏距离（如果传感器安装在同一平面则只需要x、y坐标）。把传感器索引作为图节点，传感器位置坐标作为节点特征的一部分。PathGraph的边数远小于KNNGraph和RadiusGraph，计算开销最小，但要求你对设备结构足够了解——传感器位置标错，全链路效果直接崩塌。
三种构图方式的选择原则：KNN适合特征空间聚集性强的数据；Radius适合分布均匀的数据；PathGraph适合有明确物理拓扑的数据。实际调试时我通常先跑KNNGraph拿到baseline，再用PathGraph作为论文的对比实验。
3. Generator.py与数据包装：从原始信号到PyG的Data对象
3.1 为什么不能直接拿numpy数组训练
PyTorch Geometric的模型接受的输入是torch_geometric.data.Data对象，它至少包含x（节点特征矩阵）、edge_index（边）、y（标签三个字段）。如果做图分类，还需要batch字段来区分多个图实例。Generator.py解决的就是这个问题——把原始振动信号切成窗口、提特征、构图，然后打包成Data对象。
import torch from torch_geometric.data import Data class VibrationDataGenerator: def __init__(self, window_size=1024, stride=512, feature_extractor=None): self.window_size = window_size self.stride = stride self.feature_extractor = feature_extractor def sliding_window(self, signal): """将一维信号切成重叠窗口""" windows = [] n = len(signal) for start in range(0, n - self.window_size + 1, self.stride): windows.append(signal[start:start + self.window_size]) return np.array(windows) def to_graph_data(self, signal, label): windows = self.sliding_window(signal) if self.feature_extractor: # 每个窗口提取统计特征 features = [self.feature_extractor(w) for w in windows] x = torch.tensor(features, dtype=torch.float) else: # 直接用窗口的原始波形作为节点特征 x = torch.tensor(windows, dtype=torch.float).unsqueeze(-1) # 在特征空间构图 edge_index = knn_graph(x.numpy(), k=self.k) return Data(x=x, edge_index=edge_index, y=torch.tensor([label], dtype=torch.long))逻辑说明：to_graph_data先做滑动窗口切分，默认窗口大小1024，步长512，重叠率50%。每个窗口要么直接用原始波形（维度是[1024]），要么提取统计特征（均值、方差、峰值因子、峭度等）。然后调用前面的knn_graph生成edge_index。对图分类任务，整个窗口集合构成一个图，label是窗口对应的故障类别。
滑动窗口参数需要根据信号特性调整。CWRU数据的轴承转速在1730~1797rpm之间，转频约30Hz，1024个采样点覆盖约21个旋转周期，足够捕捉一个完整的故障冲击周期。如果转速更高，窗口要相应缩短；如果采样频率不同，窗口尺寸不能直接搬用。
3.2 构建多图Batch：torch_geometric的collate逻辑
训练时通常需要把多个Data对象合成一个batch——本质是沿节点维度拼接所有图，edge_index的第二维加上节点偏移量。
from torch_geometric.loader import DataLoader list_of_graphs = [] for signal, label in zip(all_signals, all_labels): graph = vg.to_graph_data(signal, label) list_of_graphs.append(graph) loader = DataLoader(list_of_graphs, batch_size=32, shuffle=True) for batch in loader: print(batch) # batch.edge_index会自动完成节点索引偏移DataLoader内部实现了Batch.from_data_list，返回的batch对象包含batch属性，表示每个节点属于哪个图。GNN模型里的global_mean_pool或global_max_pool层就是靠它把节点级别的输出汇总成图级别的预测。
4. GNN模型与训练循环的设计
4.1 图卷积层怎么选：GCN、GraphSAGE还是GAT
框架里train_graph_utils.py封装了两类任务：节点分类用于识别传感器或子段的故障状态，图分类用于判断整个样本的故障类型。这两种任务的模型骨干可以共用，区别只在最后的读出层。
import torch.nn as nn from torch_geometric.nn import GCNConv, SAGEConv, GATConv, global_mean_pool class FaultGNN(nn.Module): def __init__(self, in_channels, hidden_channels=64, out_channels=10, conv_type='gcn', num_layers=3): super().__init__() self.conv_type = conv_type if conv_type == 'gcn': conv_layer = GCNConv elif conv_type == 'sage': conv_layer = SAGEConv elif conv_type == 'gat': conv_layer = GATConv else: raise ValueError(f"Unsupported conv_type: {conv_type}") assert num_layers >= 2 self.convs = nn.ModuleList() self.convs.append(conv_layer(in_channels, hidden_channels)) for _ in range(num_layers - 2): self.convs.append(conv_layer(hidden_channels, hidden_channels)) self.convs.append(conv_layer(hidden_channels, out_channels)) # 用于聚合成图级别表示的层 self.pool = global_mean_pool def forward(self, x, edge_index, batch=None): for i, conv in enumerate(self.convs): x = conv(x, edge_index) x = nn.functional.relu(x) x = nn.functional.dropout(x, p=0.2, training=self.training) if batch is not None: # 图分类: 将所有节点表示聚合成一个图表示 x = self.pool(x, batch) return nn.functional.log_softmax(x, dim=-1)FaultGNN的forward支持两种模式。batch=None时走节点分类，输出每个节点的类别分布；batch参数传入时，global_mean_pool把同一图的所有节点表示取平均，得到图级别表示，接log_softmax做多分类。conv_type参数控制底层卷积算子，切换GCN/GraphSAGE/GAT只需要改一个字符串，方便做消融实验。
选择建议：GCN是baseline首选，参数少、收敛快；GraphSAGE擅长处理大图和inductive场景（测试时有新节点出现）；GAT的多头注意力机制能学习边的权重，在传感器相关性随工况变化的任务中表现更稳。
4.2 训练循环：监控什么指标才有意义
from torch_geometric.loader import DataLoader import mlflow # 用于实验跟踪 def train_model(model, train_loader, val_loader, epochs=200, lr=0.001, device='cuda'): model = model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) criterion = nn.CrossEntropyLoss() for epoch in range(epochs): model.train() total_loss = 0 for batch in train_loader: batch = batch.to(device) optimizer.zero_grad() if batch.y.dim() == 1: # 图分类任务 out = model(batch.x, batch.edge_index, batch.batch) else: # 节点分类任务 out = model(batch.x, batch.edge_index) loss = criterion(out, batch.y) loss.backward() optimizer.step() total_loss += loss.item() # 验证 model.eval() val_acc = evaluate(model, val_loader, device) mlflow.log_metrics({'train_loss': total_loss / len(train_loader), 'val_acc': val_acc}, step=epoch) scheduler.step()训练技巧：这个框架使用了CosineAnnealing学习率调度，优点是前期快速下降到局部最优附近，后期用小学习率精细搜索，配合Adam比固定学习率稳定很多。weight_decay设为5e-4是为了抑制过平滑（over-smoothing）——GNN层数加深后节点表示趋于一致，正则化能延缓这个过程。
验证集指标别只看准确率。故障诊断场景下类别不平衡很常见（正常样本远多于故障样本），要同时监控macro-F1和混淆矩阵。如果某个故障类别的recall偏低，先检查这个类别的样本数是不是太少，再考虑加class weights或Focal Loss。
5. 实战调优：从CWRU换到SEU、PU数据集时改什么
5.1 跨数据集需要动的三处地方
SEURadius.py和PURadius.py的存在说明这个框架设计时考虑了跨数据集复用。但直接换数据集训练GNN通常会遇到三个问题：采样频率不同导致有效频率bandwidth不同，传感器布点不同导致图的拓扑结构不同，故障类型和标签体系不同。
处理方式是统一信号预处理：首先对原始振动信号做带通滤波——CWRU是12kHz采样，SEU是20kHz采样，滤波截止频率按采样率比例折算，保留2kHz到5kHz之间的故障冲击频段。然后做包络谱分析，取包络谱的前若干阶特征频率幅值作为节点特征，而不是直接用原始波形。这样可以部分消除采样率差异带来的特征分布偏移。
跨数据集训练时的节点特征设计要更通用，统计特征（峭度、峰值因子、波形因子）比原始波形更稳。
5.2 四个容易踩的坑
第一个坑是孤立点问题。RadiusGraph在特征空间密度不均匀时会产生大量孤立节点，这些节点的embedding经过消息传递后仍然是初始值，相当于噪声。解决办法是加一个self_loop标志，确保每个节点至少有一条自环边。
from torch_geometric.utils import add_self_loops edge_index, _ = add_self_loops(edge_index)self-loop只保证消息传递不会把孤立节点漏掉，不能解决图质量本身的问题。如果超过10%的节点是孤立的，说明构图参数需要重新调。
第二个坑是过平滑。GNN层数超过4层后，所有节点的表示趋同，准确率断崖式下跌。CWRU这种小规模数据集（几百个有效样本），2~3层GCN就足够了。训练时如果发现加深网络后验证集准确率不再提升甚至下降，优先怀疑过平滑，加残差连接或JK-Net结构。
第三个坑是标签泄露。如果对整段信号提取特征再切窗口，前后窗口共享了全局统计量，会引入测试集信息。正确做法是先切窗口，再在每个窗口内部提取特征。Generator.py里的顺序是先sliding_window再feature_extractor，实现上已经规避了这个风险，但如果自己改写特征提取逻辑，务必保持这个顺序。
验证是否泄露的方法是：把训练集随机打乱，如果准确率下降很多，说明模型依赖了样本顺序相关的信息，检查特征提取是否有跨样本的统计操作。
第四个坑是batch_size对GNN的影响。和CNN不同，DataLoader的batch表现是拼接N个图，一个batch的节点总数等于所有图节点数之和。batch_size设得过大，GPU显存会突然爆掉，这是因为邻接矩阵的大小随节点数平方增长。按经验，batch_size=16、每图100个节点左右是30系显卡的舒适区。必要时用NeighborSampler做mini-batch采样。
5.3 部署时的推理加速技巧
模型训练完要落地的话，有一个容易被忽略的点：PyTorch Geometric的Data对象转换成ONNX时，edge_index的动态shape会阻塞导出。常见的做法是固定图大小——把所有图的节点数pad到同一数值K，不足的用mask机制屏蔽。
def pad_graph(data, max_nodes): n = data.x.shape[0] if n < max_nodes: pad_len = max_nodes - n data.x = torch.cat([data.x, torch.zeros(pad_len, data.x.shape[1])]) data.mask = torch.cat([torch.ones(n), torch.zeros(pad_len)]).bool() return data再加上adjacency矩阵预先计算好，推理时直接查表，而不是每次重新构图。CWRU的推理延迟可以压到毫秒级，满足产线实时诊断的基本要求。
故障诊断工程的验收标准是召回率优先还是准确率优先，取决于误报和漏报的成本：风电叶片卯焊开裂漏报一次可能造成整机报废，代价远大于误报停机。GNN的图结构天然支持在节点级别输出置信度，这比CNN的黑盒全连接层更容易做边界样本的兜底判断——你可以设定一个阈值，当某个传感器节点的softmax置信度低于0.7时，把该节点的预测标记为“待人工复核”，而不是直接归入某个故障类别。这个能力是从图结构里白捡的，CNN做不了。
本文还有配套的精品资源，点击获取
