使用pytorch实现深度可分离卷积改进模型的实战实践

这一篇是前面一篇的续集，使用深度可分离进行改进，使得网络尽量能在低端的设备快速运行。先简单解释深度可分离，然后将这个思想应用于改进网络。这样的做法同样可以用深度可分离之外的很多网络结构去修改原来的网络，实现性能或者准确度的目标。

前面的文章在这：

pytorch复现经典生成对抗式的超分辨率网络

深度可分离：

深度可分离分为深度卷积部分和点卷积部分。深度卷积部分使用分组卷积，分组数等于输入的通道数。然后使用$1\times 1$的普通点卷积实现不同通道信息的融合。然后在深度卷积和点卷积中间加上批归一化层和激活层即可。改成深度可分离卷积后的网络如下图所示：

pytorch实现深度可分离的生成模型：

代码中的：class GeneratorDepthwiseBlock(nn.Module):就是深度可分离卷积模块。然后Generator中使用config.G.base_block_type来区别不同的基本模块类型。

import torch
from torch import nn
import torchvision
from torch import Tensor

class GeneratorDepthwiseBlock(nn.Module):
    """
    生成器重复的部分,使用深度可分离卷积的思想
    """

    def __init__(self, channel, kernel_size) -> None:
        super(GeneratorDepthwiseBlock, self).__init__()

        self.channel = channel
        self.conv11 = nn.Conv2d(in_channels=channel, out_channels=channel,
                                kernel_size=(kernel_size, kernel_size),
                                stride=(1, 1), padding=(1, 1), groups=channel)
        self.bn11 = nn.BatchNorm2d(num_features=channel)
        self.p_relu11 = nn.PReLU()
        self.conv12 = nn.Conv2d(in_channels=channel, out_channels=channel,
                                kernel_size=(1, 1), stride=(1, 1))
        self.bn12 = nn.BatchNorm2d(num_features=channel)
        self.p_relu12 = nn.PReLU()

        self.conv21 = nn.Conv2d(in_channels=channel, out_channels=channel,
                                kernel_size=(kernel_size, kernel_size),
                                stride=(1, 1), padding=(1, 1), groups=channel)
        self.bn21 = nn.BatchNorm2d(num_features=channel)
        self.conv22 = nn.Conv2d(in_channels=channel, out_channels=channel,
                                kernel_size=(1, 1), stride=(1, 1))
        self.bn22 = nn.BatchNorm2d(num_features=channel)

    def forward(self, x: Tensor) -> Tensor:
        """
        前向推断
        :param x:
        :return:
        """
        identity = x
        out = self.conv11(x)
        out = self.bn11(out)
        out = self.p_relu11(out)
        out = self.conv12(out)
        out = self.bn12(out)
        out = self.p_relu12(out)

        out = self.conv21(out)
        out = self.bn21(out)
        out = self.conv22(out)
        out = self.bn22(out)
        out += identity
        return out


class GeneratorBasicBlock(nn.Module):
    """
    生成器重复的部分
    """

    def __init__(self, channel, kernel_size) -> None:
        super(GeneratorBasicBlock, self).__init__()

        self.channel = channel
        self.conv1 = nn.Conv2d(in_channels=channel, out_channels=channel,
                               kernel_size=(kernel_size, kernel_size),
                               stride=(1, 1), padding=(1, 1))
        self.bn1 = nn.BatchNorm2d(num_features=channel)
        self.p_relu1 = nn.PReLU()
        self.conv2 = nn.Conv2d(in_channels=channel, out_channels=channel,
                               kernel_size=(kernel_size, kernel_size),
                               stride=(1, 1), padding=(1, 1))
        self.bn2 = nn.BatchNorm2d(num_features=channel)

    def forward(self, x: Tensor) -> Tensor:
        """
        前向推断
        :param x:
        :return:
        """
        identity = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.p_relu1(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out += identity
        return out


class PixelShufflerBlock(nn.Module):
    """
    生成器最后的pixelshuffler
    """

    def __init__(self, in_channel, out_channel) -> None:
        super(PixelShufflerBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
                               kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        self.pixels_shuffle = nn.PixelShuffle(upscale_factor=2)
        self.prelu = nn.PReLU()

    def forward(self, x: Tensor) -> Tensor:
        """
        前向
        """
        out = self.conv1(x)
        out = self.pixels_shuffle(out)
        out = self.prelu(out)
        return out


class Generator(nn.Module):
    """
    生成器
    """

    def __init__(self, config) -> None:
        # Generator parameters
        super(Generator, self).__init__()
        large_kernel_size = config.G.large_kernel_size  # = 9
        small_kernel_size = config.G.small_kernel_size  # = 3
        n_channels = config.G.n_channels  # = 64
        n_blocks = config.G.n_blocks  # = 16
        base_block_type = config.G.base_block_type  # 'depthwise_conv_residual'  # 'conv_residual' or 'depthwise_conv_residual'

        # base block
        if base_block_type == 'depthwise_conv_residual':
            self.repeat_block = GeneratorDepthwiseBlock
        if base_block_type == 'conv_residual':
            self.repeat_block = GeneratorBasicBlock

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=n_channels,
                               kernel_size=(large_kernel_size, large_kernel_size),
                               stride=(1, 1), padding=(4, 4))
        self.prelu1 = nn.PReLU()
        self.B_residul_block = self._make_layer(self.repeat_block, n_channels,
                                                n_blocks, small_kernel_size)
        self.conv2 = nn.Conv2d(in_channels=n_channels, out_channels=n_channels,
                               kernel_size=(small_kernel_size, small_kernel_size),
                               stride=(1, 1), padding=(1, 1))
        self.bn1 = nn.BatchNorm2d(n_channels)
        self.pixel_shuffle_block1 = PixelShufflerBlock(n_channels, 4 * n_channels)
        self.pixel_shuffle_block2 = PixelShufflerBlock(n_channels, 4 * n_channels)
        self.conv3 = nn.Conv2d(in_channels=n_channels, out_channels=3,
                               kernel_size=(large_kernel_size, large_kernel_size),
                               stride=(1, 1), padding=(4, 4))

    def _make_layer(self, base_block, n_channels, n_block, kernel_size) -> nn.Sequential:
        """
        构建重复的B个基本块
        :param base_block: 基本块
        :param n_channels: 块里面的通道数
        :param n_block: 块数
        :return:
        """
        layers = []
        self.base_block = base_block
        for _ in range(n_block):
            layers.append(self.base_block(n_channels, kernel_size))
        return nn.Sequential(*layers)

    def _forward_impl(self, x: Tensor) -> Tensor:
        """
        前向的实现
        """
        out = self.conv1(x)
        out = self.prelu1(out)
        identity = out
        out = self.B_residul_block(out)
        out = self.conv2(out)
        out = self.bn1(out)
        out += identity
        out = self.pixel_shuffle_block1(out)
        out = self.pixel_shuffle_block2(out)
        out = self.conv3(out)

        return out

    def forward(self, x: Tensor) -> Tensor:
        """
        前向
        """
        return self._forward_impl(x)