def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv1x1(self.inplanes, planes * block.expansion,
stride),
nn.BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes))
return nn.Sequential(*layers)
def __init__(self, C_in, C_out, affine=True):
super(FactorizedReduce, self).__init__()
self.relu = nn.ReLU()
self.conv1 = nn.Conv2d(C_in,
C_out // 2,
1,
stride=2,
padding=0,
bias=False)
self.conv2 = nn.Conv2d(C_in,
C_out // 2,
1,
stride=2,
padding=0,
bias=False)
self.bn = nn.BatchNorm2d(C_out, affine=affine)
def __init__(self, kernel_size, in_ch, out_ch, skip, exp_ratio, stride):
super().__init__()
self.kernel_size = kernel_size
self.in_ch = in_ch
self.out_ch = out_ch
self.skip = skip
self.exp_ratio = exp_ratio
self.stride = stride
self.conv = nn.Conv2d(in_ch,
out_ch,
kernel_size,
padding=kernel_size // 2,
stride=stride,
bias=False)
self.relu = nn.ReLU(inplace=False)
self.bn = nn.BatchNorm2d(out_ch, momentum=BN_MOMENTUM)
def __init__(self,
input_size,
in_channels,
channels,
n_classes,
n_layers,
n_nodes=4,
stem_multiplier=3,
auxiliary=False):
super().__init__()
self.in_channels = in_channels
self.channels = channels
self.n_classes = n_classes
self.n_layers = n_layers
self.aux_pos = 2 * n_layers // 3 if auxiliary else -1
c_cur = stem_multiplier * self.channels
self.stem = nn.Sequential(
nn.Conv2d(in_channels, c_cur, 3, 1, 1, bias=False),
nn.BatchNorm2d(c_cur))
# for the first cell, stem is used for both s0 and s1
# [!] channels_pp and channels_p is output channel size, but c_cur is input channel size.
channels_pp, channels_p, c_cur = c_cur, c_cur, channels
self.cells = nn.ModuleList()
reduction_p, reduction = False, False
for i in range(n_layers):
reduction_p, reduction = reduction, False
# Reduce featuremap size and double channels in 1/3 and 2/3 layer.
if i in [n_layers // 3, 2 * n_layers // 3]:
c_cur *= 2
reduction = True
cell = Cell(n_nodes, channels_pp, channels_p, c_cur, reduction_p,
reduction)
self.cells.append(cell)
c_cur_out = c_cur * n_nodes
channels_pp, channels_p = channels_p, c_cur_out
#if i == self.aux_pos:
# self.aux_head = AuxiliaryHead(input_size // 4, channels_p, n_classes)
self.gap = nn.AdaptiveAvgPool2d(1)
self.linear = nn.Linear(channels_p, n_classes)
def __init__(self,
pool_type,
C,
kernel_size,
stride,
padding,
affine=True):
super().__init__()
if pool_type.lower() == 'max':
self.pool = nn.MaxPool2d(kernel_size, stride, padding)
elif pool_type.lower() == 'avg':
self.pool = nn.AvgPool2d(kernel_size,
stride,
padding,
count_include_pad=False)
else:
raise ValueError()
self.bn = nn.BatchNorm2d(C, affine=affine)
def __init__(self,
C_in,
C_out,
kernel_length,
stride,
padding,
affine=True):
super().__init__()
self.net = nn.Sequential(
nn.ReLU(),
nn.Conv2d(C_in,
C_in, (kernel_length, 1),
stride,
padding,
bias=False),
nn.Conv2d(C_in,
C_out, (1, kernel_length),
stride,
padding,
bias=False), nn.BatchNorm2d(C_out, affine=affine))
def __init__(self,
C_in,
C_out,
kernel_size,
stride,
padding,
dilation,
affine=True):
super().__init__()
self.net = nn.Sequential(
nn.ReLU(),
nn.Conv2d(C_in,
C_in,
kernel_size,
stride,
padding,
dilation=dilation,
groups=C_in,
bias=False),
nn.Conv2d(C_in, C_out, 1, stride=1, padding=0, bias=False),
nn.BatchNorm2d(C_out, affine=affine))
def __init__(self,
C_in,
C_out,
kernel_size,
stride,
padding,
dilation,
affine=True):
super().__init__(
nn.ReLU(inplace=False),
nn.Conv2d(C_in,
C_in,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=C_in,
bias=False),
nn.Conv2d(C_in, C_out, kernel_size=1, padding=0, bias=False),
nn.BatchNorm2d(C_out, affine=affine),
)
def __init__(self, C_in, C_out, affine=True):
super().__init__()
if isinstance(C_out, int):
assert C_out % 2 == 0
else: # is a value choice
assert all(c % 2 == 0 for c in C_out.all_options())
self.relu = nn.ReLU(inplace=False)
self.conv_1 = nn.Conv2d(C_in,
C_out // 2,
1,
stride=2,
padding=0,
bias=False)
self.conv_2 = nn.Conv2d(C_in,
C_out // 2,
1,
stride=2,
padding=0,
bias=False)
self.bn = nn.BatchNorm2d(C_out, affine=affine)
self.pad = nn.ConstantPad2d((0, 1, 0, 1), 0)
def __init__(self):
super().__init__()
self.m = nn.BatchNorm2d(2)
'max_pool_2x2':
lambda C, stride, affine: nn.MaxPool2d(2, stride=stride, padding=0),
'max_pool_3x3':
lambda C, stride, affine: nn.MaxPool2d(3, stride=stride, padding=1),
'max_pool_5x5':
lambda C, stride, affine: nn.MaxPool2d(5, stride=stride, padding=2),
'max_pool_7x7':
lambda C, stride, affine: nn.MaxPool2d(7, stride=stride, padding=3),
'skip_connect':
lambda C, stride, affine: nn.Identity()
if stride == 1 else FactorizedReduce(C, C, affine=affine),
'conv_1x1':
lambda C, stride, affine: nn.Sequential(
nn.ReLU(inplace=False),
nn.Conv2d(C, C, 1, stride=stride, padding=0, bias=False),
nn.BatchNorm2d(C, affine=affine)),
'conv_3x3':
lambda C, stride, affine: nn.Sequential(
nn.ReLU(inplace=False),
nn.Conv2d(C, C, 3, stride=stride, padding=1, bias=False),
nn.BatchNorm2d(C, affine=affine)),
'sep_conv_3x3':
lambda C, stride, affine: SepConv(C, C, 3, stride, 1, affine=affine),
'sep_conv_5x5':
lambda C, stride, affine: SepConv(C, C, 5, stride, 2, affine=affine),
'sep_conv_7x7':
lambda C, stride, affine: SepConv(C, C, 7, stride, 3, affine=affine),
'dil_conv_3x3':
lambda C, stride, affine: DilConv(C, C, 3, stride, 2, 2, affine=affine),
'dil_conv_5x5':
lambda C, stride, affine: DilConv(C, C, 5, stride, 4, 2, affine=affine),
def __init__(self,
op_candidates: List[str],
merge_op: Literal['all', 'loose_end'] = 'all',
num_nodes_per_cell: int = 4,
width: Union[Tuple[int], int] = 16,
num_cells: Union[Tuple[int], int] = 20,
dataset: Literal['cifar', 'imagenet'] = 'imagenet',
auxiliary_loss: bool = False):
super().__init__()
self.dataset = dataset
self.num_labels = 10 if dataset == 'cifar' else 1000
self.auxiliary_loss = auxiliary_loss
# preprocess the specified width and depth
if isinstance(width, Iterable):
C = nn.ValueChoice(list(width), label='width')
else:
C = width
if isinstance(num_cells, Iterable):
num_cells = nn.ValueChoice(list(num_cells), label='depth')
num_cells_per_stage = [
i * num_cells // 3 - (i - 1) * num_cells // 3 for i in range(3)
]
# auxiliary head is different for network targetted at different datasets
if dataset == 'imagenet':
self.stem0 = nn.Sequential(
nn.Conv2d(3,
C // 2,
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.BatchNorm2d(C // 2),
nn.ReLU(inplace=True),
nn.Conv2d(C // 2, C, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(C),
)
self.stem1 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(C, C, 3, stride=2, padding=1, bias=False),
nn.BatchNorm2d(C),
)
C_pprev = C_prev = C_curr = C
last_cell_reduce = True
elif dataset == 'cifar':
self.stem = nn.Sequential(
nn.Conv2d(3, 3 * C, 3, padding=1, bias=False),
nn.BatchNorm2d(3 * C))
C_pprev = C_prev = 3 * C
C_curr = C
last_cell_reduce = False
self.stages = nn.ModuleList()
for stage_idx in range(3):
if stage_idx > 0:
C_curr *= 2
# For a stage, we get C_in, C_curr, and C_out.
# C_in is only used in the first cell.
# C_curr is number of channels for each operator in current stage.
# C_out is usually `C * num_nodes_per_cell` because of concat operator.
cell_builder = CellBuilder(op_candidates, C_pprev, C_prev, C_curr,
num_nodes_per_cell, merge_op,
stage_idx > 0, last_cell_reduce)
stage = nn.Repeat(cell_builder, num_cells_per_stage[stage_idx])
self.stages.append(stage)
# C_pprev is output channel number of last second cell among all the cells already built.
if len(stage) > 1:
# Contains more than one cell
C_pprev = len(stage[-2].output_node_indices) * C_curr
else:
# Look up in the out channels of last stage.
C_pprev = C_prev
# This was originally,
# C_prev = num_nodes_per_cell * C_curr.
# but due to loose end, it becomes,
C_prev = len(stage[-1].output_node_indices) * C_curr
# Useful in aligning the pprev and prev cell.
last_cell_reduce = cell_builder.last_cell_reduce
if stage_idx == 2:
C_to_auxiliary = C_prev
if auxiliary_loss:
assert isinstance(
self.stages[2], nn.Sequential
), 'Auxiliary loss can only be enabled in retrain mode.'
self.stages[2] = SequentialBreakdown(self.stages[2])
self.auxiliary_head = AuxiliaryHead(C_to_auxiliary,
self.num_labels,
dataset=self.dataset)
self.global_pooling = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Linear(C_prev, self.num_labels)
def __init__(self):
super(Fuse, self).__init__()
self.conv = nn.Conv2d(3, 2, kernel_size=1, stride=2, padding=3, bias=False)
self.bn = nn.BatchNorm2d(2)
def __init__(self,
alpha,
depths,
convops,
kernel_sizes,
num_layers,
skips,
num_classes=1000,
dropout=0.2):
super().__init__()
assert alpha > 0.0
assert len(depths) == len(convops) == len(kernel_sizes) == len(
num_layers) == len(skips) == 7
self.alpha = alpha
self.num_classes = num_classes
depths = _get_depths([_FIRST_DEPTH] + depths, alpha)
base_filter_sizes = [16, 24, 40, 80, 96, 192, 320]
exp_ratios = [3, 3, 3, 6, 6, 6, 6]
strides = [1, 2, 2, 2, 1, 2, 1]
layers = [
# First layer: regular conv.
nn.Conv2d(3, depths[0], 3, padding=1, stride=2, bias=False),
nn.BatchNorm2d(depths[0], momentum=_BN_MOMENTUM),
nn.ReLU(inplace=False),
]
count = 0
# for conv, prev_depth, depth, ks, skip, stride, repeat, exp_ratio in \
# zip(convops, depths[:-1], depths[1:], kernel_sizes, skips, strides, num_layers, exp_ratios):
for filter_size, exp_ratio, stride in zip(base_filter_sizes,
exp_ratios, strides):
# TODO: restrict that "choose" can only be used within mutator
ph = nn.Placeholder(
label=f'mutable_{count}',
**{
'kernel_size_options': [1, 3, 5],
'n_layer_options': [1, 2, 3, 4],
'op_type_options': [
'__mutated__.base_mnasnet.RegularConv',
'__mutated__.base_mnasnet.DepthwiseConv',
'__mutated__.base_mnasnet.MobileConv'
],
# 'se_ratio_options': [0, 0.25],
'skip_options': ['identity', 'no'],
'n_filter_options':
[int(filter_size * x) for x in [0.75, 1.0, 1.25]],
'exp_ratio':
exp_ratio,
'stride':
stride,
'in_ch':
depths[0] if count == 0 else None
})
layers.append(ph)
'''if conv == "mconv":
# MNASNet blocks: stacks of inverted residuals.
layers.append(_stack_inverted_residual(prev_depth, depth, ks, skip,
stride, exp_ratio, repeat, _BN_MOMENTUM))
else:
# Normal conv and depth-separated conv
layers += _stack_normal_conv(prev_depth, depth, ks, skip, conv == "dconv",
stride, repeat, _BN_MOMENTUM)'''
count += 1
if count >= 2:
break
layers += [
# Final mapping to classifier input.
nn.Conv2d(depths[7], 1280, 1, padding=0, stride=1, bias=False),
nn.BatchNorm2d(1280, momentum=_BN_MOMENTUM),
nn.ReLU(inplace=False),
]
self.layers = nn.Sequential(*layers)
self.classifier = nn.Sequential(nn.Dropout(p=dropout, inplace=False),
nn.Linear(1280, num_classes))
self._initialize_weights()
def __init__(self):
super().__init__(
nn.Conv2d(3, 3, 1, 1, bias=False),
nn.BatchNorm2d(3),
nn.ReLU(inplace=False)
)
def __init__(self,
num_labels: int = 1000,
channel_search: bool = False,
affine: bool = False):
super().__init__()
self.num_labels = num_labels
self.channel_search = channel_search
self.affine = affine
# the block number in each stage. 4 stages in total. 20 blocks in total.
self.stage_repeats = [4, 4, 8, 4]
# output channels for all stages, including the very first layer and the very last layer
self.stage_out_channels = [-1, 16, 64, 160, 320, 640, 1024]
# building first layer
out_channels = self.stage_out_channels[1]
self.first_conv = nn.Sequential(
nn.Conv2d(3, out_channels, 3, 2, 1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
)
self.features = []
global_block_idx = 0
for stage_idx, num_repeat in enumerate(self.stage_repeats):
for block_idx in range(num_repeat):
# count global index to give names to choices
global_block_idx += 1
# get ready for input and output
in_channels = out_channels
out_channels = self.stage_out_channels[stage_idx + 2]
stride = 2 if block_idx == 0 else 1
# mid channels can be searched
base_mid_channels = out_channels // 2
if self.channel_search:
k_choice_list = [
int(base_mid_channels * (.2 * k)) for k in range(1, 9)
]
mid_channels = nn.ValueChoice(
k_choice_list, label=f'channel_{global_block_idx}')
else:
mid_channels = int(base_mid_channels)
choice_block = nn.LayerChoice(
[
ShuffleNetBlock(in_channels,
out_channels,
mid_channels=mid_channels,
kernel_size=3,
stride=stride,
affine=affine),
ShuffleNetBlock(in_channels,
out_channels,
mid_channels=mid_channels,
kernel_size=5,
stride=stride,
affine=affine),
ShuffleNetBlock(in_channels,
out_channels,
mid_channels=mid_channels,
kernel_size=7,
stride=stride,
affine=affine),
ShuffleXceptionBlock(in_channels,
out_channels,
mid_channels=mid_channels,
stride=stride,
affine=affine)
],
label=f'layer_{global_block_idx}')
self.features.append(choice_block)
self.features = nn.Sequential(*self.features)
# final layers
last_conv_channels = self.stage_out_channels[-1]
self.conv_last = nn.Sequential(
nn.Conv2d(out_channels, last_conv_channels, 1, 1, 0, bias=False),
nn.BatchNorm2d(last_conv_channels, affine=affine),
nn.ReLU(inplace=True),
)
self.globalpool = nn.AdaptiveAvgPool2d((1, 1))
self.dropout = nn.Dropout(0.1)
self.classifier = nn.Sequential(
nn.Linear(last_conv_channels, num_labels, bias=False), )
self._initialize_weights()
def __init__(self, C_in, C_out, kernel_size, stride, padding, affine=True):
super(StdConv, self).__init__()
self.net = nn.Sequential(
nn.ReLU(),
nn.Conv2d(C_in, C_out, kernel_size, stride, padding, bias=False),
nn.BatchNorm2d(C_out, affine=affine))
def __init__(self, some_ch):
super().__init__()
self.some_ch = some_ch
self.batch_norm = nn.BatchNorm2d(some_ch)
def __init__(self,
op_candidates: List[str],
merge_op: Literal['all', 'loose_end'] = 'all',
num_nodes_per_cell: int = 4,
width: Union[Tuple[int, ...], int] = 16,
num_cells: Union[Tuple[int, ...], int] = 20,
dataset: Literal['cifar', 'imagenet'] = 'imagenet',
auxiliary_loss: bool = False):
super().__init__()
self.dataset = dataset
self.num_labels = 10 if dataset == 'cifar' else 1000
self.auxiliary_loss = auxiliary_loss
# preprocess the specified width and depth
if isinstance(width, Iterable):
C = nn.ValueChoice(list(width), label='width')
else:
C = width
self.num_cells: nn.MaybeChoice[int] = cast(int, num_cells)
if isinstance(num_cells, Iterable):
self.num_cells = nn.ValueChoice(list(num_cells), label='depth')
num_cells_per_stage = [
(i + 1) * self.num_cells // 3 - i * self.num_cells // 3
for i in range(3)
]
# auxiliary head is different for network targetted at different datasets
if dataset == 'imagenet':
self.stem0 = nn.Sequential(
nn.Conv2d(3,
cast(int, C // 2),
kernel_size=3,
stride=2,
padding=1,
bias=False),
nn.BatchNorm2d(cast(int, C // 2)),
nn.ReLU(inplace=True),
nn.Conv2d(cast(int, C // 2),
cast(int, C),
3,
stride=2,
padding=1,
bias=False),
nn.BatchNorm2d(C),
)
self.stem1 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(cast(int, C),
cast(int, C),
3,
stride=2,
padding=1,
bias=False),
nn.BatchNorm2d(C),
)
C_pprev = C_prev = C_curr = C
last_cell_reduce = True
elif dataset == 'cifar':
self.stem = nn.Sequential(
nn.Conv2d(3, cast(int, 3 * C), 3, padding=1, bias=False),
nn.BatchNorm2d(cast(int, 3 * C)))
C_pprev = C_prev = 3 * C
C_curr = C
last_cell_reduce = False
else:
raise ValueError(f'Unsupported dataset: {dataset}')
self.stages = nn.ModuleList()
for stage_idx in range(3):
if stage_idx > 0:
C_curr *= 2
# For a stage, we get C_in, C_curr, and C_out.
# C_in is only used in the first cell.
# C_curr is number of channels for each operator in current stage.
# C_out is usually `C * num_nodes_per_cell` because of concat operator.
cell_builder = CellBuilder(op_candidates, C_pprev, C_prev, C_curr,
num_nodes_per_cell, merge_op,
stage_idx > 0, last_cell_reduce)
stage: Union[NDSStage, nn.Sequential] = NDSStage(
cell_builder, num_cells_per_stage[stage_idx])
if isinstance(stage, NDSStage):
stage.estimated_out_channels_prev = cast(int, C_prev)
stage.estimated_out_channels = cast(
int, C_curr * num_nodes_per_cell)
stage.downsampling = stage_idx > 0
self.stages.append(stage)
# NOTE: output_node_indices will be computed on-the-fly in trial code.
# When constructing model space, it's just all the nodes in the cell,
# which happens to be the case of one-shot supernet.
# C_pprev is output channel number of last second cell among all the cells already built.
if len(stage) > 1:
# Contains more than one cell
C_pprev = len(cast(nn.Cell,
stage[-2]).output_node_indices) * C_curr
else:
# Look up in the out channels of last stage.
C_pprev = C_prev
# This was originally,
# C_prev = num_nodes_per_cell * C_curr.
# but due to loose end, it becomes,
C_prev = len(cast(nn.Cell, stage[-1]).output_node_indices) * C_curr
# Useful in aligning the pprev and prev cell.
last_cell_reduce = cell_builder.last_cell_reduce
if stage_idx == 2:
C_to_auxiliary = C_prev
if auxiliary_loss:
assert isinstance(
self.stages[2], nn.Sequential
), 'Auxiliary loss can only be enabled in retrain mode.'
self.stages[2] = SequentialBreakdown(
cast(nn.Sequential, self.stages[2]))
self.auxiliary_head = AuxiliaryHead(
C_to_auxiliary, self.num_labels,
dataset=self.dataset) # type: ignore
self.global_pooling = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Linear(cast(int, C_prev), self.num_labels)
def __init__(self):
super().__init__()
self.m = nn.BatchNorm2d(128, affine=False, momentum=0.3)
