def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
self.use_res_connect = self.stride == 1 and inp == oup
self.conv = nn.Sequential(
# pw
Conv2dWithTime(inp, inp * expand_ratio, 1, 1, 0, bias=False),
BatchNorm2dWithTime(inp * expand_ratio),
ReLU6WithTime(inplace=True),
# dw
Conv2dWithTime(inp * expand_ratio,
inp * expand_ratio,
3,
stride,
1,
groups=inp * expand_ratio,
bias=False),
BatchNorm2dWithTime(inp * expand_ratio),
ReLU6WithTime(inplace=True),
# pw-linear
Conv2dWithTime(inp * expand_ratio, oup, 1, 1, 0, bias=False),
BatchNorm2dWithTime(oup),
)
def _conv_dw(in_channels, out_channels, stride):
return nn.Sequential(
Conv2dWithTime(in_channels=in_channels, out_channels=in_channels,
kernel_size=3, stride=stride, padding=1, groups=in_channels,
bias=False),
BatchNorm2dWithTime(in_channels),
ReLUWithTime(inplace=True),
Conv2dWithTime(in_channels=in_channels, out_channels=out_channels,
kernel_size=1, stride=1, padding=0),
ReLUWithTime(inplace=True)
)
def _conv_dw(in_channels, out_channels, stride, layer_name):
conv_dw = Conv2dWithTime(in_channels=in_channels, out_channels=in_channels,
kernel_size=3, stride=stride, padding=1, groups=in_channels,
bias=False)
bn = BatchNorm2dWithTime(in_channels)
relu_1 = ReLU6WithTime(inplace=True)
conv_pw = Conv2dWithTime(in_channels=in_channels, out_channels=out_channels,
kernel_size=1, stride=1, padding=0)
relu_2 = ReLU6WithTime(inplace=True)
conv_dw.name, bn.name, relu_1.name, conv_pw.name, relu_2.name = layer_name+'_'+conv_dw.layer_type, \
layer_name+'_'+bn.layer_type, \
layer_name+'_'+relu_1.layer_type+'_after_dw', \
layer_name+'_'+conv_pw.layer_type, \
layer_name+'_'+relu_2.layer_type+'_after_pw'
return nn.Sequential(
conv_dw, bn, relu_1, conv_pw, relu_2
)
def __init__(self, inp, oup, stride, expand_ratio):
super(InvertedResidual, self).__init__()
self.stride = stride
assert stride in [1, 2]
self.use_res_connect = self.stride == 1 and inp == oup
# self.conv = nn.Sequential(
# # pw
# Conv2dWithTime(inp, inp * expand_ratio, 1, 1, 0, bias=False),
# BatchNorm2dWithTime(inp * expand_ratio),
# ReLU6WithTime(inplace=True),
# # dw
# Conv2dWithTime(inp * expand_ratio, inp * expand_ratio, 3,
# stride, 1, groups=inp * expand_ratio, bias=False),
# BatchNorm2dWithTime(inp * expand_ratio),
# ReLU6WithTime(inplace=True),
# # pw-linear
# Conv2dWithTime(inp * expand_ratio, oup, 1, 1, 0, bias=False),
# BatchNorm2dWithTime(oup),
# )
self.pw = nn.Sequential(
Conv2dWithTime(inp, inp * expand_ratio, 1, 1, 0, bias=False),
BatchNorm2dWithTime(inp * expand_ratio),
ReLU6WithTime(inplace=True))
self.dw = nn.Sequential(
Conv2dWithTime(inp * expand_ratio,
inp * expand_ratio,
3,
stride,
1,
groups=inp * expand_ratio,
bias=False),
BatchNorm2dWithTime(inp * expand_ratio),
ReLU6WithTime(inplace=True))
self.g = inp if self.use_res_connect else 1
self.pw_linear = nn.Sequential(
Conv2dWithTime(inp * expand_ratio,
oup,
1,
1,
0,
groups=self.g,
bias=False), BatchNorm2dWithTime(oup))
def conv1x1(in_channels, out_channels, groups=1):
"""1x1 convolution with padding
- Normal pointwise convolution When groups == 1
- Grouped pointwise convolution when groups > 1
"""
return Conv2dWithTime(in_channels,
out_channels,
kernel_size=1,
groups=groups,
stride=1)
def _conv_bn(in_channels, out_channels, stride, layer_name):
conv = Conv2dWithTime(in_channels=in_channels, out_channels=out_channels,
kernel_size=3, stride=stride, padding=1, bias=False)
bn = BatchNorm2dWithTime(out_channels)
relu = ReLU6WithTime(inplace=True)
conv.name, bn.name, relu.name = layer_name+'_'+conv.layer_type, \
layer_name+'_'+bn.layer_type, \
layer_name+'_'+relu.layer_type
return nn.Sequential(
conv, bn, relu
)
def conv3x3(in_channels,
out_channels,
stride=1,
padding=1,
bias=True,
groups=1):
"""3x3 convolution with padding
"""
return Conv2dWithTime(in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=padding,
bias=bias,
groups=groups)
def conv_1x1_bn(inp, oup):
return nn.Sequential(Conv2dWithTime(inp, oup, 1, 1, 0, bias=False),
BatchNorm2dWithTime(oup), ReLU6WithTime(inplace=True))
def conv_bn(inp, oup, stride):
return nn.Sequential(Conv2dWithTime(inp, oup, 3, stride, 1, bias=False),
BatchNorm2dWithTime(oup), ReLU6WithTime(inplace=True))