Inverted residuals block、Shufflenet block的Pytorch实现

---

0. 目前卷积的三种实现方式：

---

图源：https://blog.csdn.net/tintinetmilou/article/details/81607721

Basic Convolution

Depthwise Convolution

Pointwise Convolution

1. MobileNet v2

---

MobileNetv2采用先升维度再降维度的方法，先将通道维度上升到原来的t倍，然后进行depthwise conv，最后再通过1x1 conv将通道数为恢复到原来维度。

class mobilev2_block(nn.Module):
	def __init__(self, in_channels, out_channels, ratio=6, stride=1):
		super(mobilev2_block, self).__init__()
		self.in_channels = in_channels
		self.out_channels = out_channels
		self.stride = stride
		self.ratio = ratio
		
		if stride == 1:
			assert out_channels == in_channels, "Under the stride=1 input and output channel number should be equal."
		mid_channels = in_channels*ratio
		self.branch = nn.Sequential(
			nn.Conv2d(in_channels, mid_channels, kernel_size=1, bias=False),
			nn.BatchNorm2d(mid_channels),
			nn.ReLU6(inplace=True),

			nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=stride, padding=1, groups=mid_channels, bias=False),
			
			nn.Conv2d(mid_channels, out_channels, kernel_size=1, bias=False),
			nn.BatchNorm2d(out_channels),
			nn.ReLU6(inplace=True),

		)
	
	def forward(self, x):
		if self.stride == 1:
			x = x + self.branch(x)
		elif self.stride == 2:
			x = self.branch(x)
		return x

fake_data = torch.randn(10, 32, 128, 128)
net = mobilev2_block(in_channels=32, out_channels=64, stride=2)
print(net(fake_data).size())

2. ShuffleNet v2

---

上图右边两个结构是shufflenet v2的基础残差块，图（c）的输入和输出的特征图完全一致：NxCxHxW->NxCxHxW；图（d）输入和输出的特征图宽高变为原来的1/2，同时通道数可以为之前的1,2,3倍。

import random
import torch
import torch.nn as nn


class BasicConv2d(nn.Module):
	def __init__(self, in_channels, out_channels, **kwargs):
		super(BasicConv2d, self).__init__()
		self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
		self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
		self.relu = nn.ReLU(inplace=True) # 不确定是ReLU还是relu6

	def forward(self, x):
		x = self.conv(x)
		x = self.bn(x)

		return self.relu(x)

class Identity_layer(nn.Module):
	"""
	恒等映射
	"""
	def __init__(self):
		super(Identity_layer, self).__init__()
		pass

	def forward(self, x):
		return x


class shufflev2_block(nn.Module):
	def __init__(self, in_channels, out_channels, mode=1, group=2):
		super(shufflev2_block, self).__init__()
		self.group = group
		self.in_channels = in_channels
		self.mode = mode
		self.out_channels = out_channels
		
		# HxW->HxW
		if self.mode == 1: 
			assert in_channels == out_channels, "Under the MODE 1 input and output channel number should be equal."
			mid_channels = self.in_channels // group
			self.branch1 = Identity_layer()
			self.branch2 = nn.Sequential(
				# pw
				BasicConv2d(mid_channels, mid_channels, kernel_size=1),
				# dw
				nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, groups=mid_channels, bias=False),
				nn.BatchNorm2d(mid_channels),
				# pw linner
				BasicConv2d(mid_channels, mid_channels, kernel_size=1)
			)
		# HxW->H/2xW/2
		elif mode == 2:
			mid_channels = self.out_channels // group
			self.branch1 = nn.Sequential(
				# dw
				nn.Conv2d(self.in_channels, self.in_channels, kernel_size=3, stride=2, padding=1, groups=self.in_channels, bias=False),
				nn.BatchNorm2d(self.in_channels),
				# pw linner
				BasicConv2d(self.in_channels, mid_channels, kernel_size=1)

			)
			self.branch2 = nn.Sequential(
				# pw
				BasicConv2d(self.in_channels, mid_channels, kernel_size=1),
				# dw
				nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=2, padding=1, groups=mid_channels, bias=False),
				nn.BatchNorm2d(mid_channels),
				# pw linner
				BasicConv2d(mid_channels, mid_channels, kernel_size=1)
			)
		
	def forward(self, x):
		if self.mode == 1:
			channel_per_group = self.in_channels // self.group
			x = torch.cat([self.branch1(x[:, :channel_per_group]), self.branch2(x[:, channel_per_group:])], dim=1)
		elif self.mode == 2:
			x = torch.cat([self.branch1(x), self.branch2(x)], dim=1)
		return self._shuffle(x)
	

	def _shuffle(self, x):
		channel = x.size(1)
		# 随机切换通道排序
		shuffle_list = random.sample(range(channel), channel)
		return x[:, shuffle_list]


fake_data = torch.randn(10, 512, 128, 128)
net = shufflev2_block(in_channels=512, out_channels=512, mode=1)
print(net(fake_data).size())

可以通过nn.conv2d()内部的groups参数轻松进行depthwise conv操作。

参考

---

【知乎】《轻量化神经网络综述》：https://zhuanlan.zhihu.com/p/45496826

【CSDN】Depthwise卷积与Pointwise卷积:https://blog.csdn.net/tintinetmilou/article/details/81607721

【GitHub】Shufflenet-v2-Pytorch

【Arxiv】ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design

【Arxiv】MobileNetV2: Inverted Residuals and Linear Bottlenecks