总结几个好用的CNN模块（Pytorch）

r囧r小猫 2023-01-10 14:45 285阅读 0赞

总结几个比较好的CNN模块。

*  # SEBlock #

代码：

class SEBlock(nn.Module):
    
    
    
        def __init__(self, input_channels, internal_neurons):
    
            super(SEBlock, self).__init__()
    
            self.down = nn.Conv2d(in_channels=input_channels, out_channels=internal_neurons, kernel_size=1, stride=1,
    
                                  bias=True, padding_mode='same')
    
            self.up = nn.Conv2d(in_channels=internal_neurons, out_channels=input_channels, kernel_size=1, stride=1,
    
                                bias=True, padding_mode='same')
    
    
    
        def forward(self, inputs):
    
            x = F.avg_pool2d(inputs, kernel_size=inputs.size(3))
    
            x = self.down(x)
    
            x = F.leaky_relu(x)
    
            x = self.up(x)
    
            x = F.sigmoid(x)
    
            x = x.repeat(1, 1, inputs.size(2), inputs.size(3))
    
            return inputs * x

*  # ACBlock #

代码

class CropLayer(nn.Module):
    
    
    
        #   E.g., (-1, 0) means this layer should crop the first and last rows of the feature map. And (0, -1) crops the first and last columns
    
        def __init__(self, crop_set):
    
            super(CropLayer, self).__init__()
    
            self.rows_to_crop = - crop_set[0]
    
            self.cols_to_crop = - crop_set[1]
    
            assert self.rows_to_crop >= 0
    
            assert self.cols_to_crop >= 0
    
    
    
        def forward(self, input):
    
            if self.rows_to_crop == 0 and self.cols_to_crop == 0:
    
                return input
    
            elif self.rows_to_crop > 0 and self.cols_to_crop == 0:
    
                return input[:, :, self.rows_to_crop:-self.rows_to_crop, :]
    
            elif self.rows_to_crop == 0 and self.cols_to_crop > 0:
    
                return input[:, :, :, self.cols_to_crop:-self.cols_to_crop]
    
            else:
    
                return input[:, :, self.rows_to_crop:-self.rows_to_crop, self.cols_to_crop:-self.cols_to_crop]
    
    class ACBlock(nn.Module):
    
    
    
        def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=1, dilation=1, groups=1,
    
                     padding_mode='same', deploy=False,
    
                     use_affine=True, reduce_gamma=False, use_last_bn=False, gamma_init=None):
    
            super(ACBlock, self).__init__()
    
            self.deploy = deploy
    
            if deploy:
    
                self.fused_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
    
                                            kernel_size=(kernel_size, kernel_size), stride=stride,
    
                                            padding=padding, dilation=dilation, groups=groups, bias=True,
    
                                            padding_mode=padding_mode)
    
            else:
    
                self.square_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
    
                                             kernel_size=(kernel_size, kernel_size), stride=stride,
    
                                             padding=padding, dilation=dilation, groups=groups, bias=False,
    
                                             padding_mode=padding_mode)
    
                self.square_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
    
    
    
                center_offset_from_origin_border = padding - kernel_size // 2
    
                ver_pad_or_crop = (padding, center_offset_from_origin_border)
    
                hor_pad_or_crop = (center_offset_from_origin_border, padding)
    
                if center_offset_from_origin_border >= 0:
    
                    self.ver_conv_crop_layer = nn.Identity()
    
                    ver_conv_padding = ver_pad_or_crop
    
                    self.hor_conv_crop_layer = nn.Identity()
    
                    hor_conv_padding = hor_pad_or_crop
    
                else:
    
                    self.ver_conv_crop_layer = CropLayer(crop_set=ver_pad_or_crop)
    
                    ver_conv_padding = (0, 0)
    
                    self.hor_conv_crop_layer = CropLayer(crop_set=hor_pad_or_crop)
    
                    hor_conv_padding = (0, 0)
    
                self.ver_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(kernel_size, 1),
    
                                          stride=stride,
    
                                          padding=ver_conv_padding, dilation=dilation, groups=groups, bias=False,
    
                                          padding_mode=padding_mode)
    
    
    
                self.hor_conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, kernel_size),
    
                                          stride=stride,
    
                                          padding=hor_conv_padding, dilation=dilation, groups=groups, bias=False,
    
                                          padding_mode=padding_mode)
    
                self.ver_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
    
                self.hor_bn = nn.BatchNorm2d(num_features=out_channels, affine=use_affine)
    
    
    
                if reduce_gamma:
    
                    assert not use_last_bn
    
                    self.init_gamma(1.0 / 3)
    
    
    
                if use_last_bn:
    
                    assert not reduce_gamma
    
                    self.last_bn = nn.BatchNorm2d(num_features=out_channels, affine=True)
    
    
    
                if gamma_init is not None:
    
                    assert not reduce_gamma
    
                    self.init_gamma(gamma_init)
    
    
    
        def init_gamma(self, gamma_value):
    
            init.constant_(self.square_bn.weight, gamma_value)
    
            init.constant_(self.ver_bn.weight, gamma_value)
    
            init.constant_(self.hor_bn.weight, gamma_value)
    
            print('init gamma of square, ver and hor as ', gamma_value)
    
    
    
        def single_init(self):
    
            init.constant_(self.square_bn.weight, 1.0)
    
            init.constant_(self.ver_bn.weight, 0.0)
    
            init.constant_(self.hor_bn.weight, 0.0)
    
            print('init gamma of square as 1, ver and hor as 0')
    
    
    
        def forward(self, input):
    
            if self.deploy:
    
                return self.fused_conv(input)
    
            else:
    
                square_outputs = self.square_conv(input)
    
                square_outputs = self.square_bn(square_outputs)
    
                vertical_outputs = self.ver_conv_crop_layer(input)
    
                vertical_outputs = self.ver_conv(vertical_outputs)
    
                vertical_outputs = self.ver_bn(vertical_outputs)
    
                horizontal_outputs = self.hor_conv_crop_layer(input)
    
                horizontal_outputs = self.hor_conv(horizontal_outputs)
    
                horizontal_outputs = self.hor_bn(horizontal_outputs)
    
                result = square_outputs + vertical_outputs + horizontal_outputs
    
                if hasattr(self, 'last_bn'):
    
                    return self.last_bn(result)
    
                return result

*  #  eca\_layer #

class eca_layer(nn.Module):
        """Constructs a ECA module.
    
        Args:
            channel: Number of channels of the input feature map
            k_size: Adaptive selection of kernel size
        """
        def __init__(self, channel, k_size=3):
            super(eca_layer, self).__init__()
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
            self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False) 
            self.sigmoid = nn.Sigmoid()
    
        def forward(self, x):
            # x: input features with shape [b, c, h, w]
            b, c, h, w = x.size()
    
            # feature descriptor on the global spatial information
            y = self.avg_pool(x)
    
            # Two different branches of ECA module
            y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
    
            # Multi-scale information fusion
            y = self.sigmoid(y)
    
            return x * y.expand_as(x)

*  # ChannelAttention #

class ChannelAttention(nn.Module):
        def __init__(self, in_planes, ratio=16):
            super(ChannelAttention, self).__init__()
            self.avg_pool = nn.AdaptiveAvgPool2d(1)
            self.max_pool = nn.AdaptiveMaxPool2d(1)
            self.fc1 = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
            self.relu1 = nn.LeakyReLU(negative_slope=0.01, inplace=False)
            self.fc2 = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)
    
            self.sigmoid = nn.Sigmoid()
    
        def forward(self, x):
            avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
            max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
            out = avg_out + max_out
            return self.sigmoid(out)

*  #  ConvBN #

class ConvBN(nn.Sequential):
        def __init__(self, in_planes, out_planes, kernel_size, stride=1, groups=1):
            if not isinstance(kernel_size, int):
                padding = [(i - 1) // 2 for i in kernel_size]
            else:
                padding = (kernel_size - 1) // 2
            super(ConvBN, self).__init__(OrderedDict([
                ('conv', nn.Conv2d(in_planes, out_planes, kernel_size, stride,
                                   padding=padding, groups=groups, bias=False)),
                ('bn', nn.BatchNorm2d(out_planes)),
                #('Mish', Mish())
                ('Mish', nn.LeakyReLU(negative_slope=0.3, inplace=False))
            ]))

*  # ResBlock #
    
        class ResBlock(nn.Module):
            """
            Sequential residual blocks each of which consists of \
            two convolution layers.
            Args:
                ch (int): number of input and output channels.
                nblocks (int): number of residual blocks.
                shortcut (bool): if True, residual tensor addition is enabled.
            """
        
            def __init__(self, ch, nblocks=1, shortcut=True):
                super().__init__()
                self.shortcut = shortcut
                self.module_list = nn.ModuleList()
                for i in range(nblocks):
                    resblock_one = nn.ModuleList()
                    resblock_one.append(ConvBN(ch, ch, 1))
                    resblock_one.append(Mish())
                    resblock_one.append(ConvBN(ch, ch, 3))
                    resblock_one.append(Mish())
                    self.module_list.append(resblock_one)
        
            def forward(self, x):
                for module in self.module_list:
                    h = x
                    for res in module:
                        h = res(h)
                    x = x + h if self.shortcut else h
                return x