多模态特征融合的轴承故障诊断混合深度学习框架：时频域协同分析与神经ODE动态建模

开始
│
├─ 初始化设置（随机种子、目录创建）
├─ 自定义层定义（SEBlock1, NeuralODEBlock1, EnhancedGatedAttention1等）
├─ 混合模型构建（build_hybrid_model）
│   ├─ 输入层
│   ├─ 共享卷积层
│   ├─ 三条并行路径：
│   │   ├─ 时域路径（卷积+注意力+神经ODE）
│   │   ├─ 频域路径（傅里叶神经算子+注意力）
│   │   └─ LSTM路径（LSTM+SE块）
│   ├─ 多路径特征融合（注意力加权）
│   └─ 输出层（全连接+分类）
├─ 数据加载与预处理（load_surf_dataset）
│   ├─ 读取.mat/.csv文件
│   ├─ 信号长度标准化
│   └─ 标签编码
├─ 模型编译与训练
│   ├─ 自定义F1评估指标
│   ├─ 类别权重平衡
│   └─ 模型训练
├─ 性能评估与可视化
│   ├─ 训练指标曲线
│   ├─ 测试集评估
│   ├─ 噪声鲁棒性测试
│   ├─ 特征空间可视化（UMAP/t-SNE）
│   ├─ 混淆矩阵/ROC曲线
│   └─ 层激活可视化
└─ 高级信号分析
    ├─ 时域/频域分析
    ├─ 时频分析（STFT/小波）
    └─ 特征分布可视化
结束

# 导入必要的库
import os  # 提供操作系统相关功能
import numpy as np  # 科学计算库
import tensorflow as tf  # 深度学习框架
from tensorflow.keras.layers import (  # Keras层组件
    Input, Conv1D, BatchNormalization, Dense, Dropout, LSTM, Concatenate, Add,
    LayerNormalization, Activation, GlobalAveragePooling1D, Multiply, Reshape, Lambda)
from tensorflow.keras.models import Model  # 模型构建
from tensorflow.keras.utils import to_categorical  # 类别编码
from sklearn.preprocessing import LabelEncoder  # 标签编码
from sklearn.utils import class_weight  # 类别权重计算
import matplotlib.pyplot as plt  # 绘图库
import scipy.io  # MATLAB文件处理
import pandas as pd  # 数据处理
import pickle  # 数据序列化
from scipy.signal import stft, hilbert  # 信号处理
import seaborn as sns  # 统计可视化
import umap  # 降维可视化
from sklearn.manifold import TSNE  # t-SNE降维
from sklearn.metrics import (  # 评估指标
    accuracy_score, precision_score, recall_score, f1_score, roc_auc_score,
    confusion_matrix, ConfusionMatrixDisplay, roc_curve, auc)


# 设置随机种子确保结果可重现
seed_value = 1234
np.random.seed(seed_value)
tf.manual_seed(seed_value)


# 创建报告目录用于保存结果
os.makedirs("Report", exist_ok=True)


# 定义信号维度参数
n_u, n_y = 2, 2  # 输入输出通道数
seq_len = 1024  # 信号长度
input_shape = (seq_len, 1)  # 输入形状
num_classes = 3  # 故障类别数


### 自定义神经网络层实现 ###


class SEBlock1(Layer):
    """挤压激励注意力模块 (Squeeze-and-Excitation Block)"""
    def __init__(self, reduction_ratio=8, **kwargs):
        """
        初始化SE块
        :param reduction_ratio: 通道压缩比例
        """
        super().__init__(**kwargs)
        self.reduction_ratio = reduction_ratio


    def build(self, input_shape):
        """构建层权重"""
        channels = input_shape[-1]  # 获取输入通道数
        # 使用1x1卷积替代全连接层加速计算
        self.conv1 = Conv1D(channels//self.reduction_ratio, 1, activation='relu')  # 降维卷积
        self.conv2 = Conv1D(channels, 1, activation='sigmoid')  # 重建卷积
        super().build(input_shape)


    def call(self, inputs):
        """前向传播逻辑"""
        # 全局平均池化替代全连接层
        se = tf.reduce_mean(inputs, axis=1, keepdims=True)  # 空间维度压缩
        se = self.conv1(se)  # 通道压缩
        se = self.conv2(se)  # 通道重建
        return Multiply()([inputs, se])  # 通道加权


    def compute_output_shape(self, input_shape):
        """输出形状计算"""
        return input_shape




class NeuralODEBlock1(Layer):
    """神经常微分方程块 (Neural Ordinary Differential Equations Block)"""
    def __init__(self, units, time_steps=10, **kwargs):
        """
        初始化神经ODE块
        :param units: 隐藏单元数
        :param time_steps: 时间步长
        """
        super().__init__(**kwargs)
        self.units = units
        self.time_steps = time_steps
        self.dense1 = Dense(units, activation='tanh')  # 非线性变换层
        self.dense2 = Dense(units)  # 导数计算层


    def call(self, x):
        """前向传播实现欧拉积分"""
        outputs = []  # 存储各时间步输出
        h = x  # 初始状态


        # 通过欧拉方法模拟ODE
        for t in range(self.time_steps):
            dx = self.dense2(self.dense1(h))  # 计算导数
            h = h + dx  # 状态更新
            outputs.append(h)  # 保存当前状态


        # 拼接所有时间步输出
        stacked = tf.concat(outputs, axis=1)
        return stacked


    def compute_output_shape(self, input_shape):
        """输出形状计算"""
        return (input_shape[0], self.time_steps, self.units)




class EnhancedGatedAttention1(Layer):
    """增强门控注意力机制 (Enhanced Gated Attention)"""
    def __init__(self, d_model, **kwargs):
        """
        初始化注意力层
        :param d_model: 特征维度
        """
        super().__init__(**kwargs)
        self.d_model = d_model
        # 使用较少注意力头加速计算
        self.mha = tf.keras.layers.MultiHeadAttention(
            num_heads=2, key_dim=d_model//2)  # 多头注意力
        # 使用1D卷积替代全连接加速门控计算
        self.gate = Conv1D(1, kernel_size=1, activation='sigmoid')  # 门控生成
        # 使用批归一化替代层归一化加速训练
        self.norm = BatchNormalization()  # 归一化层


    def call(self, x):
        """前向传播逻辑"""
        attn = self.mha(x, x)  # 自注意力计算
        gate = self.gate(x)  # 门控信号生成
        # 残差连接+门控注意力
        return self.norm(x + attn * gate)  # 特征更新




class VanillaSelfAttention(Layer):
    """标准自注意力机制 (Vanilla Self-Attention)"""
    def __init__(self, d_model, **kwargs):
        """初始化标准注意力层"""
        super().__init__(**kwargs)
        self.d_model = d_model
        # 使用与增强门控注意力相同的配置
        self.mha = tf.keras.layers.MultiHeadAttention(
            num_heads=2, key_dim=d_model//2)
        self.norm = BatchNormalization()


    def call(self, x):
        """前向传播逻辑"""
        attn = self.mha(x, x)  # 自注意力计算
        return self.norm(x + attn)  # 残差连接




class FourierNeuralOperator1(Layer):
    """傅里叶神经算子 (Fourier Neural Operator)"""
    def __init__(self, modes, filters, **kwargs):
        """
        初始化FNO层
        :param modes: 保留的傅里叶模式数
        :param filters: 滤波器数量
        """
        super().__init__(**kwargs)
        self.modes = modes
        self.filters = filters
        # 使用更窄的全连接层加速计算
        self.fft_dense = Dense(filters//4, activation='gelu')  # 傅里叶域变换
        self.ifft_dense = Dense(filters//2)  # 逆傅里叶域变换


    def call(self, x):
        """前向传播逻辑"""
        # 1. 快速傅里叶变换
        x_fft = tf.signal.fft(tf.cast(x, tf.complex64))
        x_fft = tf.math.real(x_fft[..., :self.modes])  # 保留主要模式


        # 2. 傅里叶域特征变换
        x_fft = self.fft_dense(x_fft)
        x_fft = self.ifft_dense(x_fft)


        # 3. 填充并逆变换回时域
        x_fft = tf.pad(x_fft, [[0,0],[0,0],[0,tf.shape(x)[1]-self.modes]])
        return tf.math.real(tf.signal.ifft(tf.cast(x_fft, tf.complex64)))




### 混合模型构建函数 ###


def build_hybrid_model(input_shape, num_classes):
    """
    构建三路径混合故障诊断模型
    :param input_shape: 输入形状 (序列长度, 通道数)
    :param num_classes: 分类类别数
    :return: 构建好的Keras模型
    """
    # 1. 输入层
    inputs = Input(shape=input_shape)


    # ===== 共享预处理层 =====
    shared_conv = Conv1D(32, 3, padding='same', name='shared_conv')(inputs)


    # ===== 路径1: 时域特征提取 =====
    # 1.1 基础特征提取
    x_time = Conv1D(64, 3, padding='same', name='time_conv1')(shared_conv)
    x_time = EnhancedGatedAttention1(d_model=64, name='time_attn1')(x_time)
    # 1.2 中间特征保存用于跨路径连接
    time_mid = Conv1D(64, 3, padding='same', name='time_mid')(x_time)
    # 1.3 神经ODE时间演化
    x_time = GlobalAveragePooling1D(name='time_gap1')(x_time)
    x_time = Reshape((1, 64))(x_time)
    x_time = NeuralODEBlock1(units=64, time_steps=50, name='time_ode')(x_time)
    x_time = GlobalAveragePooling1D(name='time_gap2')(x_time)


    # ===== 路径2: 频域特征提取 =====
    # 2.1 傅里叶神经算子
    x_freq = FourierNeuralOperator1(modes=16, filters=64, name='fno')(shared_conv)
    # 2.2 跨路径连接时域特征
    x_freq = Concatenate(axis=-1, name='freq_concat')([x_freq, time_mid])
    x_freq = Conv1D(64, 1, name='freq_merge')(x_freq)
    # 2.3 频域注意力增强
    x_freq = EnhancedGatedAttention1(d_model=64, name='freq_attn')(x_freq)
    x_freq = GlobalAveragePooling1D(name='freq_gap')(x_freq)


    # ===== 路径3: LSTM时序建模 =====
    # 3.1 LSTM时序特征提取
    x_lstm = LSTM(64, return_sequences=True, name='lstm1')(shared_conv)
    # 3.2 时域特征适配
    time_mid_adjusted = Conv1D(64, 1, name='time_mid_adjust')(time_mid)
    # 3.3 频域特征适配
    x_freq_expanded = Lambda(
        lambda x: tf.repeat(x[0], tf.shape(x[1])[1], axis=1),
        name='freq_expand')([x_freq[:, tf.newaxis, :], x_lstm])
    # 3.4 多源特征融合
    x_lstm = Concatenate(axis=-1, name='lstm_concat')([
        x_lstm, time_mid_adjusted, x_freq_expanded])
    x_lstm = Conv1D(64, 1, name='lstm_merge')(x_lstm)
    # 3.5 通道注意力增强
    x_lstm = SEBlock1(name='lstm_se')(x_lstm)
    x_lstm = GlobalAveragePooling1D(name='lstm_gap')(x_lstm)


    # ===== 多路径特征融合 =====
    # 4.1 特征拼接
    fused = Concatenate(name='final_concat')([x_time, x_freq, x_lstm])
    # 4.2 注意力加权融合
    attention_units = fused.shape[-1]  # 自动获取特征维度
    attention_weights = Dense(attention_units, activation='softmax')(fused)
    fused = Multiply(name='attention_scale')([fused, attention_weights])


    # ===== 输出层 =====
    # 5.1 全连接层
    out = Dense(128, activation='gelu', name='dense1')(fused)
    out = Dropout(0.5, name='dropout1')(out)
    # 5.2 分类输出层
    out = Dense(num_classes, activation='softmax', name='output')(out)


    return Model(inputs, out)




### 数据加载函数 ###


def load_surf_dataset(folder, seq_len=1024):
    """
    加载SURF轴承故障数据集
    :param folder: 数据文件夹路径
    :param seq_len: 信号长度
    :return: 信号数组和标签数组
    """
    X = []  # 存储信号
    y = []  # 存储标签
    class_names = sorted(os.listdir(folder))  # 获取故障类别


    # 遍历每个故障类别
    for label in class_names:
        class_path = os.path.join(folder, label)
        # 遍历类别文件夹中的文件
        for fname in os.listdir(class_path):
            file_path = os.path.join(class_path, fname)
            # 处理MATLAB数据文件
            if fname.endswith('.mat'):
                mat = scipy.io.loadmat(file_path)
                for key in mat:
                    if not key.startswith("__"):  # 跳过系统变量
                        signal = mat[key].squeeze()
                        break
            # 处理CSV数据文件
            elif fname.endswith('.csv'):
                df = pd.read_csv(file_path, header=None)
                signal = df.values.squeeze()
            else:
                continue


            # 信号长度标准化
            if len(signal) >= seq_len:
                signal = signal[:seq_len]  # 截断
            else:
                # 填充不足部分
                signal = np.pad(signal, (0, seq_len - len(signal)))


            X.append(signal)
            y.append(label)


    return np.array(X), np.array(y)




### 自定义评估指标 ###


def f1_score(y_true, y_pred):
    """自定义F1分数计算函数"""
    y_pred = tf.round(y_pred)  # 预测概率转类别
    # 计算真阳性、假阳性、假阴性
    tp = tf.reduce_sum(tf.cast(y_true * y_pred, 'float'), axis=0)
    fp = tf.reduce_sum(tf.cast((1 - y_true) * y_pred, 'float'), axis=0)
    fn = tf.reduce_sum(tf.cast(y_true * (1 - y_pred), 'float'), axis=0)


    # 计算精确率和召回率
    precision = tp / (tp + fp + tf.keras.backend.epsilon())
    recall = tp / (tp + fn + tf.keras.backend.epsilon())


    # 计算F1分数
    f1 = 2 * precision * recall / (precision + recall + tf.keras.backend.epsilon())
    return tf.reduce_mean(f1)  # 返回宏平均F1




### 主执行流程 ###


if __name__ == "__main__":
    # === 数据准备 ===
    # 加载训练集和测试集
    X_train, y_train = load_surf_dataset("Veriseti_Surf/train", seq_len=seq_len)
    X_test, y_test = load_surf_dataset("Veriseti_Surf/test", seq_len=seq_len)


    # 添加通道维度 (N, seq_len) -> (N, seq_len, 1)
    X_train = X_train[..., np.newaxis]
    X_test = X_test[..., np.newaxis]


    # 标签编码和one-hot转换
    encoder = LabelEncoder()
    y_train_enc = encoder.fit_transform(y_train)
    y_test_enc = encoder.transform(y_test)
    y_train_cat = to_categorical(y_train_enc, num_classes=num_classes)
    y_test_cat = to_categorical(y_test_enc, num_classes=num_classes)


    # === 模型构建与编译 ===
    model = build_hybrid_model(input_shape=input_shape, num_classes=num_classes)


    # 计算类别权重处理不平衡数据
    class_weights = class_weight.compute_class_weight(
        class_weight='balanced',
        classes=np.unique(y_train_enc),
        y=y_train_enc
    )
    class_weights_dict = dict(enumerate(class_weights))


    # 模型编译
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4),
        loss='categorical_crossentropy',
        metrics=[
            'accuracy',
            tf.keras.metrics.Precision(name='precision'),
            tf.keras.metrics.Recall(name='recall'),
            f1_score  # 使用自定义F1指标
        ]
    )


    # 打印模型结构
    model.summary()


    # === 模型训练 ===
    history = model.fit(
        X_train, y_train_cat,
        validation_data=(X_test, y_test_cat),
        epochs=50,
        batch_size=256,
        class_weight=class_weights_dict  # 类别权重
    )


    # === 训练结果分析 ===
    # 提取训练指标
    train_acc = history.history['accuracy']
    train_prec = history.history['precision']
    train_rec = history.history['recall']
    train_f1 = history.history['f1_score']
    train_loss = history.history['loss']


    # 提取验证指标
    val_acc = history.history['val_accuracy']
    val_prec = history.history['val_precision']
    val_rec = history.history['val_recall']
    val_f1 = history.history['val_f1_score']
    val_loss = history.history['val_loss']


    # 保存指标结果
    with open("ablation/ab1_acc.pkl", "wb") as f:
        pickle.dump(train_acc, f)
    with open("ablation/ab1_prec.pkl", "wb") as f:
        pickle.dump(train_prec, f)
    with open("ablation/ab1_rec.pkl", "wb") as f:
        pickle.dump(train_rec, f)
    with open("ablation/ab1_f1.pkl", "wb") as f:
        pickle.dump(train_f1, f)


    # === 模型可视化分析 ===
    # 1. 训练过程指标曲线
    plt.figure(figsize=(10, 6))
    epochs = range(1, len(train_acc)+1)
    plt.plot(epochs, train_acc, label='Accuracy')
    plt.plot(epochs, train_prec, label='Precision')
    plt.plot(epochs, train_rec, label='Recall')
    plt.plot(epochs, train_f1, label='F1-Score')
    plt.title('Training Metrics Over Epochs')
    plt.xlabel('Epochs')
    plt.ylabel('Score')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("Report/training_metrics.pdf")
    plt.show()


    # 2. 特征空间可视化 (UMAP)
    feature_extractor = Model(inputs=model.input, outputs=model.layers[-3].output)
    features = feature_extractor.predict(X_test)
    reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, metric='euclidean', random_state=42)
    embedding = reducer.fit_transform(features)


    plt.figure(figsize=(8, 6))
    scatter = plt.scatter(embedding[:, 0], embedding[:, 1], c=y_test_enc, cmap='Set1', alpha=0.8)
    handles, _ = scatter.legend_elements()
    plt.legend(handles=handles, labels=encoder.classes_.tolist())
    plt.title("UMAP Visualization of Feature Space")
    plt.xlabel("UMAP 1")
    plt.ylabel("UMAP 2")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("Report/feature_space_umap.pdf")
    plt.show()


    # 3. 混淆矩阵
    y_pred = model.predict(X_test)
    y_pred_classes = np.argmax(y_pred, axis=1)
    cm = confusion_matrix(y_test_enc, y_pred_classes)
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=encoder.classes_)
    disp.plot(cmap=plt.cm.Blues)
    plt.title("Confusion Matrix")
    plt.savefig("Report/confusion_matrix.pdf")
    plt.show()


    # 4. ROC曲线
    y_test_onehot = to_categorical(y_test_enc, num_classes=num_classes)
    plt.figure(figsize=(8, 6))
    for i in range(num_classes):
        fpr, tpr, _ = roc_curve(y_test_onehot[:, i], y_pred[:, i])
        roc_auc = auc(fpr, tpr)
        plt.plot(fpr, tpr, label=f'{encoder.classes_[i]} (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve for Each Class')
    plt.legend()
    plt.grid(True)
    plt.savefig("Report/roc_curve.pdf")
    plt.show()


    # === 噪声鲁棒性测试 ===
    def add_gaussian_noise(signal, snr_db):
        """添加高斯白噪声"""
        signal_power = np.mean(signal ** 2)
        snr_linear = 10 ** (snr_db / 10)
        noise_power = signal_power / snr_linear
        noise = np.random.normal(0, np.sqrt(noise_power), signal.shape)
        return signal + noise


    # 测试不同信噪比下的性能
    snr_levels = [5, 10, 15, 20]
    X_test_noisy = {}


    # 生成加噪测试集
    for snr in snr_levels:
        noisy_signals = []
        for signal in X_test.squeeze():
            noisy_signal = add_gaussian_noise(signal, snr)
            noisy_signals.append(noisy_signal)
        X_test_noisy[snr] = np.array(noisy_signals)[..., np.newaxis]


    # 评估各信噪比下的模型性能
    snr_results = {}
    for snr in [np.inf] + snr_levels:  # 包含无噪声情况
        if snr == np.inf:
            X_input = X_test
            snr_label = "∞"
        else:
            X_input = X_test_noisy[snr]
            snr_label = f"{snr} dB"


        # 模型预测
        y_pred = model.predict(X_input)
        y_pred_cls = np.argmax(y_pred, axis=1)


        # 计算评估指标
        acc = accuracy_score(y_test_enc, y_pred_cls)
        prec = precision_score(y_test_enc, y_pred_cls, average='macro', zero_division=0)
        rec = recall_score(y_test_enc, y_pred_cls, average='macro', zero_division=0)
        f1 = f1_score(y_test_enc, y_pred_cls, average='macro', zero_division=0)


        snr_results[snr_label] = {'acc': acc, 'prec': prec, 'rec': rec, 'f1': f1}


    # 可视化噪声鲁棒性结果
    snr_labels = list(snr_results.keys())
    acc_list = [snr_results[snr]['acc'] for snr in snr_labels]
    prec_list = [snr_results[snr]['prec'] for snr in snr_labels]
    rec_list = [snr_results[snr]['rec'] for snr in snr_labels]
    f1_list = [snr_results[snr]['f1'] for snr in snr_labels]


    plt.figure(figsize=(10, 6))
    plt.plot(snr_labels, acc_list, marker='o', label='Accuracy')
    plt.plot(snr_labels, prec_list, marker='s', label='Precision')
    plt.plot(snr_labels, rec_list, marker='^', label='Recall')
    plt.plot(snr_labels, f1_list, marker='d', label='F1-Score')
    plt.title('Model Performance under Varying SNR Levels')
    plt.xlabel('SNR (dB)')
    plt.ylabel('Score')
    plt.ylim(0.0, 1.05)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig("Report/noise_robustness.pdf")
    plt.show()


    # === 信号分析 ===
    # 1. 时域分析
    def plot_time_domain(signal, title="Time-Domain Signal", sampling_rate=1000):
        time = np.linspace(0, len(signal) / sampling_rate, len(signal))
        plt.figure(figsize=(10, 3))
        plt.plot(time, signal)
        plt.title(title)
        plt.xlabel("Time (s)")
        plt.ylabel("Amplitude")
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"Report/{title.replace(' ', '_')}.pdf")
        plt.show()


    # 2. 频域分析
    def plot_fft(signal, sampling_rate=1000, title="Frequency Spectrum"):
        N = len(signal)
        freq = np.fft.fftfreq(N, d=1/sampling_rate)
        fft_vals = np.fft.fft(signal)
        plt.figure(figsize=(10, 3))
        plt.plot(freq[:N//2], np.abs(fft_vals)[:N//2])
        plt.title(title)
        plt.xlabel("Frequency (Hz)")
        plt.ylabel("Amplitude")
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"Report/{title.replace(' ', '_')}.pdf")
        plt.show()


    # 应用分析函数
    sample_signal = X_test[0].squeeze()
    plot_time_domain(sample_signal, "Normal Bearing Signal")
    plot_fft(sample_signal, "Normal Bearing Spectrum")


    # === 模型解释性分析 ===
    # 1. 注意力权重可视化
    attention_layer = model.get_layer('time_attn1')
    attention_extractor = Model(inputs=model.input, outputs=attention_layer.output)
    attention_weights = attention_extractor.predict(X_test[:1]).squeeze()


    plt.figure(figsize=(12, 4))
    plt.plot(sample_signal, label='Input Signal')
    plt.plot(attention_weights * np.max(sample_signal), alpha=0.7, label='Attention Weights')
    plt.title('Attention Weights over Time Domain Signal')
    plt.xlabel('Time step')
    plt.legend()
    plt.savefig("Report/attention_weights.pdf")
    plt.show()


    # 2. 神经ODE状态演化可视化
    ode_extractor = Model(inputs=model.input, outputs=model.get_layer('time_ode').output)
    ode_outputs = ode_extractor.predict(X_test[:1]).squeeze()


    plt.figure(figsize=(12, 6))
    for i in range(5):  # 可视化5个特征
        plt.plot(ode_outputs[:, i], label=f'Feature {i+1}')
    plt.xlabel('Neural ODE Time Steps')
    plt.ylabel('Feature value')
    plt.title('Neural ODE Feature Evolution Over Time')
    plt.legend()
    plt.savefig("Report/neural_ode_evolution.pdf")
    plt.show()