DeepBear-Health：基于迁移学习和可解释时频分析的滚动轴承故障诊断（Python）

Ne0inhk

28 Feb 2026 — 16 min read

首先，利用短时傅里叶变换将原始振动信号转换为时频图像，并通过动态阈值增强故障特征；然后，采用预训练的VGG16网络进行迁移学习，在保留通用视觉特征的基础上微调分类层，实现对10种健康状态的高精度识别。

为了增强模型的可信度，算法引入Grad-CAM热力图可视化模型的决策依据，并设计Faithfulness指标量化解释的忠实性。此外，结合故障机理知识和大语言模型，自动生成符合专家习惯的诊断报告，并在检测到严重故障时通过邮件实时告警。该方法在CWRU轴承数据集上表现出优异的分类性能，同时为工业场景下的智能运维提供了透明、可解释的决策支持。

算法步骤

数据加载与类别映射遍历指定目录下的所有.mat文件，根据文件名中的关键词（如'Normal'、'B007'、'IR014'、'OR021'等）将样本分为10类：正常、三种尺寸的滚动体故障、三种尺寸的内圈故障、三种尺寸的外圈故障。提取每个文件中的驱动端振动信号（DE_time）并存储。

信号预处理

下采样：将原始采样率（12000 Hz）降为6000 Hz，减少数据量并保留故障特征频段。

滑动窗口分割：对每段下采样后的长信号，以1024点长度、512点步长切分为多个短样本，增加训练样本数量。

时频变换（STFT）：对每个短样本计算短时傅里叶变换，得到时频谱图，窗长128点、重叠64点，兼顾时间和频率分辨率。

图像增强：对幅值谱取dB值（对数变换），然后通过动态阈值裁剪底部35%噪声，增强故障冲击成分的对比度。

归一化与尺寸调整：将谱图线性归一化至[0,1]，再缩放到224×224像素，最后堆叠为三通道灰度图，满足ImageNet预训练模型的输入要求。

数据集划分将生成的时频图像及对应标签按照类别分层抽样，划分为训练集（70%）、验证集（15%）和测试集（15%），确保各类别分布均衡。

迁移学习模型构建

加载在ImageNet上预训练的VGG16模型，去除顶部分类层，保留卷积基。

冻结所有卷积层参数，避免微调初期破坏通用特征。

在卷积基后添加全局平均池化层、256维全连接层（ReLU激活）、Dropout（0.1）以及10维Softmax输出层。

使用Adam优化器（学习率0.001）、分类交叉熵损失函数编译模型。

模型训练与验证使用训练集训练模型25个轮次，每批32个样本，并在验证集上监控准确率和损失，保存最佳模型。训练过程中保持卷积基冻结，仅更新新增的全连接层参数。

模型评估

在测试集上计算整体分类准确率。

输出每个类别的精确率、召回率、F1-score，生成分类报告。

绘制混淆矩阵、多类别ROC曲线及AUC值，直观评估各类别的区分能力。

可解释性分析

Grad-CAM可视化：针对测试样本，计算最后一个卷积层（block5_conv3）的梯度，生成热力图，高亮模型决策时关注的时频区域。

Faithfulness指标计算：随机选取50个高置信度样本，删除热力图中最关键的20%像素（置零），重新预测并计算置信度下降百分比，量化解释的忠实度。

专家级报告生成结合Grad-CAM提取的显著频率、模型预测的故障类别及预定义的物理知识（如不同故障对应的特征频率），构建结构化提示，调用本地大语言模型（LLaMA 3.2）生成包含诊断结论、严重等级、维护建议的文本报告。

自动告警通知当诊断出严重故障（如尺寸021）或模型置信度极高时，通过SendGrid API将生成的报告以邮件形式发送给指定维护人员，实现实时预警。

# ==================== 1. 导入所需库 ==================== import os import numpy as np import scipy.io from scipy import signal import matplotlib.pyplot as plt import cv2 from sklearn.model_selection import train_test_split from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc import seaborn as sns import tensorflow as tf from tensorflow.keras.applications import VGG16 from tensorflow.keras.models import Model from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, Dropout from tensorflow.keras.optimizers import Adam from tensorflow.keras.utils import to_categorical import copy from tqdm import tqdm import ollama # 用于调用本地大语言模型生成报告 from datetime import datetime import sendgrid # 用于发送邮件通知 from sendgrid.helpers.mail import Mail, Email, To, Content # 设置随机种子以保证结果可重复 np.random.seed(42) tf.random.set_seed(42) # ==================== 2. 数据加载与标签映射 ==================== data_dir = '../data' # 数据存放目录，请根据实际情况修改 file_list = [f for f in os.listdir(data_dir) if f.endswith('.mat')] raw_signals = [] # 原始振动信号列表 labels = [] # 对应的标签（整数 0-9） print(f"找到 {len(file_list)} 个 .mat 文件") # 定义标签映射字典：根据文件名中的关键词确定类别 def get_label_from_filename(filename): if 'Normal' in filename or 'Time' in filename: return 0 elif 'B007' in filename: return 1 elif 'B014' in filename: return 2 elif 'B021' in filename: return 3 elif 'IR007' in filename: return 4 elif 'IR014' in filename: return 5 elif 'IR021' in filename: return 6 elif 'OR007' in filename: return 7 elif 'OR014' in filename: return 8 elif 'OR021' in filename: return 9 else: return None # 未知类别 for file_name in file_list: label = get_label_from_filename(file_name) if label is None: print(f"跳过未定义类别的文件: {file_name}") continue # 加载 .mat 文件 file_path = os.path.join(data_dir, file_name) mat_data = scipy.io.loadmat(file_path) # 查找驱动端振动信号变量（通常为 DE_time） signal_key = [key for key in mat_data.keys() if 'DE_time' in key] if len(signal_key) == 0: print(f"跳过 {file_name} (未找到 DE_time 变量)") continue sig = mat_data[signal_key[0]].flatten() # 展平为一维数组 raw_signals.append(sig) labels.append(label) print(f"已加载 {file_name} -> 类别 {label}") print(f"\n数据加载完成。共 {len(labels)} 个样本") # ==================== 3. 信号预处理：下采样 + STFT 生成时频图像 ==================== ORIGINAL_FS = 12000 # 原始采样频率 (Hz)，CWRU 驱动端数据通常为 12k target_fs = 6000 # 目标采样频率 (Hz)，下采样至 6k 以减少计算量 # 3.1 下采样 downsampled_signals = [] for sig in raw_signals: num_samples = int(len(sig) * (target_fs / ORIGINAL_FS)) resampled = signal.resample(sig, num_samples) downsampled_signals.append(resampled) print(f"信号已下采样至 {target_fs/ORIGINAL_FS*100:.1f}%") # 3.2 滑动窗口参数 segment_length = 1024 # 每个片段的长度（点数） step_size = 512 # 滑动步长（点数） target_image_size = (224, 224) # VGG16 要求的输入尺寸 nperseg = 128 # STFT 每段点数 noverlap = 64 # STFT 重叠点数（50% 重叠） stft_images = [] # 存储生成的 RGB 图像 stft_labels = [] # 对应的标签 print("开始生成 STFT 时频图像...") for idx, sig in enumerate(downsampled_signals): current_label = labels[idx] # 对每段长信号进行滑动窗口切片 for i in range(0, len(sig) - segment_length, step_size): segment = sig[i : i + segment_length] # 计算 STFT f, t, Zxx = signal.stft(segment, fs=target_fs, nperseg=nperseg, noverlap=noverlap) magnitude = np.abs(Zxx) # 幅值谱 magnitude_db = 20 * np.log10(magnitude + 1e-6) # 转换为 dB 单位，增强对比度 # 动态阈值：抑制底部 35% 的噪声区域 range_db = magnitude_db.max() - magnitude_db.min() threshold = magnitude_db.min() + 0.35 * range_db magnitude_db[magnitude_db < threshold] = threshold # 归一化到 [0,1] mag_norm = (magnitude_db - magnitude_db.min()) / (magnitude_db.max() - magnitude_db.min()) # 缩放到 VGG16 输入尺寸 (224,224) resized = cv2.resize(mag_norm, target_image_size) # 堆叠为三通道灰度图（满足 ImageNet 预训练模型输入要求） img_rgb = np.stack((resized,) * 3, axis=-1) stft_images.append(img_rgb) stft_labels.append(current_label) X = np.array(stft_images) # 特征数组，形状 (N, 224, 224, 3) y = np.array(stft_labels) # 标签数组 print(f"特征提取完成。输入形状: {X.shape}") # ==================== 4. 数据集划分 ==================== # 将标签转换为 one-hot 编码 y_encoded = to_categorical(y, num_classes=10) # 先划分 70% 训练，30% 临时（再平分验证和测试） X_train, X_temp, y_train, y_temp = train_test_split( X, y_encoded, test_size=0.3, random_state=42, stratify=y ) # 从临时集中划分 15% 验证，15% 测试 stratify_labels = np.argmax(y_temp, axis=1) # 用于分层抽样 X_val, X_test, y_val, y_test = train_test_split( X_temp, y_temp, test_size=0.5, random_state=42, stratify=stratify_labels ) print(f"训练集形状: {X_train.shape}") print(f"验证集形状: {X_val.shape}") print(f"测试集形状: {X_test.shape}") # ==================== 5. 构建迁移学习模型 (VGG16) ==================== # 加载预训练的 VGG16（不含顶层） base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3)) # 冻结所有卷积层（初期不更新这些权重） for layer in base_model.layers: layer.trainable = False # 添加自定义分类头 x = base_model.output x = GlobalAveragePooling2D()(x) # 全局平均池化，减少参数量 x = Dense(256, activation='relu')(x) x = Dropout(0.1)(x) predictions = Dense(10, activation='softmax')(x) # 10 个输出类别 model = Model(inputs=base_model.input, outputs=predictions) # 编译模型 model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy']) model.summary() # ==================== 6. 模型训练 ==================== history = model.fit( X_train, y_train, epochs=25, batch_size=32, validation_data=(X_val, y_val), verbose=1 ) # ==================== 7. 模型评估 ==================== # 绘制训练曲线 plt.figure(figsize=(12, 4)) plt.subplot(1, 2, 1) plt.plot(history.history['accuracy'], label='训练集') plt.plot(history.history['val_accuracy'], label='验证集') plt.title('准确率') plt.legend() plt.subplot(1, 2, 2) plt.plot(history.history['loss'], label='训练集') plt.plot(history.history['val_loss'], label='验证集') plt.title('损失') plt.legend() plt.show() # 在测试集上评估 y_pred_probs = model.predict(X_test) y_pred_classes = np.argmax(y_pred_probs, axis=1) y_true_classes = np.argmax(y_test, axis=1) # 类别名称 target_names = [ 'Normal', 'Ball_007', 'Ball_014', 'Ball_021', 'IR_007', 'IR_014', 'IR_021', 'OR_007', 'OR_014', 'OR_021' ] print("分类报告:") print(classification_report(y_true_classes, y_pred_classes, target_names=target_names)) # 混淆矩阵 cm = confusion_matrix(y_true_classes, y_pred_classes) plt.figure(figsize=(12, 10)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=target_names, yticklabels=target_names) plt.title('混淆矩阵') plt.ylabel('真实类别') plt.xlabel('预测类别') plt.xticks(rotation=45, ha='right') plt.tight_layout() plt.show() # ROC 曲线 fpr = dict() tpr = dict() roc_auc = dict() n_classes = 10 for i in range(n_classes): fpr[i], tpr[i], _ = roc_curve(y_test[:, i], y_pred_probs[:, i]) roc_auc[i] = auc(fpr[i], tpr[i]) plt.figure(figsize=(12, 8)) colors = plt.cm.tab10(np.linspace(0, 1, 10)) for i, color in zip(range(n_classes), colors): plt.plot(fpr[i], tpr[i], color=color, lw=2, label=f'{target_names[i]} (AUC = {roc_auc[i]:0.2f})') plt.plot([0, 1], [0, 1], 'k--', lw=2) plt.xlim([0.0, 1.0]) plt.ylim([0.0, 1.05]) plt.xlabel('假阳性率') plt.ylabel('真阳性率') plt.title('多类别 ROC 曲线') plt.legend(loc="lower right") plt.show() # ==================== 8. 可解释性分析 (Grad-CAM) ==================== def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None): """ 计算 Grad-CAM 热力图 img_array: 输入图像 (已扩展 batch 维度，形状 (1,224,224,3)) model: Keras 模型 last_conv_layer_name: 最后一层卷积层的名称 pred_index: 要解释的类别索引，若为 None 则使用预测概率最高的类别 """ # 构建一个从输入到目标卷积层输出及模型输出的新模型 grad_model = tf.keras.models.Model( [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output] ) with tf.GradientTape() as tape: last_conv_output, preds = grad_model(img_array) if pred_index is None: pred_index = tf.argmax(preds[0]) class_channel = preds[:, pred_index] # 计算梯度 grads = tape.gradient(class_channel, last_conv_output) pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2)) # 全局平均池化 # 加权求和 last_conv_output = last_conv_output[0] heatmap = last_conv_output @ pooled_grads[..., tf.newaxis] heatmap = tf.squeeze(heatmap) # 归一化到 [0,1] heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap) return heatmap.numpy() def display_gradcam(img, heatmap, title="Grad-CAM", alpha=0.4): """ 显示原始图像、热力图及叠加后的图像 """ # 将图像从 [0,1] 转换回 [0,255] 整数 img_uint8 = np.uint8(255 * img) # 热力图转伪彩色 heatmap_uint8 = np.uint8(255 * heatmap) jet = plt.get_cmap("jet") jet_colors = jet(np.arange(256))[:, :3] jet_heatmap = jet_colors[heatmap_uint8] jet_heatmap = cv2.resize(jet_heatmap, (img.shape[1], img.shape[0])) jet_heatmap = np.uint8(255 * jet_heatmap) # 叠加 superimposed = jet_heatmap * alpha + img_uint8 superimposed = np.clip(superimposed, 0, 255).astype("uint8") # 绘图 plt.figure(figsize=(12, 4)) plt.subplot(1, 3, 1) plt.imshow(img_uint8) plt.title("原始时频图") plt.axis("off") plt.subplot(1, 3, 2) plt.imshow(heatmap, cmap='jet') plt.title("热力图") plt.axis("off") plt.subplot(1, 3, 3) plt.imshow(superimposed) plt.title(f"叠加: {title}") plt.axis("off") plt.show() # 选择一个测试样本进行 Grad-CAM 可视化 last_conv_layer = "block5_conv3" # VGG16 最后一层卷积的名称 sample_idx = 42 # 可根据需要修改 img_sample = X_test[sample_idx] true_label_idx = np.argmax(y_test[sample_idx]) true_label_name = target_names[true_label_idx] img_batch = np.expand_dims(img_sample, axis=0) heatmap = make_gradcam_heatmap(img_batch, model, last_conv_layer) display_gradcam(img_sample, heatmap, title=true_label_name) # ==================== 9. Faithfulness 指标计算 ==================== def calculate_faithfulness(model, X_data, y_data, layer_name, num_samples=50, mask_top_percent=20): """ 计算 Faithfulness（删除最显著像素后置信度下降百分比） """ print(f"\n--- 开始 Faithfulness 评估，随机选取 {num_samples} 个样本 ---") drops = [] conf_after = [] max_samples = len(X_data) actual_samples = min(num_samples, max_samples) indices = np.random.choice(max_samples, size=actual_samples, replace=False) for idx in tqdm(indices): img = X_data[idx] img_batch = np.expand_dims(img, axis=0) true_label = np.argmax(y_data[idx]) # 原始预测置信度 pred_orig = model.predict(img_batch, verbose=0)[0] conf_orig = pred_orig[true_label] if conf_orig < 0.5: # 只考虑模型原本就有把握的样本 continue # 生成热力图（针对真实类别） heat = make_gradcam_heatmap(img_batch, model, layer_name, pred_index=true_label) heat_resized = cv2.resize(heat, (img.shape[1], img.shape[0])) # 生成掩膜：取热力值最高的 mask_top_percent% 像素 threshold = np.percentile(heat_resized, 100 - mask_top_percent) mask = heat_resized > threshold # 遮挡这些像素（置零） masked_img = copy.deepcopy(img) for c in range(3): masked_img[:, :, c][mask] = 0.0 # 重新预测 masked_batch = np.expand_dims(masked_img, axis=0) pred_new = model.predict(masked_batch, verbose=0)[0] conf_new = pred_new[true_label] # 计算相对下降百分比 drop = max(0, (conf_orig - conf_new) / conf_orig) * 100 drops.append(drop) conf_after.append(conf_new * 100) if len(drops) == 0: print("没有满足置信度要求的样本") return avg_drop = np.mean(drops) avg_conf_after = np.mean(conf_after) print("\n" + "="*40) print("Faithfulness 评估结果") print("="*40) print(f"平均置信度下降: {avg_drop:.2f}%") print(f"遮挡后平均置信度: {avg_conf_after:.2f}%") print("="*40) # 执行 Faithfulness 评估 calculate_faithfulness(model, X_test, y_test, layer_name=last_conv_layer, num_samples=50) # ==================== 10. 专家级报告生成 (LLaMA + 邮件通知) ==================== def extract_heatmap_evidence(heatmap, target_fs=6000): """ 从热力图中提取显著频率信息 """ # 沿时间轴求和得到频率能量分布 freq_profile = np.sum(heatmap, axis=1) peak_freq_idx = np.argmax(freq_profile) nyquist = target_fs / 2 dominant_freq = (peak_freq_idx / heatmap.shape[0]) * nyquist avg_intensity = np.mean(heatmap) * 100 peak_intensity = np.max(heatmap) * 100 return { "dominant_freq": round(dominant_freq, 1), "peak_confidence": round(peak_intensity, 1), "avg_activation": round(avg_intensity, 1) } def generate_llm_report(predicted_label, evidence): """ 利用本地 LLaMA 模型生成诊断报告 """ # 根据预测标签确定故障类型和物理描述 if "Normal" in predicted_label: fault_type = "健康状态" severity = "无" physics = "旋转平稳，未检测到冲击。" action = "继续正常运行，30天后重新评估。" elif "Ball" in predicted_label: fault_type = "滚动体故障" physics = "滚动体表面出现剥落或点蚀，产生以滚动体自转频率为特征的冲击。" if "007" in predicted_label: severity = "早期（0.007英寸）" action = "密切监视，可在下次维护窗口安排检查。" elif "014" in predicted_label: severity = "发展期（0.014英寸）" action = "计划维护，建议两周内更换轴承。" elif "021" in predicted_label: severity = "严重（0.021英寸）" action = "立即停机更换，存在失效风险。" elif "IR" in predicted_label: fault_type = "内圈故障" physics = "内圈滚道出现损伤，球体通过缺陷时产生以BPFI为特征的冲击。" if "007" in predicted_label: severity = "早期（0.007英寸）" action = "密切监视，可在下次维护窗口安排检查。" elif "014" in predicted_label: severity = "发展期（0.014英寸）" action = "计划维护，建议两周内更换轴承。" elif "021" in predicted_label: severity = "严重（0.021英寸）" action = "立即停机更换，存在失效风险。" elif "OR" in predicted_label: fault_type = "外圈故障" physics = "外圈滚道出现损伤，产生以BPFO为特征的冲击。" if "007" in predicted_label: severity = "早期（0.007英寸）" action = "密切监视，可在下次维护窗口安排检查。" elif "014" in predicted_label: severity = "发展期（0.014英寸）" action = "计划维护，建议两周内更换轴承。" elif "021" in predicted_label: severity = "严重（0.021英寸）" action = "立即停机更换，存在失效风险。" else: fault_type = "未知异常" physics = "无法确定具体故障类型。" severity = "未知" action = "立即进行详细检测。" current_date = datetime.now().strftime("%Y-%m-%d") prompt = f""" 你是一名资深振动分析师。请根据以下信息撰写一份诊断报告。 元数据： - 日期：{current_date} - 分析系统：DeepBear-Health AI 输入数据： - 故障类型：{fault_type} - 严重程度：{severity} - 显著频率：~{evidence['dominant_freq']} Hz - 模型置信度：{evidence['peak_confidence']}% 物理背景： {physics} 请撰写一份包含“诊断结论”和“维护建议”两部分的报告。 维护建议请严格遵循："{action}" """ response = ollama.chat(model='llama3.2', messages=[{'role': 'user', 'content': prompt}]) return response['message']['content'] def send_plain_text_email(subject, body, api_key, from_email, to_email): """ 通过 SendGrid 发送纯文本邮件 """ sg = sendgrid.SendGridAPIClient(api_key=api_key) from_email = Email(from_email) to_email = To(to_email) content = Content("text/plain", body) mail = Mail(from_email, to_email, subject, content).get() try: response = sg.client.mail.send.post(request_body=mail) if response.status_code == 202: print("✅ 邮件发送成功") else: print("❌ 邮件发送失败") except Exception as e: print(f"❌ 发送邮件异常: {e}") # 使用之前 Grad-CAM 的样本生成报告（也可另选样本） evidence = extract_heatmap_evidence(heatmap, target_fs=6000) report = generate_llm_report(true_label_name, evidence) print("-"*50) print("自动生成的诊断报告") print("-"*50) print(report) print("-"*50) # 如果检测到严重故障（如 021 级别），则发送邮件通知 if "严重" in report or "Critical" in report: # 请替换为您的 SendGrid 信息 SENDGRID_API_KEY = "YOUR_SENDGRID_API_KEY" FROM_EMAIL = "[email protected]" TO_EMAIL = "[email protected]" send_plain_text_email("【轴承故障告警】严重故障需立即处理", report, SENDGRID_API_KEY, FROM_EMAIL, TO_EMAIL) else: print("未检测到严重故障，不发送邮件")