在旋转机械的运行过程中，采集到的振动信号往往夹杂着大量的环境噪声。在强噪声干扰下，传统的深度神经网络经常难以提取出判别性强的故障特征。论文“Deep Residual Shrinkage Networks for Fault Diagnosis”针对这一痛点，提出了一种深度残差收缩网络（DRSN）。其核心思想是将软阈值化作为非线性变换层集成到残差单元中，通过自动学习阈值来消除不重要的噪声特征，从而在极低信噪比条件下依然保持较高的诊断准确率。

一、逐通道收缩块（RSBU-CW）的结构解析

论文提出的核心模块是“逐通道阈值的残差收缩构建单元”（RSBU-CW）。如图所示，该单元在传统残差路径的基础上增加了一个子网络用于确定阈值。该子网络首先对输入特征图进行绝对值处理，通过全局平均池化（GAP）获取通道统计量。通过两个全连接层和 Sigmoid 得到缩放参数 ，最后计算出每个通道的自适应阈值：$\tau =\alpha \cdot \text{average}(|x|)$。

软阈值化的数学表达为：$y=\mathrm{sign}(x)\cdot \max (|x|-\tau ,0)$ 。这种机制让网络能够根据输入信号的实时噪声水平，动态地”收缩”掉近零的干扰成分，尽量保留真正的冲击特征。

二、实验数据与预处理

本复现所采用的数据源自经典的西储大学（CWRU）轴承数据集。实验涵盖了正常状态、内圈故障、外圈故障及滚动体故障等 10 类健康状态。

为了模拟真实工业场景中的恶劣环境，向原始信号中注入了 -8 dB 的加性高斯白噪声。数据处理采用滑动窗口策略，将一维振动信号切分为长度为 1024 的样本序列。在进入模型前，所有样本均经过标准化处理，以加速模型收敛并提升训练稳定性。

三、代码实现与结果评估

以下是基于 TensorFlow 框架复现的 DRSN-CW 完整代码。

# ==============================================================================
# 基于深度残差收缩网络 (DRSN-CW) 的故障诊断模型实现与验证
#
# 本代码基于 TensorFlow 框架，复现了论文《Deep Residual Shrinkage Networks
# for Fault Diagnosis》中提出的逐通道阈值变体（DRSN-CW）。该实现专注于在
# 高噪声环境下，通过自适应特征收缩机制提升振动信号的故障诊断性能。
# 注意：本复现使用指定路径下的 CWRU 格式数据集进行流程验证。
#
# 核心功能模块：
# 1. 信号引擎：实现原始信号的加载、基于时域的滑动窗口切片操作，并集成
#    可控参数的高斯白噪声注入功能。
# 2. 核心架构单元：构建了残差收缩构建单元 (RSBU-CW)，该单元集成通道
#    注意力子网络，以自动学习并应用软阈值，从而抑制冗余噪声特征。
# 3. 训练与评估：训练流程内置在线数据增强策略，并设计了极低信噪比
#    (-8dB SNR) 干扰下的模型鲁棒性检验流程。
#
# 论文引用格式：
# 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.
# ==============================================================================

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

# =============================================================================
# 硬件与运行的环境初始化
# =============================================================================

def init_env():
    """
    配置深度学习执行环境。
    主要任务：验证核心依赖库，过滤冗余日志，并针对硬件加速器启用显存动态分配。
    """
    logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')
    
    # 验证关键第三方科学计算库
    try:
        import sklearn, scipy
    except ImportError as error:
        logging.critical("运行环境缺失核心库依赖: {}".format(error))
        sys.exit(1)

    # 抑制 TensorFlow 非关键性冗余日志
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    
    # 获取物理加速设备列表
    accelerators = tf.config.list_physical_devices('GPU')
    if accelerators:
        try:
            for device in accelerators:
                tf.config.experimental.set_memory_growth(device, True)
            logging.info("发现 {} 个计算加速设备，已优化显存分配策略。".format(len(accelerators)))
        except RuntimeError as dev_err:
            logging.warning("加速设备配置失败: {}".format(dev_err))
    else:
        logging.info("未发现可用 GPU，系统切换至 CPU 计算模式。")

# 执行环境初始化
init_env()

# =============================================================================
# 数据集构建模块
# =============================================================================

class CWRUDataLoader:
    """
    实验数据加载模块。
    负责从原始磁盘文件提取机械振动信号，并执行基于滑动时间窗口的样本特征切分。
    """
    def __init__(self, data_root, window_size=1024):
        """
        :param data_root: 原始 .mat 文件存放目录
        :param window_size: 单次采样的时间窗口长度
        """
        self.directory_path = os.path.abspath(data_root)
        self.window_size = window_size
        self.sampling_stride = window_size

    def _load_mat(self, file_path):
        """
        解析 CWRU 格式的 .mat 文件，提取指定驱动端（DE）的采样序列。
        """
        try:
            raw_content = sio.loadmat(file_path)
            for identifier in raw_content.keys():
                if 'DE_time' in identifier:
                    return raw_content[identifier].flatten()
        except Exception:
            return None
        return None

    def load_data(self, FAULT_LABEL_MAP):
        """
        遍历数据集目录，构建用于模型训练与评估的特征矩阵及对应的故障类别标签。
        :param FAULT_LABEL_MAP: 故障类别与文件名的映射字典
        :return: (特征矩阵, 标签集合)
        """
        feature_collection, label_collection = [], []
        is_data_present = False
        
        for class_id, file_names in FAULT_LABEL_MAP.items():
            for name in file_names:
                full_path = os.path.join(self.directory_path, "{}.mat".format(name))
                if not os.path.exists(full_path):
                    continue
                
                signal_stream = self._load_mat(full_path)
                if signal_stream is None:
                    continue
                
                is_data_present = True
                # 滑动窗口切片逻辑
                for offset in range(0, len(signal_stream) - self.window_size + 1, self.sampling_stride):
                    clip = signal_stream[offset : offset + self.window_size]
                    feature_collection.append(clip)
                    label_collection.append(class_id)
        
        if not is_data_present:
            raise FileNotFoundError("在目标路径下未检索到匹配的信号文件: {}".format(self.directory_path))
            
        return np.array(feature_collection, dtype='float32'), np.array(label_collection, dtype='int32')

def add_awgn(input_signals, target_snr):
    """
    向输入信号注入指定信噪比的高斯白噪声。
    数学公式: P_noise = P_signal / (10^(SNR/10))
    """
    input_signals = np.array(input_signals)
    random_state = np.random.default_rng()
    
    # 确定信噪比参数
    snr_val = target_snr if isinstance(target_snr, (int, float)) else random_state.uniform(target_snr[0], target_snr[1])
    
    # 计算功率谱与噪声幅值
    signal_power = np.mean(input_signals**2, axis=1, keepdims=True)
    noise_variance = signal_power / (10**(snr_val / 10))
    noise_amplitude = random_state.normal(0, np.sqrt(noise_variance), input_signals.shape)
    
    return (input_signals + noise_amplitude).astype('float32')

# =============================================================================
# 深度残差收缩网络 (DRSN) 架构实现
# =============================================================================

class SoftThresholding(layers.Layer):
    """
    软阈值算子层。
    实现基于自适应学习阈值的软阈值算子，旨在自动消除不重要的噪声特征。
    """
    def __init__(self, **kwargs):
        super(SoftThresholding, self).__init__(**kwargs)

    def call(self, inputs):
        """
        执行收缩变换：y = sign(x) * max(|x| - τ, 0)
        """
        latent_features, threshold = inputs
        # 维度对齐处理
        threshold_expanded = tf.expand_dims(threshold, axis=1)
        return tf.sign(latent_features) * tf.maximum(tf.abs(latent_features) - threshold_expanded, 0.0)

class RSBU_CW(layers.Layer):
    """
    逐个通道阈值的残差收缩构建单元（RSBU-CW）。
    集成子网络以自动确定通道级的自适应学习阈值。
    """
    def __init__(self, out_channels, kernel_size, stride=1, **kwargs):
        super(RSBU_CW, self).__init__(**kwargs)
        self.out_channels = out_channels
        self.stride = stride
        self.kernel_size = kernel_size
        self.weight_decay = regularizers.l2(1e-4)

        self.bypass = None
        
        # 级联卷积层定义
        self.norm_a = layers.BatchNormalization()
        self.relu_a = layers.Activation('relu')
        self.conv_a = layers.Conv1D(out_channels, kernel_size, strides=stride, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        self.norm_b = layers.BatchNormalization()
        self.relu_b = layers.Activation('relu')
        self.conv_b = layers.Conv1D(out_channels, kernel_size, strides=1, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        # 阈值学习子网络
        self.pooling = layers.GlobalAveragePooling1D()
        self.dense_a = layers.Dense(out_channels, kernel_initializer='he_normal')
        self.norm_dense = layers.BatchNormalization()
        self.relu_dense = layers.Activation('relu')
        self.dense_b = layers.Dense(out_channels, activation='sigmoid')
        self.threshold_op = SoftThresholding()

    def build(self, input_shape):
        """
        构建捷径分支。
        """
        if self.stride != 1 or input_shape[-1] != self.out_channels:
            self.bypass = models.Sequential([
                layers.Conv1D(self.out_channels, 1, strides=self.stride, padding='same'),
            ])
        super(RSBU_CW, self).build(input_shape)

    def call(self, x_input):
        """
        前向推理流程：执行恒等映射残差连接，并应用通道注意力生成的阈值进行特征收缩。
        """
        shortcut_connection = x_input
        if self.bypass:
            shortcut_connection = self.bypass(x_input)

        # 特征提取路径
        feat = self.norm_a(x_input)
        feat = self.relu_a(feat)
        feat = self.conv_a(feat)
        feat = self.norm_b(feat)
        feat = self.relu_b(feat)
        feat = self.conv_b(feat)

        # 计算通道级自适应阈值
        abs_feat = tf.abs(feat)
        global_stat = self.pooling(abs_feat)
        
        alpha = self.dense_a(global_stat)
        alpha = self.norm_dense(alpha)
        alpha = self.relu_dense(alpha)
        alpha = self.dense_b(alpha)
        
        tau = tf.multiply(alpha, global_stat)
        
        # 应用软阈值化并合并残差
        denoised_feat = self.threshold_op([feat, tau])
        return layers.Add()([denoised_feat, shortcut_connection])

class DRSN_CW(models.Model):
    """
    深度残差收缩网络的完整架构。
    通过堆叠残差收缩单元，实现在强噪声干扰下的端到端故障诊断。
    """
    def __init__(self, num_classes):
        super(DRSN_CW, self).__init__(name="DRSN_CW_Model")
        
        # 入口层
        self.stem_conv = layers.Conv1D(32, 15, strides=2, padding='same', kernel_initializer='he_normal')
        self.stem_norm = layers.BatchNormalization()
        self.stem_act = layers.Activation('relu')
        
        # 深度残差收缩块堆叠
        self.stage_blocks = [
            RSBU_CW(32, 5, stride=2),
            RSBU_CW(32, 5, stride=1),
            RSBU_CW(64, 5, stride=2),
            RSBU_CW(64, 5, stride=1),
            RSBU_CW(128, 5, stride=2),
            RSBU_CW(128, 5, stride=1)
        ]
        
        # 决策层
        self.head_norm = layers.BatchNormalization()
        self.head_act = layers.Activation('relu')
        self.global_avg = layers.GlobalAveragePooling1D()
        self.output_layer = layers.Dense(num_classes, activation='softmax')

    def call(self, network_input):
        """
        执行深度特征提取。
        """
        h = self.stem_conv(network_input)
        h = self.stem_norm(h)
        h = self.stem_act(h)
        
        for stage in self.stage_blocks:
            h = stage(h)
            
        h = self.head_norm(h)
        h = self.head_act(h)
        h = self.global_avg(h)
        return self.output_layer(h)

# =============================================================================
# 训练引擎与评估流水线
# =============================================================================

def train(data_dir, window_len=1024):
    """
    执行完整的模型训练与测试评估流程。
    """
    
    # 故障标签映射定义
    FAULT_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']
    }
    
    data_loader = CWRUDataLoader(data_root=data_dir, window_size=window_len)
    
    try:
        raw_x, raw_y = data_loader.load_data(FAULT_LABEL_MAP)
    except Exception as data_err:
        logging.error("数据预处理阶段异常: {}".format(data_err))
        return

    # 训练、验证、测试集分割
    tr_x, temp_x, tr_y, temp_y = train_test_split(raw_x, raw_y, test_size=0.3, random_state=42)
    val_x_raw, ts_x_raw, val_y_raw, ts_y_raw = train_test_split(temp_x, temp_y, test_size=0.5, random_state=42)
    
    # 执行标准化处理 (Standardization)
    avg, std_dev = np.mean(tr_x), np.std(tr_x)
    
    def preprocess(data):
        return ((data - avg) / std_dev).reshape(-1, window_len, 1)

    train_set_x = preprocess(tr_x)
    val_set_x = preprocess(val_x_raw)
    test_set_x = preprocess(ts_x_raw)
    
    # 标签向量化
    num_classes = len(FAULT_LABEL_MAP)
    train_set_y = tf.keras.utils.to_categorical(tr_y, num_classes).astype('float32')
    val_set_y = tf.keras.utils.to_categorical(val_y_raw, num_classes).astype('float32')
    test_set_y = tf.keras.utils.to_categorical(ts_y_raw, num_classes).astype('float32')

    # 构建极端噪声测试集。通过注入 -8dB 的加性噪声，模拟恶劣工况下的传感器采集信号。
    val_set_noisy = add_awgn(val_set_x, target_snr=-8)
    test_set_noisy = add_awgn(test_set_x, target_snr=-8)

    def augment_data(features, labels):
        """
        应用动态数据增强，提升深度残差收缩网络的泛化性能。
        包含：时域循环平移、模拟瞬时冲击干扰、以及多量级的噪声合成。
        """
        gen = np.random.default_rng()
        augmented_x = features.copy()
        count, length, _ = augmented_x.shape

        # 变换1：随机时间轴循环移位
        for idx in range(count):
            offset = gen.integers(0, length)
            augmented_x[idx, :, 0] = np.roll(augmented_x[idx, :, 0], offset)

        # 变换2：稀疏瞬时冲击干扰
        if gen.random() > 0.9: 
            for idx in range(count):
                if gen.random() > 0.5: 
                    impulse_count = gen.integers(1, 3) 
                    indices = gen.integers(0, length, impulse_count)
                    magnitude = np.std(augmented_x[idx]) * gen.uniform(1.5, 2.5) 
                    augmented_x[idx, indices, 0] += magnitude * gen.choice([-1, 1], size=impulse_count)

        # 变换3：概率性添加混合信噪比噪声
        if gen.random() > 0.5: 
            augmented_x = add_awgn(augmented_x, target_snr=(-8, 8))

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

    def set_shapes(f_tensor, l_tensor):
        f_tensor.set_shape([None, window_len, 1])
        l_tensor.set_shape([None, num_classes])
        return f_tensor, l_tensor

    # 构建 TensorFlow 异步数据流水线
    train_ds = tf.data.Dataset.from_tensor_slices((train_set_x.astype('float32'), train_set_y))
    train_ds = train_ds.shuffle(len(train_set_x)).batch(64)
    train_ds = train_ds.map(
        lambda x, y: tf.numpy_function(augment_data, [x, y], [tf.float32, tf.float32]),
        num_parallel_calls=tf.data.AUTOTUNE
    ).map(set_shapes).prefetch(tf.data.AUTOTUNE)

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

    logging.info("基于“深度残差收缩网络”的诊断系统初始化完毕。类别数: {}, 窗口规模: {}".format(num_classes, window_len))
    
    # 训练优化策略：包含动态学习率衰减（基于验证集损失）以及防止过拟合的提前停止机制。
    learning_callbacks = [
        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)
    ]

    # 启动模型拟合
    model.fit(
        train_ds,
        epochs=100,
        validation_data=(val_set_noisy, val_set_y),
        callbacks=learning_callbacks,
        verbose=2
    )

    # 性能评估：在极端高噪声工业场景（-8dB SNR）下对训练完毕的模型进行泛化性能测试。
    final_loss, final_acc = model.evaluate(test_set_noisy, test_set_y, verbose=0)
    print("\n高噪声工业场景 (-8dB SNR) 下的识别准确率: {:.2f}%".format(final_acc * 100))

# =============================================================================
# 程序入口执行
# =============================================================================

if __name__ == "__main__":
    # 预设数据搜索路径
    data_path = os.path.join(os.getcwd(), 'bearing_data')
    
    if not os.path.exists(data_path):
        logging.info("默认数据路径未找到: {}".format(data_path))
        user_input_path = input("请输入 CWRU 数据集所在的本地文件夹路径: ").strip()
        if user_input_path:
            data_path = user_input_path
        else:
            logging.error("路径输入为空，任务终止。")
            sys.exit(0)

    # 开启任务工作流
    train(data_path, window_len=1024)

在验证环节引入了极端噪声干扰，从实验结果图中可以看到，在 -8dB 信噪比的严苛条件下，模型经过 100 个周期的训练，测试集准确率超过了 90% 。这初步体现了深度残差收缩网络在复杂噪声环境下的特征自动提取与降噪能力。

论文标题: Deep residual shrinkage networks for fault diagnosis
出版期刊: IEEE Transactions on Industrial Informatics. 2020, 16(7): 4681-4690.
DOI: 10.1109/TII.2019.2943898

​