使用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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184