深度残差收缩网络PyTorch代码复现（转载）

本文为转载，首发于知乎

深度残差收缩网络（DRSN）理解、PyTorch代码复现 - 去否的文章 - 知乎 https://zhuanlan.zhihu.com/p/485134830

首先附上原文地址和github链接：

原文：Deep Residual Shrinkage Networks for Fault Diagnosis | IEEE Journals & Magazine

github：GitHub - zhao62/Deep-Residual-Shrinkage-Networks: The deep residual shrinkage network is a variant of deep residual networks.

研究进展 

 大量的深度学习方法被用于机器故障诊断中。例如，Ince使用一维卷积神经网络从电流信号中实施诊断电机故障。Shao将一个深度信念卷积神经网络用于电机轴承的故障诊断中。Ma等人使用具有解调时间-频率特征的ResNet来诊断非平稳运行条件下的行星齿轮箱。Zhao等人使用ResNet融合多组小波包系数进行故障诊断。

以往方法缺点

当面对大量高噪声振动信号时，ResNet的学习能力会有所下降。因为ResNet使用卷积核作为局部特征提取器，由于噪声的干扰，有可能无法检测到故障相关的特征，进而输出层学习到的高层次特征判别力较差，不足以实现故障的准确分类。因此，针对强背景噪声下的旋转电机基于振动的故障诊断，提出一种新的深度学习方法是必要的。

本文创新

本文提出了两种深度残差收缩网络（DRSN），一种DRSN使用通道共享阈值（DRSN-CS），另一种使用通道wise阈值（DRSN-CW），在最终达到较高的准确率的目标下，提升ResNet从高噪声振动信号中提取特征的能力。主要的贡献有： 1. 软阈值被插入深度结构中，作为非线性转换层，目的是有效地排除噪声相关的特征。 2. 阈值通过特殊设计的子网络来适应性的获得，因此每个振动信号都有自己的一系列阈值。 3. 在软阈值处理中，考虑了两种阈值，即通道共享阈值和通道wise阈值。

详细说明

软阈值

将信号软阈值化在过去20年来一直是很多信号去噪方法的关键步骤。一般来说，原始信号会被转换到一个范围内，在该范围内，接近0的数字不太重要，因此软阈值化就是将该接近0的特征变为0.例如，小波软阈值化作为一个经典的去噪方法，由三个部分组成：小波分解、软阈值化、小波重构。为了保证良好的信号去噪效果，在小波阈值化中一个关键的任务是设计一个滤波器，该滤波器可以将有用的信息转换为正负值很大的特征，并且将噪声信息转换为接近0的特征。然而设计这样的一个滤波器需要很多信号处理领域专业的知识，长久以来是一个挑战。深度学习提供了一个新的方法解决该问题。深度学习通过使用梯度下降算法自动学习过滤器，而不是让专家人工设计。因此，将软阈值化和深度学习组合起来是一种有前途的消除噪声相关信息和构建有判别力特征的方法。软阈值化方法如下所示：

其中x表示输入特征，y表示输出特征，为阈值。软阈值化将接近0的数转换为0，不像ReLU将负数变为0，因此有用的负数特征也可以被保存。可以看到该软阈值化函数的导数不是0就是1，可以很好的防止梯度消失和梯度爆炸问题，如下式所示：

软阈值化的示意图如下所示：

 在传统信号去噪方中，为阈值设置一个合适的值是困难的。针对这个问题，提出的DRSN方法会在深度结构中自动的获取该值，以避免人工操作带来的错误。

 DRSN网络

DRSN-CS

DRSN-CS是ResNet网络的一个变种，使用软阈值化来消除噪声相关的特征。软阈值化作为一个非线性转换层被插入结构单元中。同时，该阈值可以在结构单元中学习得到。该结构单元命名为（RSBU-CS，即Residual shrinkage building unit）。不同于RBU（Residual building unit，ResNet网络中的基本单元），RSBU-CS有一个特殊的模块用于估计阈值。在该模块中，GAP被用于特征图x的绝对值中，来得到一维向量。随后，该一维向量会传入两层全连接层来获得尺度参数。在两层全连接层后会使用一个sigmoid函数，此时尺度参数会被转换为（0，1）中的某个值。该尺度参数随后会与特征图|x|的平均值相乘来得到阈值。这一步是考虑到阈值不仅需要是正数，而且不能太大，否则软阈值化后的特征将全为0.RSBU-CS和DRSN-CS的网络结构如下图所示：

DRSN-CW 

DRSN-CW是ResNet的另一个变种，DRSN-CS将一个阈值应用到了全部通道，而DRSN-CW为每一个通道都设置了一个阈值。特征图x被GAP降维成了一个一维向量，随后输入两个全连接层。第二个全连接层输出的神经元数量为输入特征图通道的数量，随后每一个神经元都进行sigmoid化。最后特定通道的尺度值与输入特征图的特定通道平均值相乘，得到每一个通道的阈值。该结构如下所示：

实验 

超参数设置

 前40个epoch的lr=0.1，中间40个epoch的lr=0.01，最后20个epoch的lr=0.001.使用的优化器为Momentum，系数为0.9.使用L2正则化，系数为0.0001，mini batch-size为128.

PyTorch代码实现

在具体实现上，很多编码上的细节并没有在论文中说明，比如特征图在一个RSBU经过下采样后，shape变为原来的一半，此时如何进行shortcut；再比如，经过卷积后通道数增长一倍，此时如何进行shortcut。这些细节需要阅读代码才能知道，很遗憾的是作者只公开了Keras和Tensorflow的代码，并没有编写Pytorch版本的代码，因此我在这里补上Pytorch版本的代码，供大家参考，如果有哪里不对的地方，烦请大家指正。

class RSBU_CW(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, down_sample=False):
        super().__init__()
        self.down_sample = down_sample
        self.in_channels = in_channels
        self.out_channels = out_channels
        stride = 1
        if down_sample:
            stride = 2
        self.BRC = nn.Sequential(
            nn.BatchNorm1d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                      padding=1),
            nn.BatchNorm1d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv1d(in_channels=out_channels, out_channels=out_channels, kernel_size=kernel_size, stride=1,
                      padding=1)
        )
        self.global_average_pool = nn.AdaptiveAvgPool1d(1)
        self.FC = nn.Sequential(
            Linear(in_features=out_channels, out_features=out_channels),
            BatchNorm1d(out_channels),
            ReLU(inplace=True),
            Linear(in_features=out_channels, out_features=out_channels),
            Sigmoid()
        )
        self.flatten = Flatten()
        self.average_pool = AvgPool1d(kernel_size=1, stride=2)
 
    def forward(self, input):
        x = self.BRC(input)
        x_abs = torch.abs(x)
        gap = self.global_average_pool(x_abs)
        gap = self.flatten(gap)
        alpha = self.FC(gap)
        threshold = torch.mul(gap, alpha)
        threshold = torch.unsqueeze(threshold, 2)
        # 软阈值化
        sub = x_abs - threshold
        zeros = sub - sub
        n_sub = torch.max(sub, zeros)
        x = torch.mul(torch.sign(x), n_sub)  
        if self.down_sample:  # 如果是下采样，则对输入进行平均池化下采样
            input = self.average_pool(input)
        if self.in_channels != self.out_channels:  # 如果输入的通道和输出的通道不一致，则进行padding,直接通过复制拼接矩阵进行padding,原代码是通过填充0
            zero_padding=torch.zeros(input.shape).cuda()
            input = torch.cat((input, zero_padding), dim=1)
 
        result = x + input
        return result
 
 
class DRSNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = Conv1d(in_channels=1, out_channels=4, kernel_size=3, stride=2, padding=1)
        self.bn = BatchNorm1d(16)
        self.relu = ReLU()
        self.softmax = Softmax(dim=1)
        self.global_average_pool = nn.AdaptiveAvgPool1d(1)
        self.flatten = Flatten()
        self.linear6_8 = Linear(in_features=256, out_features=128)
        self.linear8_4 = Linear(in_features=128, out_features=64)
        self.linear4_2 = Linear(in_features=64, out_features=32)
        self.output_center_pos = Linear(in_features=32, out_features=1)
        self.output_width = Linear(in_features=32, out_features=1)
 
        self.linear = Linear(in_features=16, out_features=8)
        self.output_class = Linear(in_features=8, out_features=3)
 
    def forward(self, input):  # 1*256
        x = self.conv1(input)  # 4*128
        x = RSBU_CW(in_channels=4, out_channels=4, kernel_size=3, down_sample=True).cuda()(x)  # 4*64
        x = RSBU_CW(in_channels=4, out_channels=4, kernel_size=3, down_sample=False).cuda()(x)  # 4*64
        x = RSBU_CW(in_channels=4, out_channels=8, kernel_size=3, down_sample=True).cuda()(x)  # 8*32
        x = RSBU_CW(in_channels=8, out_channels=8, kernel_size=3, down_sample=False).cuda()(x)  # 8*32
        x = RSBU_CW(in_channels=8, out_channels=16, kernel_size=3, down_sample=True).cuda()(x)  # 16*16
        x = RSBU_CW(in_channels=16, out_channels=16, kernel_size=3, down_sample=False).cuda()(x)  # 16*16
        x = self.bn(x)
        x = self.relu(x)
        gap = self.global_average_pool(x)  # 16*1
        gap = self.flatten(gap)  # 1*16
        linear1 = self.linear(gap)  # 1*8
        output_class = self.output_class(linear1)  # 1*3
        output_class = self.softmax(output_class)  # 1*3
 
        return output_class