面向工业噪声环境的深度残差收缩网络TensorFlow代码复现

引言

在工业现场环境中,通过加速度传感器采集到的旋转机械振动信号往往包含大量的背景噪声。传统的深度学习模型,如残差网络(Residual Network, ResNet),虽然在特征提取上表现出色,但在处理强噪声信号时,在区分故障特征与噪声成分时存在一定困难。

本文主要复现论文 “Deep Residual Shrinkage Networks for Fault Diagnosis” 的核心思想。针对工业诊断中的噪声瓶颈,该研究提出了一种改进的深度学习架构——深度残差收缩网络(DRSN)。其核心创新在于将“软阈值化(Soft Thresholding)”作为非线性层引入残差模块,并利用注意力机制(Attention Mechanism)自动确定阈值。这种设计使得模型能够在特征图中自动消除不重要的噪声特征,尽量保留与故障相关的辨识性特征,从而在极低信噪比(Signal-to-Noise Ratio, SNR)环境下依然保持较高的诊断精度。

一、网络结构图解与原理

深度残差收缩网络(DRSN)最关键的改进是提出了带有通道阈值的残差收缩构建模块(Residual Shrinkage Building Unit with Channel-wise thresholds, RSBU-CW),其结构如下图所示:

图1

该模块的运作逻辑如下:

  1. 特征提取:首先通过两次批标准化(Batch Normalization, BN)、修正线性单元(Rectified Linear Unit, ReLU)激活和卷积(Convolution, Conv)操作提取基础特征。
  2. 子网络确定阈值:特征图经过全局平均池化(Global Average Pooling, GAP)后进入一个两层的全连接层(Fully Connected layer, FC),输出一个介于 0 到 1 之间的缩放参数 $\alpha$。
  3. 软阈值化:通过公式 $\tau = \alpha \times \text{average}(|x|)$ 自动计算每个通道的阈值 $\tau$。软阈值化函数将绝对值小于 $\tau$ 的特征置为零,并对大于 $\tau$ 的特征进行收缩。
  4. 恒等映射:处理后的特征与原始输入通过恒等连接(Identity Shortcut)相加,有效缓解了深层网络中的梯度消失问题。

二、代码复现实现

以下是基于 TensorFlow 2.x 框架复现的 DRSN 核心模块(RSBU-CW)及模型构建代码。代码中实现了自动学习阈值的子网络逻辑。

# =============================================================================
# 本程序致力于复现 Zhao 等人在 IEEE Transactions on Industrial Informatics 上
# 发表的深度残差收缩网络(DRSN)方法,构建一个端到端的滚动轴承故障诊断验证框架。
# 代码基于 TensorFlow/Keras 框架,核心实现为通道级自适应软阈值机制的
# 残差收缩构建单元(RSBU-CW)。
#
# 整体工作流包括:
# 1. 计算后端环境的初始化与加速策略配置。
# 2. CWRU 轴承原始.mat 振动信号的解析、滑动窗口特征切片,以及基于训练集的
#    零均值标准化处理。
# 3. 引入在线数据增强策略与指定信噪比高斯噪声注入,以模拟复杂工况。
# 4. 深度残差收缩网络(DRSN-CW)的结构搭建、模型编译与迭代优化。
# 5. 在强噪声背景下对最终模型的诊断准确性和鲁棒性进行量化评估。
#
# 参考文献:
# 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.
# =============================================================================

import os
import sys
import logging
import numpy as np
import scipy.io as matlab_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 RuntimeConfigurator:
    """
    环境预检器:执行后端加速算力的资源调度与依赖性校验
    """
    @staticmethod
    def configure_computation_backend():
        """
        优化分布式计算后端:
        1. 抑制异构计算平台的非关键性冗余诊断输出。
        2. 强制校验科学计算环境的原子性依赖。
        3. 启用显存动态增长机制,防止算力溢出引发的死锁。
        """
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
        
        try:
            import sklearn
            import scipy
        except ImportError as exc:
            logging.critical("运行环境缺失核心科学计算库支持:{}".format(exc))
            sys.exit(1)

        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 显存动态增长。")
            except RuntimeError as err:
                logging.error("计算设备配置失败:{}".format(err))
        else:
            logging.info("未检测到硬件加速设备,系统降级至 CPU 计算模式。")

# 启动环境初始化
RuntimeConfigurator.configure_computation_backend()

# =============================================================================
# 数据加载
# =============================================================================

class CWRULoader:
    """
    信号数据存储库:实现对 CWRU 原始加速度信号的读取、解析与样本重构
    """
    def __init__(self, dataset_root, input_length=1024):
        """
        :param dataset_root: .mat 原始文件的根目录
        :param input_length: 单次采样的时间步长
        """
        self.base_dir = os.path.abspath(dataset_root)
        self.input_length = input_length
        self.hop_length = input_length  # 默认步长与采样长度一致(无重叠)

    def _load_vibration_signal(self, file_path):
        """
        低级接口:解析 MATLAB 格式数据并定位驱动端传感器的时序信号
        """
        try:
            data_dict = matlab_io.loadmat(file_path)
            for identifier in data_dict.keys():
                if 'DE_time' in identifier:
                    return data_dict[identifier].ravel()
        except Exception:
            return None
        return None

    def prepare_samples(self, fault_config):
        """
        构建训练/测试特征库。
        :param fault_config: 映射字典,关联类别标签与文件名称。
        :return: (特征矩阵, 类别向量)
        """
        features, labels = [], []
        is_accessible = False
        
        for class_idx, name_list in fault_config.items():
            for filename in name_list:
                full_path = os.path.join(self.base_dir, "{}.mat".format(filename))
                if not os.path.exists(full_path):
                    continue
                
                signal = self._load_vibration_signal(full_path)
                if signal is None:
                    continue
                
                is_accessible = True
                # 滑动窗口
                num_samples = (signal.size - self.input_length) // self.hop_length + 1
                for i in range(num_samples):
                    offset = i * self.hop_length
                    frame = signal[offset : offset + self.input_length]
                    features.append(frame)
                    labels.append(class_idx)
        
        if not is_accessible:
            raise IOError("在指定路径下未检索到符合规范的数据源:{}".format(self.base_dir))
            
        return np.array(features, dtype='float32'), np.array(labels, dtype='int32')

def apply_snr_noise(x_batch, snr):
    """
    针对输入信号批次施加加性高斯白噪声 (AWGN),用于评估模型的鲁棒性。
    计算公式:P_noise = P_signal / 10^(SNR/10)
    """
    x_batch = np.array(x_batch)
    random_engine = np.random.default_rng()
    
    # 支持固定 SNR 或 SNR 区间采样
    snr_val = snr if isinstance(snr, (int, float)) else \
                 random_engine.uniform(snr[0], snr[1])
    
    # 统计功率并生成对应的噪声张量
    mean_power = np.mean(np.square(x_batch), axis=1, keepdims=True)
    noise_variance = mean_power / (10 ** (snr_val / 10.0))
    noise = random_engine.normal(0, np.sqrt(noise_variance), x_batch.shape)
    
    return (x_batch + noise).astype('float32')

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

class SoftThresholding(layers.Layer):
    """
    非线性缩进层:实现深度残差收缩网络中的核心软阈值去噪功能。
    计算逻辑:y = sign(x) * ReLU(|x| - threshold)
    """
    def __init__(self, **kwargs):
        super(SoftThresholding, self).__init__(**kwargs)

    def call(self, inputs):
        """
        :param inputs: 包含 [原始特征, 自适应阈值向量] 的列表
        """
        x, threshold = inputs
        # 将阈值向量调整为可广播的维度 (Batch, 1, Channels)
        thres_expanded = tf.expand_dims(threshold, axis=1)
        return tf.math.sign(x) * tf.math.maximum(tf.math.abs(x) - thres_expanded, 0.0)

class RSBU_CW(layers.Layer):
    """
    RSBU-CW单元:集成通道注意力机制的残差收缩块。
    通过子网络学习特征通道的收缩比例,自适应滤除强噪声背景下的冗余成分。
    """
    def __init__(self, filters, kernel_size, strides=1, **kwargs):
        super(RSBU_CW, self).__init__(**kwargs)
        self.filters = filters
        self.strides = strides
        self.kernel_size = kernel_size
        self.weight_decay = regularizers.l2(1e-4)

        self.shortcut_path = None
        
        # 特征映射分支
        self.norm_1 = layers.BatchNormalization()
        self.relu_1 = layers.Activation('relu')
        self.conv_1 = layers.Conv1D(filters, kernel_size, strides=strides, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        self.norm_2 = layers.BatchNormalization()
        self.relu_2 = layers.Activation('relu')
        self.conv_2 = layers.Conv1D(filters, kernel_size, strides=1, padding='same', 
                                   kernel_initializer='he_normal', kernel_regularizer=self.weight_decay)
        
        # 阈值学习子网络
        self.global_pool = layers.GlobalAveragePooling1D()
        self.dense_1 = layers.Dense(filters, kernel_initializer='he_normal')
        self.norm_dense = layers.BatchNormalization()
        self.relu_dense = layers.Activation('relu')
        self.dense_2 = layers.Dense(filters, activation='sigmoid')
        self.shrinkage_op = SoftThresholding()

    def build(self, input_shape):
        """
        根据输入输出特征图尺寸自动调整残差连接 (Identity Mapping)
        """
        if self.strides != 1 or input_shape[-1] != self.filters:
            self.shortcut_path = models.Sequential([
                layers.Conv1D(self.filters, 1, strides=self.strides, padding='same'),
            ])
        super(RSBU_CW, self).build(input_shape)

    def call(self, inputs):
        """
        残差前向传播逻辑
        """
        identity = inputs
        if self.shortcut_path:
            identity = self.shortcut_path(inputs)

        # 非线性特征转换
        net = self.norm_1(inputs)
        net = self.relu_1(net)
        net = self.conv_1(net)
        net = self.norm_2(net)
        net = self.relu_2(net)
        net = self.conv_2(net)

        # 阈值自动提取流
        abs_x = tf.abs(net)
        gap_val = self.global_pool(abs_x)
        
        alpha = self.dense_1(gap_val)
        alpha = self.norm_dense(alpha)
        alpha = self.relu_dense(alpha)
        alpha = self.dense_2(alpha)
        
        thres = tf.multiply(alpha, gap_val)
        
        # 执行软阈值收缩与残差融合
        x_denoised = self.shrinkage_op([net, thres])
        return layers.Add()([x_denoised, identity])

class DRSN_CW_Model(models.Model):
    """
    深度残差收缩网络完整分类器。
    堆叠多个 RSBU_CW 模块,通过多层抽象提取高阶判别特征。
    """
    def __init__(self, num_classes):
        super(DRSN_CW_Model, self).__init__(name="DRSN_Classifier_Model")
        
        # 初始预处理层
        self.stem_conv = layers.Conv1D(32, 15, strides=2, padding='same', kernel_initializer='he_normal')
        self.stem_bn = layers.BatchNormalization()
        self.stem_relu = layers.Activation('relu')
        
        # 深度残差收缩模块序列 (逐步下采样并增加通道数)
        self.backbone = [
            RSBU_CW(32, 5, strides=2),
            RSBU_CW(32, 5, strides=1),
            RSBU_CW(64, 5, strides=2),
            RSBU_CW(64, 5, strides=1),
            RSBU_CW(128, 5, strides=2),
            RSBU_CW(128, 5, strides=1)
        ]
        
        # 分类
        self.post_bn = layers.BatchNormalization()
        self.post_relu = layers.Activation('relu')
        self.avg_pool = layers.GlobalAveragePooling1D()
        self.logits = layers.Dense(num_classes, activation='softmax')

    def call(self, inputs):
        """
        端到端前向推理
        """
        x = self.stem_conv(inputs)
        x = self.stem_bn(x)
        x = self.stem_relu(x)
        
        for stage in self.backbone:
            x = stage(x)
            
        x = self.post_bn(x)
        x = self.post_relu(x)
        x = self.avg_pool(x)
        return self.logits(x)

# =============================================================================
# 训练引擎与评估逻辑
# =============================================================================

def run_diagnosis_workflow(dataset_path, input_length=1024):
    """
    主控工作流:执行从原始信号加载到深度残差收缩网络训练的全过程。
    """
    
    # 定义 10 类故障诊断体系
    fault_config = {
        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(dataset_root=dataset_path, input_length=input_length)
    
    try:
        x_raw, y_raw = loader.prepare_samples(fault_config)
    except Exception as failure:
        logging.error("数据流水线加载失败:{}".format(failure))
        return

    # 随机化数据集划分
    train_x, temp_x, train_y, temp_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(temp_x, temp_y, test_size=0.5, random_state=42)
    
    # 零均值标准化 (Z-score Normalization)
    data_mu, data_std = np.mean(train_x), np.std(train_x)
    
    def normalize_tensor(data):
        return ((data - data_mu) / data_std).reshape(-1, input_length, 1)

    train_x = normalize_tensor(train_x)
    val_x = normalize_tensor(val_x)
    test_x = normalize_tensor(test_x)
    
    num_classes = len(fault_config)
    train_y_onehot = tf.keras.utils.to_categorical(train_y, num_classes).astype('float32')
    val_y_onehot = tf.keras.utils.to_categorical(val_y, num_classes).astype('float32')
    test_y_onehot = tf.keras.utils.to_categorical(test_y, num_classes).astype('float32')

    # 为验证与测试集引入 -8dB 的固定强背景噪声
    val_x_noisy = apply_snr_noise(val_x, snr=-8)
    test_x_noisy = apply_snr_noise(test_x, snr=-8)

    def dynamic_augmentation_strategy(features, labels):
        """
        在线数据增强:提升深度残差收缩网络在极端环境下的泛化能力。
        """
        rng = np.random.default_rng()
        aug_x = features.copy()
        batch_size, seq_len, _ = aug_x.shape

        # 变换 1:随机时域循环位移
        for i in range(batch_size):
            step = rng.integers(0, seq_len)
            aug_x[i, :, 0] = np.roll(aug_x[i, :, 0], step)

        # 变换 2:瞬态冲击噪声模拟
        if rng.random() > 0.9: 
            for i in range(batch_size):
                if rng.random() > 0.5: 
                    spike_num = rng.integers(1, 3) 
                    indices = rng.integers(0, seq_len, spike_num)
                    impact_mag = np.std(aug_x[i]) * rng.uniform(1.5, 2.5) 
                    aug_x[i, indices, 0] += impact_mag * rng.choice([-1, 1], size=spike_num)

        # 变换 3:混合动态信噪比噪声
        if rng.random() > 0.5: 
            aug_x = apply_snr_noise(aug_x, snr=(-8, 8))

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

    def _set_dataset_shapes(feat, lab):
        feat.set_shape([None, input_length, 1])
        lab.set_shape([None, num_classes])
        return feat, lab

    # 构建高性能 tf.data 数据流
    train_ds = tf.data.Dataset.from_tensor_slices((train_x.astype('float32'), train_y_onehot))
    train_ds = train_ds.shuffle(train_x.shape[0]).batch(64)
    train_ds = train_ds.map(
        lambda x, y: tf.numpy_function(dynamic_augmentation_strategy, [x, y], [tf.float32, tf.float32]),
        num_parallel_calls=tf.data.AUTOTUNE
    ).map(_set_dataset_shapes).prefetch(tf.data.AUTOTUNE)

    # 编译深度残差收缩模型
    model_instance = DRSN_CW_Model(num_classes=num_classes)
    
    model_instance.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), 
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )

    logging.info("诊断模型准备就绪。目标分类数:{},窗口跨度:{}".format(num_classes, input_length))
    
    # 训练监控回调
    training_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_instance.fit(
        train_ds,
        epochs=100,
        validation_data=(val_x_noisy, val_y_onehot),
        callbacks=training_callbacks,
        verbose=2
    )

    # 最终性能测试(强噪声测试环境)
    eval_loss, eval_acc = model_instance.evaluate(test_x_noisy, test_y_onehot, verbose=0)
    print("\n" + "="*50)
    print("模型最终性能指标 (SNR = -8dB):")
    print("测试准确率: {:.2f}%".format(eval_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("缺省路径不可用:{}".format(DEFAULT_DATA_DIR))
        input_path = input("请手动指定 CWRU 数据集所在的目录路径: ").strip()
        if input_path:
            DEFAULT_DATA_DIR = input_path
        else:
            logging.error("路径为空,系统强制退出。")
            sys.exit(0)

    # 执行主程序
    run_diagnosis_workflow(DEFAULT_DATA_DIR, input_length=1024)

三、实验设置与结果分析

为了验证 DRSN 的性能,设计了针对滚动轴承不同故障部位及损伤尺寸的诊断实验。

1. 数据集描述

数据集包含正常状态及三类典型故障:内圈故障、滚动体故障和外圈故障。每种故障对应三种不同的损伤尺寸(0.007英寸、0.014英寸、0.021英寸),总计 10 类标签。为了模拟极端工业环境,在原始信号中人工加入了高强度的随机噪声。

图2

2. 实验结果

实验在信噪比(SNR)为 -8dB 的极端严苛条件下进行。通过下方训练日志可以看到,随着迭代轮数(Epoch)的增加,模型在验证集上的准确率快速提升,并最终稳定在 90%以上,展现了极强的抗噪性能。

图3

参考文献

  • 论文标题: Deep residual shrinkage networks for fault diagnosis
  • 出版期刊: IEEE Transactions on Industrial Informatics. 2020, 16(7): 4681-4690.
  • DOI: 10.1109/TII.2019.2943898
posted on 2026-03-02 13:23  AAA对方正在偷人  阅读(3)  评论(0)    收藏  举报