234 lines
7.2 KiB
Python
234 lines
7.2 KiB
Python
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import thop
|
|
|
|
# try:
|
|
# import softpool_cuda
|
|
# from SoftPool import soft_pool2d, SoftPool2d
|
|
# except ImportError:
|
|
# print('Please install SoftPool first: https://github.com/alexandrosstergiou/SoftPool')
|
|
# exit(0)
|
|
|
|
NET_CONFIG = {
|
|
# k, in_c, out_c, s, use_se
|
|
"blocks2": [[3, 16, 32, 1, False]],
|
|
"blocks3": [[3, 32, 64, 2, False], [3, 64, 64, 1, False]],
|
|
"blocks4": [[3, 64, 128, 2, False], [3, 128, 128, 1, False]],
|
|
"blocks5": [[3, 128, 256, 2, False], [5, 256, 256, 1, False],
|
|
[5, 256, 256, 1, False], [5, 256, 256, 1, False],
|
|
[5, 256, 256, 1, False], [5, 256, 256, 1, False]],
|
|
"blocks6": [[5, 256, 512, 2, True], [5, 512, 512, 1, True]]
|
|
}
|
|
|
|
|
|
def autopad(k, p=None):
|
|
if p is None:
|
|
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
|
|
return p
|
|
|
|
|
|
def make_divisible(v, divisor=8, min_value=None):
|
|
if min_value is None:
|
|
min_value = divisor
|
|
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
|
|
if new_v < 0.9 * v:
|
|
new_v += divisor
|
|
return new_v
|
|
|
|
|
|
class HardSwish(nn.Module):
|
|
def __init__(self, inplace=True):
|
|
super(HardSwish, self).__init__()
|
|
self.relu6 = nn.ReLU6(inplace=inplace)
|
|
|
|
def forward(self, x):
|
|
return x * self.relu6(x+3) / 6
|
|
|
|
|
|
class HardSigmoid(nn.Module):
|
|
def __init__(self, inplace=True):
|
|
super(HardSigmoid, self).__init__()
|
|
self.relu6 = nn.ReLU6(inplace=inplace)
|
|
|
|
def forward(self, x):
|
|
return (self.relu6(x+3)) / 6
|
|
|
|
|
|
class SELayer(nn.Module):
|
|
def __init__(self, channel, reduction=16):
|
|
super(SELayer, self).__init__()
|
|
self.avgpool = nn.AdaptiveAvgPool2d(1)
|
|
self.fc = nn.Sequential(
|
|
nn.Linear(channel, channel // reduction, bias=False),
|
|
nn.ReLU(inplace=True),
|
|
nn.Linear(channel // reduction, channel, bias=False),
|
|
HardSigmoid()
|
|
)
|
|
|
|
def forward(self, x):
|
|
b, c, h, w = x.size()
|
|
y = self.avgpool(x).view(b, c)
|
|
y = self.fc(y).view(b, c, 1, 1)
|
|
return x * y.expand_as(x)
|
|
|
|
|
|
class DepthwiseSeparable(nn.Module):
|
|
def __init__(self, inp, oup, dw_size, stride, use_se=False):
|
|
super(DepthwiseSeparable, self).__init__()
|
|
self.use_se = use_se
|
|
self.stride = stride
|
|
self.inp = inp
|
|
self.oup = oup
|
|
self.dw_size = dw_size
|
|
self.dw_sp = nn.Sequential(
|
|
nn.Conv2d(self.inp, self.inp, kernel_size=self.dw_size, stride=self.stride,
|
|
padding=autopad(self.dw_size, None), groups=self.inp, bias=False),
|
|
nn.BatchNorm2d(self.inp),
|
|
HardSwish(),
|
|
|
|
nn.Conv2d(self.inp, self.oup, kernel_size=1, stride=1, padding=0, bias=False),
|
|
nn.BatchNorm2d(self.oup),
|
|
HardSwish(),
|
|
)
|
|
self.se = SELayer(self.oup)
|
|
|
|
def forward(self, x):
|
|
x = self.dw_sp(x)
|
|
if self.use_se:
|
|
x = self.se(x)
|
|
return x
|
|
|
|
|
|
class PP_LCNet(nn.Module):
|
|
def __init__(self, scale=1.0, class_num=10, class_expand=1280, dropout_prob=0.2):
|
|
super(PP_LCNet, self).__init__()
|
|
self.scale = scale
|
|
self.conv1 = nn.Conv2d(3, out_channels=make_divisible(16 * self.scale),
|
|
kernel_size=3, stride=2, padding=1, bias=False)
|
|
# k, in_c, out_c, s, use_se inp, oup, dw_size, stride, use_se=False
|
|
self.blocks2 = nn.Sequential(*[
|
|
DepthwiseSeparable(inp=make_divisible(in_c * self.scale),
|
|
oup=make_divisible(out_c * self.scale),
|
|
dw_size=k, stride=s, use_se=use_se)
|
|
for i, (k, in_c, out_c, s, use_se) in enumerate(NET_CONFIG["blocks2"])
|
|
])
|
|
|
|
self.blocks3 = nn.Sequential(*[
|
|
DepthwiseSeparable(inp=make_divisible(in_c * self.scale),
|
|
oup=make_divisible(out_c * self.scale),
|
|
dw_size=k, stride=s, use_se=use_se)
|
|
for i, (k, in_c, out_c, s, use_se) in enumerate(NET_CONFIG["blocks3"])
|
|
])
|
|
|
|
self.blocks4 = nn.Sequential(*[
|
|
DepthwiseSeparable(inp=make_divisible(in_c * self.scale),
|
|
oup=make_divisible(out_c * self.scale),
|
|
dw_size=k, stride=s, use_se=use_se)
|
|
for i, (k, in_c, out_c, s, use_se) in enumerate(NET_CONFIG["blocks4"])
|
|
])
|
|
# k, in_c, out_c, s, use_se inp, oup, dw_size, stride, use_se=False
|
|
self.blocks5 = nn.Sequential(*[
|
|
DepthwiseSeparable(inp=make_divisible(in_c * self.scale),
|
|
oup=make_divisible(out_c * self.scale),
|
|
dw_size=k, stride=s, use_se=use_se)
|
|
for i, (k, in_c, out_c, s, use_se) in enumerate(NET_CONFIG["blocks5"])
|
|
])
|
|
|
|
self.blocks6 = nn.Sequential(*[
|
|
DepthwiseSeparable(inp=make_divisible(in_c * self.scale),
|
|
oup=make_divisible(out_c * self.scale),
|
|
dw_size=k, stride=s, use_se=use_se)
|
|
for i, (k, in_c, out_c, s, use_se) in enumerate(NET_CONFIG["blocks6"])
|
|
])
|
|
|
|
self.GAP = nn.AdaptiveAvgPool2d(1)
|
|
|
|
self.last_conv = nn.Conv2d(in_channels=make_divisible(NET_CONFIG["blocks6"][-1][2] * scale),
|
|
out_channels=class_expand,
|
|
kernel_size=1, stride=1, padding=0, bias=False)
|
|
|
|
self.hardswish = HardSwish()
|
|
self.dropout = nn.Dropout(p=dropout_prob)
|
|
|
|
self.fc = nn.Linear(class_expand, class_num)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
print(x.shape)
|
|
x = self.blocks2(x)
|
|
print(x.shape)
|
|
x = self.blocks3(x)
|
|
print(x.shape)
|
|
x = self.blocks4(x)
|
|
print(x.shape)
|
|
x = self.blocks5(x)
|
|
print(x.shape)
|
|
x = self.blocks6(x)
|
|
print(x.shape)
|
|
|
|
x = self.GAP(x)
|
|
x = self.last_conv(x)
|
|
x = self.hardswish(x)
|
|
x = self.dropout(x)
|
|
x = torch.flatten(x, start_dim=1, end_dim=-1)
|
|
x = self.fc(x)
|
|
return x
|
|
|
|
|
|
def PPLCNET_x0_25(**kwargs):
|
|
model = PP_LCNet(scale=0.25, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x0_35(**kwargs):
|
|
model = PP_LCNet(scale=0.35, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x0_5(**kwargs):
|
|
model = PP_LCNet(scale=0.5, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x0_75(**kwargs):
|
|
model = PP_LCNet(scale=0.75, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x1_0(**kwargs):
|
|
model = PP_LCNet(scale=1.0, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x1_5(**kwargs):
|
|
model = PP_LCNet(scale=1.5, **kwargs)
|
|
return model
|
|
|
|
|
|
def PPLCNET_x2_0(**kwargs):
|
|
model = PP_LCNet(scale=2.0, **kwargs)
|
|
return model
|
|
|
|
def PPLCNET_x2_5(**kwargs):
|
|
model = PP_LCNet(scale=2.5, **kwargs)
|
|
return model
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# input = torch.randn(1, 3, 640, 640)
|
|
# model = PPLCNET_x2_5()
|
|
# flops, params = thop.profile(model, inputs=(input,))
|
|
# print('flops:', flops / 1000000000)
|
|
# print('params:', params / 1000000)
|
|
|
|
model = PPLCNET_x1_0()
|
|
# model_1 = PW_Conv(3, 16)
|
|
input = torch.randn(2, 3, 256, 256)
|
|
print(input.shape)
|
|
output = model(input)
|
|
print(output.shape) # [1, num_class]
|
|
|