pytorch实现Sep conv卷积操作

Sep Conv

这个卷积名字起得花里胡哨的，其实总结起来就是输入通道每个通道一个卷积得到和输入通道数相同的特征图，然后再使用若干个1*1的卷积聚合每个特征图的值得到输出特征图。

假设我们输入通道是16，输出特征图是32，并且使用3*3的卷积提取特征，那么第一步一共需要16*3*3个参数，第二步需要32*16*1*1个参数，一共需要16*3*3+32*16*1*1=656个参数。

使用传统的方法进行卷积一共需要32*16*3*3=4608个参数。
可见Sep Conv在计算量上有明显的优势，下面就是pytorch实现的Sep Conv卷积

import torch
import torch.nn as nn


class SepConv(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride, padding):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channel, in_channel,kernel_size,stride,padding, groups=in_channel)
        self.conv2 = nn.Conv2d(in_channel, out_channel, kernel_size=1,stride=1,padding=0)

    def forward(self, input):
        x = self.conv1(input)
        print(self.conv1.weight.data.shape)
        x = self.conv2(x)
        print(self.conv2.weight.data.shape)
        return x

if __name__ == "__main__":
    input = torch.rand(2,16,600,800)
    Sep = SepConv(16,32,3,1,1)
    output = Sep(input)

输出结果如下