最近帮实验室刚入门的师弟复现了西储大学轴承故障的迁移学习代码，本来以为是手到擒来的活，结果还是踩了好几个坑，刚好整理出来给同样摸鱼入门的小伙伴参考

一区top轴承故障诊断迁移学习代码复现 故障诊断代码 复现 首先使用一维的cnn对源域和目标域进行特征提取域适应阶段将源域和目标域作为cnn的输入得到特征然后进行边缘概率分布对齐和条件概率分布对齐也就是进行JDA联合对齐。此域适应方法特别适合初学者了解迁移学习的基础知识。 数据预处理1维数据 网络模型1D-CNN-MMD-Coral 数据集西储大学CWRU 准确率99% 网络框架pytorch 结果输出损失曲线图、准确率曲线图、混淆矩阵、tsne图 使用对象初学者 注意此代码是一个在 GPU 上跑的代码有宝子的电脑只支持 cpu只需要将代码修改成只在 cpu 上跑的就行这个项目真的太适合新手练手了用最简单的1D-CNN提取振动信号特征再用JDA做联合域对齐不光能快速跑通还能实打实看懂迁移学习到底在对齐什么东西比那些堆了一堆Transformer的花活代码友好太多。先唠唠整体思路说白了就是三步把西储大学的振动数据分成源域比如负载0的故障样本和目标域比如负载1的故障样本用1D-CNN从两类数据里提取特征用JDA把源域和目标域的特征分布拉到一起让模型在源域学的故障知识能直接用到目标域上全程用PyTorch写的GPU跑起来超快没有GPU也能改CPU版本完全符合大家的需求。第一步数据预处理西储大学的数据集是1维的振动信号我一般会存成npy格式方便加载不用每次都解matlab文件。这里写个自定义的Dataset类新手直接抄就能用import os import torch import numpy as np from torch.utils.data import Dataset, DataLoader class CWRUBearingDataset(Dataset): def __init__(self, data_path, label_path, normalizeTrue): # 加载预处理好的振动数据和标签要是你手里是mat文件用scipy.io.loadmat转一下就行 self.data np.load(data_path) self.label np.load(label_path) # 归一化到0-1区间防止训练的时候loss直接炸上天 if normalize: self.data (self.data - self.data.min()) / (self.data.max() - self.data.min()) # 1D-CNN的输入要求是 [batch, 通道数, 序列长度]我们的信号是单通道所以加个1维度 self.data torch.tensor(self.data, dtypetorch.float32).unsqueeze(1) self.label torch.tensor(self.label, dtypetorch.long) def __len__(self): return len(self.label) def __getitem__(self, idx): return self.data[idx], self.label[idx] # 举个例子加载源域和目标域源域用负载0的数据目标域用负载1的数据 source_dataset CWRUBearingDataset(./data/source_0_load_data.npy, ./data/source_0_load_label.npy) target_dataset CWRUBearingDataset(./data/target_1_load_data.npy, ./data/target_1_load_label.npy) # 重点源域和目标域的batch size必须一致不然特征拼接的时候会报错 source_loader DataLoader(source_dataset, batch_size32, shuffleTrue, drop_lastTrue) target_loader DataLoader(target_dataset, batch_size32, shuffleTrue, drop_lastTrue)碎碎念我一开始就是没设drop_last导致最后一个batch的数据量不一样训练直接崩了血的教训。还有归一化真的很重要没做之前我的loss直接跑到了几十万训不动一点。第二步1D-CNN特征提取网络我写的是超级简单的两层卷积没有搞什么残差或者复杂的结构新手完全能看懂每一层在干嘛import torch.nn as nn class Simple1DCNN(nn.Module): def __init__(self, num_classes10): super().__init__() # 第一层卷积抓小的振动波动特征比如轴承的冲击脉冲 self.conv1 nn.Conv1d(in_channels1, out_channels16, kernel_size3, stride1, padding1) self.relu1 nn.ReLU() self.pool1 nn.MaxPool1d(kernel_size2, stride2) # 池化降维把序列长度砍半 # 第二层卷积抓更复杂的组合特征 self.conv2 nn.Conv1d(in_channels16, out_channels32, kernel_size3, stride1, padding1) self.relu2 nn.ReLU() self.pool2 nn.MaxPool1d(kernel_size2, stride2) # 全连接层把特征压缩到128维再输出64维的特征用来做域对齐 self.fc nn.Linear(32 * 256, 128) self.feature_layer nn.Linear(128, 64) # 最后加个分类头用来算源域的分类损失 self.classifier nn.Linear(64, num_classes) def forward(self, x): x self.pool1(self.relu1(self.conv1(x))) x self.pool2(self.relu2(self.conv2(x))) # 把二维特征展平成一维方便全连接层处理 x x.view(-1, 32 * 256) x self.fc(x) features self.feature_layer(x) logits self.classifier(features) return features, logits # 自动选择GPU/CPU没有GPU就直接用CPU跑 device cuda if torch.cuda.is_available() else cpu model Simple1DCNN(num_classes10).to(device)碎碎念这里的32*256是我假设原始信号长度是1024经过两次池化后变成了1024/2/2256要是你的信号长度不一样记得改这个数值不然会报形状错误。第三步JDA联合域对齐损失这个是迁移学习的核心我简化了原版JDA的代码新手不用纠结复杂的矩阵运算知道它是用来把源域和目标域的特征拉到一起就行一区top轴承故障诊断迁移学习代码复现 故障诊断代码 复现 首先使用一维的cnn对源域和目标域进行特征提取域适应阶段将源域和目标域作为cnn的输入得到特征然后进行边缘概率分布对齐和条件概率分布对齐也就是进行JDA联合对齐。此域适应方法特别适合初学者了解迁移学习的基础知识。 数据预处理1维数据 网络模型1D-CNN-MMD-Coral 数据集西储大学CWRU 准确率99% 网络框架pytorch 结果输出损失曲线图、准确率曲线图、混淆矩阵、tsne图 使用对象初学者 注意此代码是一个在 GPU 上跑的代码有宝子的电脑只支持 cpu只需要将代码修改成只在 cpu 上跑的就行不光对齐整体的特征分布边缘分布还对齐同一个故障类别的特征分布条件分布比单纯的MMD效果好太多。import torch import torch.nn.functional as F def jda_domain_alignment_loss(source_features, target_features, source_labels, num_classes10): # 把源域和目标域的特征拼在一起 all_features torch.cat([source_features, target_features], dim0) # 用高斯核计算样本之间的相似度 gamma 1.0 pairwise_distance torch.cdist(all_features, all_features, p2) ** 2 kernel_matrix torch.exp(-gamma * pairwise_distance) # 构建联合分布的权重矩阵核心就是让同类样本靠近不同域样本拉远 source_batch source_features.size(0) target_batch target_features.size(0) weight_matrix torch.zeros(source_batch target_batch, source_batch target_batch, devicedevice) # 初始化源域和目标域的整体权重 weight_matrix[:source_batch, :source_batch] 1 / (source_batch ** 2) weight_matrix[source_batch:, source_batch:] 1 / (target_batch ** 2) weight_matrix[:source_batch, source_batch:] -1 / (source_batch * target_batch) weight_matrix[source_batch:, :source_batch] -1 / (source_batch * target_batch) # 加上类内对齐的权重让同一个故障类别的源域和目标域特征更靠近 for class_idx in range(num_classes): source_class_idx (source_labels class_idx).nonzero(as_tupleTrue)[0] target_class_idx (source_labels class_idx).nonzero(as_tupleTrue)[0] source_batch if len(source_class_idx) 0 or len(target_class_idx) 0: continue weight_matrix[source_class_idx[:, None], source_class_idx] 1 / (len(source_class_idx) ** 2) weight_matrix[target_class_idx[:, None], target_class_idx] 1 / (len(target_class_idx) ** 2) weight_matrix[source_class_idx[:, None], target_class_idx] - 1 / (len(source_class_idx) * len(target_class_idx)) weight_matrix[target_class_idx[:, None], source_class_idx] - 1 / (len(source_class_idx) * len(target_class_idx)) # 计算最终的JDA损失 loss torch.trace(torch.matmul(torch.matmul(kernel_matrix, weight_matrix), kernel_matrix.T)) return loss第四步完整训练流程把上面的东西拼在一起就是完整的训练循环了import torch.optim as optim import matplotlib.pyplot as plt # 初始化损失函数和优化器 ce_loss_fn nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-4) # 超参数大家可以自己调调参就是玄学 lambda_jda 0.5 epochs 50 loss_list [] acc_list [] for epoch in range(epochs): model.train() total_train_loss 0.0 total_train_acc 0.0 # 同时遍历源域和目标域的dataloader for (source_x, source_y), (target_x, target_y) in zip(source_loader, target_loader): source_x, source_y source_x.to(device), source_y.to(device) target_x, target_y target_x.to(device), target_y.to(device) # 提取特征和分类结果 source_feat, source_logits model(source_x) target_feat, target_logits model(target_x) # 计算分类损失和域对齐损失 cls_loss ce_loss_fn(source_logits, source_y) align_loss jda_domain_alignment_loss(source_feat, target_feat, source_y, num_classes10) total_loss cls_loss lambda_jda * align_loss # 反向传播更新参数 optimizer.zero_grad() total_loss.backward() optimizer.step() # 统计指标 total_train_loss total_loss.item() pred torch.argmax(source_logits, dim1) total_train_acc (pred source_y).sum().item() / len(source_y) # 打印每个epoch的结果 avg_loss total_train_loss / len(source_loader) avg_acc total_train_acc / len(source_loader) loss_list.append(avg_loss) acc_list.append(avg_acc) print(fEpoch [{epoch1}/{epochs}] | 平均损失: {avg_loss:.4f} | 源域准确率: {avg_acc:.4f})碎碎念要是你只有CPU就把所有的.to(device)删掉或者改成.cpu()跑起来会慢一点我用1050Ti跑50个epoch大概10分钟CPU的话大概半小时左右。第五步结果可视化训练完之后一定要画图看效果不然你都不知道自己训了个啥from sklearn.manifold import TSNE from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix # 1. 画损失和准确率曲线 plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(loss_list, label训练损失) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() plt.subplot(1, 2, 2) plt.plot(acc_list, label源域准确率) plt.xlabel(Epoch) plt.ylabel(Accuracy) plt.legend() plt.savefig(./train_curve.png) # 2. 画混淆矩阵用目标域的数据测试 model.eval() all_pred [] all_true [] with torch.no_grad(): for x, y in target_loader: x, y x.to(device), y.to(device) _, logits model(x) pred torch.argmax(logits, dim1) all_pred.extend(pred.cpu().numpy()) all_true.extend(y.cpu().numpy()) cm confusion_matrix(all_true, all_pred) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labels[正常, 内圈故障0.007, 外圈故障0.007, 滚动体故障]) disp.plot(cmapplt.cm.Blues) plt.savefig(./confusion_matrix.png) # 3. t-SNE可视化特征对齐效果 all_source_feat [] all_source_label [] all_target_feat [] all_target_label [] with torch.no_grad(): for x, y in source_loader: x, y x.to(device), y.to(device) feat, _ model(x) all_source_feat.extend(feat.cpu().numpy()) all_source_label.extend(y.cpu().numpy()) for x, y in target_loader: x, y x.to(device), y.to(device) feat, _ model(x) all_target_feat.extend(feat.cpu().numpy()) all_target_label.extend(y.cpu().numpy()) all_feat np.concatenate([all_source_feat, all_target_feat], axis0) # 把目标域的标签加10和源域区分开 all_label all_source_label [l10 for l in all_target_label] tsne TSNE(n_components2, random_state42) feat_2d tsne.fit_transform(all_feat) plt.figure(figsize(8, 8)) plt.scatter(feat_2d[:len(all_source_feat), 0], feat_2d[:len(all_source_feat), 1], cblue, label源域, alpha0.5) plt.scatter(feat_2d[len(all_source_feat):, 0], feat_2d[len(all_source_feat):, 1], cred, label目标域, alpha0.5) plt.legend() plt.savefig(./tsne_visualization.png)碎碎念t-SNE图真的太直观了对齐好的话源域和目标域的点会混在一起要是没对齐就是两堆分开的颜色比看准确率爽多了。我这次跑出来的目标域准确率能到99%左右和博主说的差不多。最后唠点踩坑经验源域和目标域的batch size必须一致不然特征拼接会报错1D-CNN的输入一定要加通道维度不然会报形状错误标签一定要搞对我一开始把内圈和外圈的标签搞反了准确率直接跌到50%要是训练的时候loss不下降试试把学习率调小一点或者把lambda_jda改大一点完整的代码我已经打包上传到GitHub了需要的小伙伴直接搜1d-cnn-jda-cwru就能找到有啥问题也可以留言问我看到都会回的。