PyTorch实现深度可分离卷积(以MobileNet为例)

介绍深度可分离卷积之前首先要先介绍一下分组卷积。

分组卷积

参考链接[1]已经将分组卷积介绍的比较详细了,这里就不再赘述。原理可以参考一下参考链接。但是这篇文章对于分组卷积的具体代码讲解,感觉不太好。这里对于分组卷积着重介绍一下代码讲解。

直接调用torch.nn.Conv2d()就可以实现分组卷积。在Conv2d里面有一个参数叫groups,这个参数是实现这一功能的关键。下面详细解介绍一下这个参数。

groups

通过groups来实现分组卷积的操作。groups默认为1,也就是说将输入分为一组,此时是常规卷积。groups数值为几表示的就是将输入通道分为几组。当groups=in_channels的时候,表示的就是将输入的每一个通道都作为一组,然后分别对其进行卷积,输出通道数为k,最后再将每组输出串联,最后通道数为in_channels*k

为了详细说明这一点,我们来举几个具体例子。为了接下来方便计算,这里bias=False,另外用到了torchsummary模块,可以通过pip安装。

import torch
from torchsummary import summary
import torch.nn as nn

class CSDN_Tem(nn.Module):
    def __init__(self, in_ch, out_ch, groups):
        super(CSDN_Tem, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=groups,
            bias=False
        )

    def forward(self, input):
        out = self.conv(input)
        return out

测试时,使用的是64*64*3的输入,你可以想象为一个64*64大小的图片输入。一个3*3大小的卷积核处理,最后输出的shape是64*64*30。

我们先是group为1的常规卷积操作。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             810
================================================================
Total params: 810
Trainable params: 810
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

一共有3*3(kernel size)*3(输入图片的shape)*30(输出图片的shape)=810个参数。

下面是groups=in_channels。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 3).to(device)
print(summary(conv,  input_size=(3, 64, 64)))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             270
================================================================
Total params: 270
Trainable params: 270
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

可以发现,参数量明显的减少了。根据参考[1]的原理,过去直接卷积,卷积核大小是3*3*3,要对输入图片整体映射输出30维。所以最后算的是810个参数。但是当分组为3的时候,实际上是对输入图片的三个通道进行伞分组,最后输出也是三分组。也就是对于30维的输出而言,实际上是输入的一个通道对应的10维输出。这样的话,实际我们需要3*3*1的卷积核就可以。那么这样参数量就变成了3*3*1*(10+10+10)=270,同样的输出维度,参数却整整减小了3倍。

逐点卷积

逐点卷积就是用1*1的卷积核来对分组卷积之后的卷积网络来进行卷积最后输出成原来的样子。可以看参考资料[2]有比较详细的介绍。

深度可分离卷积

深度可分离卷积分成depthwise convolution和pointwise convolution。其中深度卷积是不只是要求group和input_channel一致。还要求input_channel和output_channel一致。也就是说对每一个通道都要单独进行卷积操作。我们举例来看一下:

需要特别注意的是图片的长和宽要始终保持一致才行。

class CSDN_Tem(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, padding, groups):
        super(CSDN_Tem, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=kernel_size,
            stride=1,
            padding=padding,
            groups=groups,
            bias=False
        )

    def forward(self, input):
        out = self.conv(input)
        return out

常规卷积是这样的:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 3, 1, 1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             810
================================================================
Total params: 810
Trainable params: 810
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

深度可分离卷积分成两步,第一步先是depthwise convolution:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 3, 3, padding=1, groups=3).to(device)
print(summary(conv,  input_size=(3, 64, 64)))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 3, 64, 64]              27
================================================================
Total params: 27
Trainable params: 27
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.09
Params size (MB): 0.00
Estimated Total Size (MB): 0.14
----------------------------------------------------------------

之后是pointwise convolution:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, kernel_size=1, padding=0, groups=1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]              90
================================================================
Total params: 90
Trainable params: 90
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.98
----------------------------------------------------------------

达到了同样的输出效果。但是参数量却基本减少了8到9倍。

参考

[1] 深度可分离卷积(Depthwise Separable Convolution)和分组卷积(Group Convolution)的理解,相互关系及PyTorch实现
[2] 卷积神经网络中的Separable Convolution

发布了189 篇原创文章 · 获赞 233 · 访问量 36万+

猜你喜欢

转载自blog.csdn.net/Einstellung/article/details/103585835
今日推荐