def __init__(self, CONFIG):
super(Supernet, self).__init__()
self.CONFIG = CONFIG
self.classes = CONFIG.classes
self.dataset = CONFIG.dataset
if self.dataset[:5] == "cifar":
first_stride = 1
elif self.dataset == "imagenet" or self.dataset == "imagenet_lmdb":
first_stride = 2
self.first = ConvBNRelu(
input_channel=3,
output_channel=32,
kernel=3,
stride=first_stride,
pad=3 // 2,
activation="relu")
input_channel = 32
output_channel = 16
self.first_mb = MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=1,
kernels=[3],
stride=1,
activation="relu",
min_expansion=1,
split_block=1,
se=False)
input_channel = output_channel
self.stages = nn.ModuleList()
for l_cfg in self.CONFIG.l_cfgs:
min_expansion = self.CONFIG.min_expansion
expansion, output_channel, kernels, stride, split_block, se = l_cfg
self.stages.append(MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=expansion,
kernels=kernels,
stride=stride,
activation="relu",
min_expansion=min_expansion,
split_block=expansion,
se=se,
search=True))
input_channel = output_channel
self.last_stage = conv_1x1_bn(input_channel, 1280)
self.classifier = nn.Linear(1280, self.classes)
self.split_block = split_block
self.architecture_param_nums = len(
self.CONFIG.l_cfgs) * (self.CONFIG.split_blocks * self.CONFIG.kernels_nums)
self._initialize_weights()
def _calculate_basic_operations(self, write_to_file=None):
model_flops = 0
if self.CONFIG.dataset[:5] == "cifar":
last_input_size = 4
first = nn.Sequential(
ConvBNRelu(input_channel=3,
output_channel=32,
kernel=3,
stride=1,
pad=3 // 2,
activation="relu"),
MBConv(input_channel=32,
output_channel=16,
min_expansion=1,
expansion=1,
kernels=[3],
stride=1,
activation="relu",
split_block=1,
se=False))
elif self.CONFIG.dataset[:8] == "imagenet":
last_input_size = 7
first = nn.Sequential(
ConvBNRelu(input_channel=3,
output_channel=32,
kernel=3,
stride=2,
pad=3 // 2,
activation="relu"),
MBConv(input_channel=32,
output_channel=16,
min_expansion=1,
expansion=1,
kernels=[3],
stride=1,
activation="relu",
split_block=1,
se=False))
model_flops += self._calculate_flops(first, 3, self.CONFIG.input_size)
last_stage = nn.Sequential(conv_1x1_bn(320, 1280), Flatten(),
nn.Linear(1280, self.CONFIG.classes))
model_flops += self._calculate_flops(last_stage, 320, last_input_size)
return model_flops
def __init__(self, l_cfgs=BASIC_CFGS, dataset="imagenet", classes=1000):
super(Model, self).__init__()
if dataset[:5] == "cifar":
self.first = ConvBNRelu(input_channel=3, output_channel=32, kernel=3, stride=1,
pad=3//2, activation="relu")
elif dataset == "imagenet":
self.first = ConvBNRelu(input_channel=3, output_channel=32, kernel=3, stride=2,
pad=3//2, activation="relu")
input_channel = 32
output_channel = 16
self.first_mb = MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=1,
kernels=[3],
stride=1,
activation="relu",
split_block=1,
se=False)
input_channel = output_channel
self.stages = nn.ModuleList()
for l_cfg in l_cfgs:
expansion, output_channel, kernel, stride, split_block, se = l_cfg
self.stages.append(MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=expansion,
kernels=kernel,
stride=stride,
activation="relu",
split_block=split_block,
se=se))
input_channel = output_channel
self.last_stage = conv_1x1_bn(input_channel, 1280)
self.classifier = nn.Linear(1280, classes)
self._initialize_weights()
def __init__(self, CONFIG):
super(Supernet, self).__init__()
self.CONFIG = CONFIG
self.classes = CONFIG.classes
self.dataset = CONFIG.dataset
if self.dataset[:
5] == "cifar": # 有cifar10和cifar100，对于cifar系列(前五个字符)，第一层的conv k3 first stride是1
first_stride = 1
elif self.dataset == "imagenet" or self.dataset == "imagenet_lmdb":
first_stride = 2
# ConvK3
self.first = ConvBNRelu(input_channel=3,
output_channel=32,
kernel=3,
stride=first_stride,
pad=3 // 2,
activation="relu")
input_channel = 32
output_channel = 16
# MB Expansion Ratio 1 K3 初始化
self.first_mb = MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=1,
kernels=[3],
stride=1,
activation="relu",
min_expansion=1,
split_block=1,
se=False)
input_channel = output_channel
# 候选骨干网存储list
self.stages = nn.ModuleList()
# 19个unified blocks初始化
for l_cfg in self.CONFIG.l_cfgs:
min_expansion = self.CONFIG.min_expansion
expansion, output_channel, kernels, stride, split_block, se = l_cfg
self.stages.append(
MBConv(
input_channel=input_channel,
output_channel=output_channel,
expansion=expansion, # 先按最大扩张率6创建子模块
kernels=kernels,
stride=stride,
activation="relu",
min_expansion=min_expansion,
split_block=expansion,
se=se,
search=True))
input_channel = output_channel
# conv K1
self.last_stage = conv_1x1_bn(input_channel, 1280)
# 1280分类器
self.classifier = nn.Linear(1280, self.classes)
self.split_block = split_block # sub-blocks数量
self.architecture_param_nums = len(self.CONFIG.l_cfgs) * (
self.CONFIG.split_blocks * self.CONFIG.kernels_nums)
self._initialize_weights()
def __init__(self, se=False, activation="relu", bn_momentum=0.1, l_cfgs_name="SGNAS_A", input_size=224, classes=1000, seg_state=False):
super(Model, self).__init__()
if l_cfgs_name == "SGNAS_A":
l_cfgs = SGNAS_A
elif l_cfgs_name == "SGNAS_B":
l_cfgs = SGNAS_B
elif l_cfgs_name == "SGNAS_C":
l_cfgs = SGNAS_C
if input_size == 32:
self.first = ConvBNRelu(input_channel=3, output_channel=32, kernel=3, stride=1,
pad=3//2, activation=activation, bn_momentum=bn_momentum)
elif input_size >= 224:
self.first = ConvBNRelu(input_channel=3, output_channel=32, kernel=3, stride=2,
pad=3//2, activation=activation, bn_momentum=bn_momentum)
input_channel = 32
output_channel = 16
self.first_mb = MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=1,
kernels=[3],
stride=1,
activation=activation,
split_block=1,
se=se,
bn_momentum=bn_momentum)
input_channel = output_channel
self.stages = nn.ModuleList()
for l_cfg in l_cfgs:
expansion, output_channel, kernel, stride, split_block, _ = l_cfg
self.stages.append(MBConv(input_channel=input_channel,
output_channel=output_channel,
expansion=expansion,
kernels=kernel,
stride=stride,
activation=activation,
split_block=split_block,
se=se,
bn_momentum=bn_momentum))
input_channel = output_channel
self.last_stage = conv_1x1_bn(input_channel, 1280, activation=activation, bn_momentum=bn_momentum)
self.classifier_new = nn.Sequential(
nn.Sequential(),
nn.Linear(1280, classes))
# ============== Segmentation ==================
self.seg_state = seg_state
if self.seg_state:
self.segmentation_conv_1 = nn.Sequential(
ConvBNRelu(input_channel=1280,
output_channel=320,
kernel=3,
stride=1,
pad=3//2),
nn.Conv2d(320, 1, 1)
)
self.segmentation_conv_2 = nn.Sequential(
ConvBNRelu(input_channel=96,
output_channel=96,
kernel=3,
stride=1,
pad=3//2),
nn.Conv2d(96, 1, 1)
)
self.segmentation_conv_3 = nn.Sequential(
ConvBNRelu(input_channel=40,
output_channel=40,
kernel=3,
stride=1,
pad=3//2),
nn.Conv2d(40, 1, 1)
)
# ==============================================
self._initialize_weights()