面向强噪声场景的深度残差收缩网络原理及Python实现

从实际工业运行环境中的机械传动系统上采集到的原始振动信号通常含有大量背景噪声，这在一定程度上掩盖了与故障相关的关键冲击特征。针对这一问题，《IEEE Transactions on Industrial Informatics》于2020年刊发了题为“Deep Residual Shrinkage Networks for Fault Diagnosis”的研究性论文，提出了一种名为深度残差收缩网络（Deep Residual Shrinkage Network, DRSN）的特征提取框架。该研究的主要思路是在传统的残差网络中引入参数化的软阈值收缩模块，并结合特征重标定机制，对前向传播过程中的特征图冗余信息进行动态约束，以此为基础开展非平稳工况下的信号模式识别。
本文基于上述文献完成了DRSN模型的代码复现，并以公开的凯斯西储大学(Case Western Reserve University, CWRU)轴承振动数据集为例开展了实验验证。在评估阶段，为测试样本添加了特定信噪比（-8 dB）的高斯白噪声，以系统性评估模型在强噪声干扰下的分类性能。
一、软阈值函数
深度残差收缩网络的核心运算单元为软阈值化。与ReLU等传统激活函数直接将负值区间特征置零不同，软阈值函数在正负特征区间均具有对称收缩特性。其数学表达式为y = sgn(x) * max(|x| -τ, 0)，其中x与y分别表示输入与输出特征，τ为软阈值参数（且恒非负）。
该非线性操作在特征提取中具有以下两方面作用：
①对低幅值噪声分量的置零（或“滤除”）：当输入特征元素的绝对幅值低于设定的阈值τ时，软阈值函数将其输出直接置为零。这种处理方式在信号去噪领域常被视为对背景噪声或无效特征的屏蔽机制，抑制局部感受野内的微弱扰动向深层网络传播。
② 对于绝对值大于设定阈值的特征元素，软阈值函数在保留原有正负极性符号的基础上，将其幅值向零点方向收缩一个固定的阈值大小τ，即在幅值上进行大小为τ的平移，避免了硬阈值函数带来的不连续截断问题。
从导数性质来看，软阈值函数的导数在除阈值点外的取值仅为0或1，这种导数特性在反向传播时能够有效缓解梯度爆炸或梯度消失问题。因此，该操作在保持特征映射连续性的同时，为深层网络构建了非线性的特征过滤机制。

二、自适应阈值参数学习
在经典的小波阈值去噪算法中，阈值参数τ通常依赖领域先验知识进行人工设定，或基于特定统计假设进行估计。为将该机制融入端到端的深度学习框架，DRSN 在基础残差模块内部集成了一个浅层的前馈神经网络结构，利用反向传播算法自适应学习特征图的通道级阈值。
该子网络的前向传播包含以下三个步骤：
① 特征空间信息的压缩：对经过若干层一维卷积操作后所得的特征图取绝对值，并沿时间/序列维度执行全局平均池化（Global Average Pooling, GAP），从而将输入张量降维为通道特征向量。
② 相对比例系数的映射：将上述通道特征向量输入至一个包含两层全连接结构的多层感知机（MLP）网络中。该模块的输出端经由Sigmoid激活函数处理，将连续特征映射至(0, 1)区间内的比例系数α，作为决定各通道阈值收缩尺度的比例因子（或权重系数）。
③ 自适应阈值的计算：采用计算所得的比例系数α乘以对应通道的特征绝对值均值，即按照计算公式τ= α* average(|x|) 计算得到最终的阈值参数。该运算逻辑限定了生成的阈值不超出当前特征图内的绝对幅值均值，确保软阈值化后仍有非零特征向后层传递。
该数据驱动的注意力机制能基于输入特征的全局统计信息，自适应生成通道收缩阈值，实现特征幅值的动态重标定。
根据阈值共享机制的差异，文献中提出了通道共享（DRSN-CS）与通道独立（DRSN-CW）两种网络结构变体。本文重点探讨的DRSN-CW模型为特征图的每一个通道均配置了独立的标量阈值，使得网络能够根据不同频段或信号分量的特征分布，动态优化各通道的阈值参数。
三、网络架构与优化策略
在网络结构设计方面，本文采用逐级扩增通道数的分层结构。起始主干层配置为单通道至32通道的升维卷积，其后依次串接包含32、64、128通道的三个残差收缩阶段。随着网络深度的增加与感受野的扩张，信号特征图在时间轴上被逐步降采样，而其原始信号中蕴含的局部波形特征则逐步转化为更具判别性的高阶抽象特征。
为防止深层参数在迭代优化过程中出现梯度震荡或在特定样本集上产生过拟合现象，残差收缩模块的卷积核引入了基于L2范数的权重衰减项（正则化系数设为1e-4）。施加L2正则化可有效限制网络权重范数，控制模型复杂度，从而避免模型对局部训练噪声过拟合，提升其在强噪声干扰下的鲁棒性。
针对输入端的数据分布，本文模型在训练阶段引入了随机数据增强策略。数据预处理阶段包含三个维度的数据增强操作：通过随机循环平移，减轻因序列截断导致的时间平移敏感性；以一定概率向序列局部添加瞬态异常脉冲扰动以模拟偶发冲击；以及在原始输入信号中叠加信噪比介于-8 dB至8 dB之间动态变化的高斯白噪声。该训练机制通过增加训练样本的多样性，提升模型对背景干扰的抗噪鲁棒性。
四、实验验证
本文在凯斯西储大学（CWRU）公开的轴承驱动端振动数据集上进行了算法验证。实验选取了不同严重程度与损伤位置的信号，构建了包含十种健康状态的分类数据集。将长度为1024的数据样本按70%、15% 和15%的比例随机划分为训练集、验证集与测试集，并执行了基于训练集统计参数的Z-Score标准化。考虑到原始纯净数据集的分类难度较低，向验证集与测试集中引入了高强度的合成高斯白噪声。

基于TensorFlow框架编写的完整深度残差收缩网络Python代码如下所示。配置好运行环境并指定.mat数据文件路径后，即可启动端到端的模型训练与评估。

# # =============================================================================
# 
# 本程序为论文《Deep Residual Shrinkage Networks for Fault Diagnosis》的
# TensorFlow/Keras 复现实现。代码基于 CWRU（Case Western Reserve University）
# 轴承振动数据集，构建一维深度残差收缩网络（DRSN-CW）模型，用于滚动轴承
# 故障多分类任务。
#
# 整体功能包括：
# 1. 运行环境与GPU资源自动配置；
# 2. CWRU原始 .mat 振动信号解析与滑动窗口分割；
# 3. 基于训练集统计量的标准化处理；
# 4. 强噪声（指定SNR）注入与在线数据增强；
# 5. DRSN核心模块（含通道级自适应软阈值收缩单元 RSBU-CW）构建；
# 6. 端到端模型训练、验证与强噪声环境下测试评估。
#
# 网络结构遵循原论文提出的“残差结构 + 通道自适应软阈值收缩”思想，
# 通过学习通道级阈值对特征图进行非线性收缩，以提升模型在强噪声环境下
# 的鲁棒性与泛化能力。
#
# 参考文献（Reference）：
# Zhao, M., Zhong, S., Fu, X., Tang, B., & Pecht, M. (2020).
# Deep residual shrinkage networks for fault diagnosis.
# IEEE Transactions on Industrial Informatics, 16(7), 4681–4690.
# https://doi.org/10.1109/TII.2019.2943898
# =============================================================================

import os
import sys
import logging
import numpy as np
import scipy.io as io
import tensorflow as tf
from tensorflow.keras import layers, models, regularizers
from sklearn.model_selection import train_test_split

# =============================================================================
# 第一部分：运行时计算环境优化与资源配置
# =============================================================================

# 全局日志记录器配置
logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s: %(message)s')

class DeviceConfig:
    """
    设备分配管理类：
    负责检测底层计算硬件（CPU/GPU）并配置 TensorFlow 的显存分配策略。
    """
    @staticmethod
    def init_device():
        """
        执行硬件上下文配置。
        采用动态显存增长策略 (Allow Growth)，避免初始化时强占全部显存。
        """
        try:
            import sklearn
            import scipy
        except ImportError as dependency_error:
            logging.error("核心依赖库缺失: {0}".format(dependency_error))
            sys.exit(1)

        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
        physical_gpus = tf.config.list_physical_devices('GPU')
        
        if physical_gpus:
            try:
                for device in physical_gpus:
                    tf.config.experimental.set_memory_growth(device, True)
                logging.info("成功初始化 GPU 环境，已启用 {0} 个设备的显存自适应模式。".format(len(physical_gpus)))
            except RuntimeError as runtime_err:
                logging.error("GPU 运行时策略设置失败: {0}".format(runtime_err))
        else:
            logging.info("未检测到可用 GPU，计算流水线将切换至 CPU 模式。")

# 初始化运行环境
DeviceConfig.init_device()

# =============================================================================
# 第二部分：信号预处理与数据集构建逻辑
# =============================================================================

class CWRULoader:
    """
    振动信号加载类：
    实现对 CWRU 原始 MATLAB 格式数据的解析、提取及滑动窗口分割。
    """
    def __init__(self, directory, window_size=1024):
        """
        初始化数据引擎。
        :param directory: 原始数据集存放的物理路径。
        :param window_size: 序列样本的时间跨度。
        """
        self.base_dir = os.path.abspath(directory)
        self.seq_len = window_size
        self.step_size = window_size  # 默认步长等于窗口长度，即无重叠采样

    def _load_mat(self, file_path):
        """
        私有方法：解析 .mat 文件并定位驱动端振动分量。
        """
        try:
            raw_content = io.loadmat(file_path)
            for identifier in raw_content.keys():
                if 'DE_time' in identifier:
                    return raw_content[identifier].ravel()
        except Exception:
            return None
        return None

    def load_data(self, LABEL_MAP):
        """
        构建带标签的特征矩阵。
        :param LABEL_MAP: 映射关系字典，格式为 {标签编号: [对应文件名列表]}。
        :return: (特征张量, 标签向量) 的 Numpy 数组。
        """
        feature_pool = []
        label_pool = []
        valid_file_counter = 0

        for class_id, resource_list in LABEL_MAP.items():
            for resource_name in resource_list:
                full_path = os.path.join(self.base_dir, "{0}.mat".format(resource_name))
                if not os.path.exists(full_path):
                    continue
                
                time_series = self._load_mat(full_path)
                if time_series is None:
                    continue
                
                valid_file_counter += 1
                # 滑动窗口切片处理
                num_segments = (len(time_series) - self.seq_len) // self.step_size + 1
                for i in range(num_segments):
                    offset = i * self.step_size
                    segment = time_series[offset : offset + self.seq_len]
                    feature_pool.append(segment)
                    label_pool.append(class_id)
        
        if valid_file_counter == 0:
            raise IOError("在指定目录 {0} 下未检索到符合条件的 .mat 信号文件。".format(self.base_dir))
            
        return np.array(feature_pool, dtype='float32'), np.array(label_pool, dtype='int32')

def add_noise(signal_batch, snr_ratio):
    """
    环境模拟工具：向输入信号注入指定信噪比的高斯白噪声。
    基于功率谱密度计算噪声幅值，支持动态 SNR 范围。
    """
    signal_batch = np.array(signal_batch)
    random_state = np.random.default_rng()
    
    # 确定当前的信噪比数值
    target_snr = snr_ratio if isinstance(snr_ratio, (int, float)) else random_state.uniform(snr_ratio[0], snr_ratio[1])
    
    # 计算信号平均功率及对应的噪声功率阈值
    avg_signal_power = np.mean(np.square(signal_batch), axis=1, keepdims=True)
    avg_noise_power = avg_signal_power / (10 ** (target_snr / 10.0))
    
    noise_component = random_state.normal(0, np.sqrt(avg_noise_power), signal_batch.shape)
    return (signal_batch + noise_component).astype('float32')

# =============================================================================
# 第三部分：深度残差收缩网络 (DRSN) 架构实现
# =============================================================================

class SoftThresholding(layers.Layer):
    """
    自适应收缩算子层：
    执行基于深度残差收缩网络核心理念的软阈值化操作。
    公式推导：y = sign(x) * max(|x| - threshold, 0)
    """
    def __init__(self, **kwargs):
        super(SoftThresholding, self).__init__(**kwargs)

    def call(self, inputs):
        """
        前向传递逻辑。
        :param inputs: 包含 [特征张量, 阈值向量] 的列表。
        """
        feature_map, tau = inputs
        # 维度对齐：将阈值向量扩展至与特征图相同的空间维度
        tau_expanded = tf.expand_dims(tau, axis=1)
        
        magnitude = tf.abs(feature_map)
        shrinkage_term = tf.maximum(magnitude - tau_expanded, 0.0)
        return tf.sign(feature_map) * shrinkage_term

class RSBU_CW(layers.Layer):
    """
    残差收缩单元 (Residual Shrinkage Building Unit - CW)：
    DRSN 的基本计算单元。通过子网络分支学习通道级阈值，并对主干特征进行非线性降噪。
    """
    def __init__(self, filters, kernel_size, strides=1, **kwargs):
        super(RSBU_CW, self).__init__(**kwargs)
        self.out_channels = filters
        self.stride = strides
        self.k_size = kernel_size
        self.weight_decay = regularizers.l2(1e-4)

        self.identity_mapping = None
        
        # 卷积变换流水线
        self.batch_norm_a = layers.BatchNormalization()
        self.relu_a = layers.Activation('relu')
        self.conv_a = layers.Conv1D(filters, kernel_size, strides=strides, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        self.batch_norm_b = layers.BatchNormalization()
        self.relu_b = layers.Activation('relu')
        self.conv_b = layers.Conv1D(filters, kernel_size, strides=1, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        # 阈值学习子网络
        self.pooling_global = layers.GlobalAveragePooling1D()
        self.dense_fc1 = layers.Dense(filters, kernel_initializer='he_normal')
        self.bn_fc = layers.BatchNormalization()
        self.relu_fc = layers.Activation('relu')
        self.dense_fc2 = layers.Dense(filters, activation='sigmoid')
        self.shrinkage_layer = SoftThresholding()

    def build(self, input_shape):
        """
        构建快捷连接。
        若维度不匹配，则引入 1x1 卷积进行空间和通道投影。
        """
        if self.stride != 1 or input_shape[-1] != self.out_channels:
            self.identity_mapping = models.Sequential([
                layers.Conv1D(self.out_channels, 1, strides=self.stride, padding='same'),
            ])
        super(RSBU_CW, self).build(input_shape)

    def call(self, inputs):
        """
        逻辑流程：输入 -> 两次卷积变换 -> 计算通道权重 -> 派生阈值 -> 软阈值降噪 -> 残差融合。
        """
        identity_feat = inputs
        if self.identity_mapping:
            identity_feat = self.identity_mapping(inputs)

        # 特征映射分支
        feat = self.batch_norm_a(inputs)
        feat = self.relu_a(feat)
        feat = self.conv_a(feat)
        feat = self.batch_norm_b(feat)
        feat = self.relu_b(feat)
        feat = self.conv_b(feat)

        # 阈值计算分支
        signal_abs = tf.abs(feat)
        abs_mean_vec = self.pooling_global(signal_abs)
        
        z = self.dense_fc1(abs_mean_vec)
        z = self.bn_fc(z)
        z = self.relu_fc(z)
        alpha = self.dense_fc2(z)
        
        # 阈值 = 绝对值均值 * 学习到的比例因子
        tau = tf.multiply(alpha, abs_mean_vec)
        
        # 执行非线性收缩与相加
        denoised_feat = self.shrinkage_layer([feat, tau])
        return layers.Add()([denoised_feat, identity_feat])

class DRSN_CW(models.Model):
    """
    深度残差收缩网络主体架构：
    包含初始特征扩展层、多级 RSBU 模块堆叠、全局池化以及 Softmax 分类器。
    """
    def __init__(self, num_classes):
        super(DRSN_CW, self).__init__(name="DRSN_Bearing_Expert")
        
        # 输入适配层
        self.stem_conv = layers.Conv1D(32, 15, strides=2, padding='same', kernel_initializer='he_normal')
        self.stem_bn = layers.BatchNormalization()
        self.stem_act = layers.Activation('relu')
        
        # 残差收缩块序列
        self.stage_1 = [RSBU_CW(32, 5, strides=2), RSBU_CW(32, 5, strides=1)]
        self.stage_2 = [RSBU_CW(64, 5, strides=2), RSBU_CW(64, 5, strides=1)]
        self.stage_3 = [RSBU_CW(128, 5, strides=2), RSBU_CW(128, 5, strides=1)]
        
        # 输出头
        self.head_bn = layers.BatchNormalization()
        self.head_act = layers.Activation('relu')
        self.gap_layer = layers.GlobalAveragePooling1D()
        self.softmax_layer = layers.Dense(num_classes, activation='softmax')

    def call(self, inputs):
        """
        端到端推理路径。
        """
        x = self.stem_conv(inputs)
        x = self.stem_act(self.stem_bn(x))
        
        for block in self.stage_1 + self.stage_2 + self.stage_3:
            x = block(x)
            
        x = self.head_act(self.head_bn(x))
        x = self.gap_layer(x)
        return self.softmax_layer(x)

# =============================================================================
# 第四部分：自动化诊断流水线与评估逻辑
# =============================================================================

def run_pipeline(data_root, input_length=1024):
    """
    主控程序：驱动数据加载、模型训练、强噪声测试的全流程。
    """
    # CWRU 故障类型元组
    LABEL_MAP = {
        0: ['Normal_0', 'Normal_1', 'Normal_2', 'Normal_3'],
        1: ['IR007_0', 'IR007_1', 'IR007_2', 'IR007_3'],
        2: ['IR014_0', 'IR014_1', 'IR014_2', 'IR014_3'],
        3: ['IR021_0', 'IR021_1', 'IR021_2', 'IR021_3'],
        4: ['B007_0', 'B007_1', 'B007_2', 'B007_3'],
        5: ['B014_0', 'B014_1', 'B014_2', 'B014_3'],
        6: ['B021_0', 'B021_1', 'B021_2', 'B021_3'],
        7: ['OR007@6_0', 'OR007@6_1', 'OR007@6_2', 'OR007@6_3'],
        8: ['OR014@6_0', 'OR014@6_1', 'OR014@6_2', 'OR014@6_3'],
        9: ['OR021@6_0', 'OR021@6_1', 'OR021@6_2', 'OR021@6_3']
    }
    
    # 实例化处理引擎
    loader = CWRULoader(directory=data_root, window_size=input_length)
    
    try:
        x_raw, y_raw = loader.load_data(LABEL_MAP)
    except Exception as fatal_err:
        logging.error("数据准备阶段发生严重错误: %s", fatal_err)
        return

    # 数据集切分 (7:1.5:1.5)
    train_x, tmp_x, train_y, tmp_y = train_test_split(x_raw, y_raw, test_size=0.3, random_state=42)
    val_x, test_x, val_y, test_y = train_test_split(tmp_x, tmp_y, test_size=0.5, random_state=42)
    
    # 基于训练统计量的 Z-Score 标准化
    train_mean = np.mean(train_x)
    train_std = np.std(train_x)
    
    def norm_and_reshape(data):
        return ((data - train_mean) / train_std).reshape(-1, input_length, 1)

    x_train_final = norm_and_reshape(train_x)
    x_val_final = norm_and_reshape(val_x)
    x_test_final = norm_and_reshape(test_x)
    
    num_classes = len(LABEL_MAP)
    y_train_encoded = tf.keras.utils.to_categorical(train_y, num_classes)
    y_val_encoded = tf.keras.utils.to_categorical(val_y, num_classes)
    y_test_encoded = tf.keras.utils.to_categorical(test_y, num_classes)

    # 在验证与测试集上预注入强噪声 (-8dB)，模拟真实恶劣工业环境
    x_val_noisy = add_noise(x_val_final, snr_ratio=-8)
    x_test_noisy = add_noise(x_test_final, snr_ratio=-8)

    def augment_batch(features, labels):
        """
        在线随机数据增强：包含循环平移、瞬态冲击模拟及混合信噪比干扰。
        """
        prng = np.random.default_rng()
        augmented_x = features.copy()
        batch_n, step_n, _ = augmented_x.shape

        # 随机时域循环平移
        for idx in range(batch_n):
            pixel_shift = prng.integers(0, step_n)
            augmented_x[idx, :, 0] = np.roll(augmented_x[idx, :, 0], pixel_shift)

        # 瞬态异常脉冲注入
        if prng.random() > 0.9: 
            for idx in range(batch_n):
                if prng.random() > 0.5: 
                    spike_count = prng.integers(1, 3) 
                    spike_pos = prng.integers(0, step_n, spike_count)
                    pulse_amp = np.std(augmented_x[idx]) * prng.uniform(1.5, 2.5) 
                    augmented_x[idx, spike_pos, 0] += pulse_amp * prng.choice([-1, 1], size=spike_count)

        # 宽动态范围信噪比扰动
        if prng.random() > 0.5: 
            augmented_x = add_noise(augmented_x, snr_ratio=(-8, 8))

        return augmented_x.astype(np.float32), labels.astype(np.float32)

    def set_tensor_shape(f_tensor, l_tensor):
        f_tensor.set_shape([None, input_length, 1])
        l_tensor.set_shape([None, num_classes])
        return f_tensor, l_tensor

    # 构建高性能 I/O 流水线
    training_pipeline = tf.data.Dataset.from_tensor_slices((x_train_final.astype('float32'), y_train_encoded))
    training_pipeline = training_pipeline.shuffle(len(train_x)).batch(64)
    training_pipeline = training_pipeline.map(
        lambda x, y: tf.numpy_function(augment_batch, [x, y], [tf.float32, tf.float32]),
        num_parallel_calls=tf.data.AUTOTUNE
    ).map(set_tensor_shape).prefetch(tf.data.AUTOTUNE)

    # 模型实例化与编译
    drsn_model = DRSN_CW(num_classes=num_classes)
    drsn_model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), 
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    logging.info("基于 DRSN 的诊断系统已就绪。分类数: %d, 采样窗口: %d", num_classes, input_length)
    
    # 训练监控回调函数
    training_monitors = [
        tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=7, min_lr=1e-6, verbose=1),
        tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=20, restore_best_weights=True)
    ]

    # 启动随机梯度下降训练过程
    drsn_model.fit(
        training_pipeline,
        epochs=100,
        validation_data=(x_val_noisy, y_val_encoded),
        callbacks=training_monitors,
        verbose=2
    )

    # 强噪声环境下最终评估
    _, final_acc = drsn_model.evaluate(x_test_noisy, y_test_encoded, verbose=0)
    print("\n" + "="*50)
    print("最终性能报告：在 -8dB SNR 极强噪声环境下，系统识别准确率为: {0:.2f}%".format(final_acc * 100))
    print("="*50)

# =============================================================================
# 程序入口
# =============================================================================

if __name__ == "__main__":
    # 默认数据源搜索路径
    DEFAULT_DATA_DIR = os.path.join(os.getcwd(), 'data_path')
    
    if not os.path.exists(DEFAULT_DATA_DIR):
        logging.warning("缺省路径未发现数据集: {0}".format(DEFAULT_DATA_DIR))
        user_input_dir = input("请手动指定包含 CWRU .mat 文件的目录路径: ").strip()
        if user_input_dir:
            DEFAULT_DATA_DIR = user_input_dir
        else:
            logging.error("未输入有效路径，程序终止运行。")
            sys.exit(0)

    # 执行主程序逻辑
    run_pipeline(DEFAULT_DATA_DIR, input_length=1024)

当对验证集与测试集全局叠加-8 dB的强噪声（即背景噪声功率大于原始信号功率）时，交叉熵损失函数可稳定收敛，测试集的最终分类准确率能够稳定达到90%以上的水平。这一测试结果基本验证了软阈值模块在特征去噪方面的有效性。后续工作可考虑结合设备运行的物理机理构造先验损失函数，并在变负载、跨设备等复杂工况下进一步验证模型的泛化能力。

方法出处：
Zhao M, Zhong S, Fu X, Tang B, Pecht M. Deep residual shrinkage networks for fault diagnosis. IEEE Transactions on Industrial Informatics. 2020, 16(7): 4681-4690.
https://ieeexplore.ieee.org/document/8850096