def __init__(self):
super().__init__()
choices = [
{'b': [3], 'bp': [6]},
{'b': [6], 'bp': [12]}
]
self.conv = nn.Conv2d(3, nn.ValueChoice(choices, label='a')['b'][0], 1)
self.conv1 = nn.Conv2d(nn.ValueChoice(choices, label='a')['bp'][0], 3, 1)
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.LayerChoice(
[nn.Conv2d(32, 64, 3, 1),
DepthwiseSeparableConv(32, 64)])
self.dropout1 = nn.Dropout(nn.ValueChoice([0.25, 0.5, 0.75]))
self.dropout2 = nn.Dropout(0.5)
feature = nn.ValueChoice([64, 128, 256])
self.fc1 = nn.Linear(9216, feature)
self.fc2 = nn.Linear(feature, 10)
def _make_stage(self, stage_idx, inp, oup, se, stride, act):
# initialize them first because they are related to layer_count.
exp, ks, se_blocks = [], [], []
for _ in range(4):
exp.append(
nn.ValueChoice(list(self.expand_ratios),
label=f'exp_{self.layer_count}'))
ks.append(nn.ValueChoice([3, 5, 7],
label=f'ks_{self.layer_count}'))
if se:
# if SE is true, assign a layer choice to SE
se_blocks.append(lambda hidden_ch: nn.LayerChoice(
[nn.Identity(), SELayer(hidden_ch)],
label=f'se_{self.layer_count}'))
else:
se_blocks.append(None)
self.layer_count += 1
blocks = [
# stride = 2
InvertedResidual(inp,
oup,
exp[0],
ks[0],
stride,
squeeze_and_excite=se_blocks[0],
activation_layer=act),
# stride = 1, residual connection should be automatically enabled
InvertedResidual(oup,
oup,
exp[1],
ks[1],
squeeze_and_excite=se_blocks[1],
activation_layer=act),
InvertedResidual(oup,
oup,
exp[2],
ks[2],
squeeze_and_excite=se_blocks[2],
activation_layer=act),
InvertedResidual(oup,
oup,
exp[3],
ks[3],
squeeze_and_excite=se_blocks[3],
activation_layer=act)
]
# mutable depth
return nn.Repeat(blocks, depth=(1, 4), label=f'depth_{stage_idx}')
def __init__(self,
n_token: int,
n_head: int = 8,
d_model: int = 512,
d_ff: int = 2048):
super().__init__()
p_dropout = nn.ValueChoice([0.1, 0.2, 0.3, 0.4, 0.5],
label='p_dropout')
n_layer = nn.ValueChoice([5, 6, 7, 8, 9], label='n_layer')
self.encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model, n_head, d_ff, p_dropout),
n_layer)
self.d_model = d_model
self.decoder = nn.Linear(d_model, n_token)
self.embeddings = nn.Embedding(n_token, d_model)
self.position = PositionalEncoding(d_model)
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
# LayerChoice is used to select a layer between Conv2d and DwConv.
self.conv2 = nn.LayerChoice(
[nn.Conv2d(32, 64, 3, 1),
DepthwiseSeparableConv(32, 64)])
# ValueChoice is used to select a dropout rate.
# ValueChoice can be used as parameter of modules wrapped in `nni.retiarii.nn.pytorch`
# or customized modules wrapped with `@basic_unit`.
self.dropout1 = nn.Dropout(nn.ValueChoice(
[0.25, 0.5, 0.75])) # choose dropout rate from 0.25, 0.5 and 0.75
self.dropout2 = nn.Dropout(0.5)
feature = nn.ValueChoice([64, 128, 256])
self.fc1 = nn.Linear(9216, feature)
self.fc2 = nn.Linear(feature, 10)
def __init__(self):
super().__init__()
self.dropout_rate = nn.ValueChoice([[
1.05,
], [
1.1,
]])
def __init__(self):
super().__init__()
channels = nn.ValueChoice([4, 6, 8])
self.conv1 = nn.Conv2d(1, channels, 5)
self.pool1 = nn.LayerChoice([
nn.MaxPool2d((2, 2)), nn.AvgPool2d((2, 2))
])
self.conv2 = nn.Conv2d(channels, 16, 5)
self.pool2 = nn.LayerChoice([
nn.MaxPool2d(2), nn.AvgPool2d(2), nn.Conv2d(16, 16, 2, 2)
])
self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5*5 from image dimension
self.fc2 = nn.Linear(120, 84)
self.fcplus = nn.Linear(84, 84)
self.shortcut = nn.InputChoice(2, 1)
self.fc3 = nn.Linear(84, 10)
def __init__(self, value_choice=True):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = LayerChoice(
[nn.Conv2d(32, 64, 3, 1),
DepthwiseSeparableConv(32, 64)])
self.dropout1 = LayerChoice(
[nn.Dropout(.25), nn.Dropout(.5),
nn.Dropout(.75)])
self.dropout2 = nn.Dropout(0.5)
if value_choice:
hidden = nn.ValueChoice([32, 64, 128])
else:
hidden = 64
self.fc1 = nn.Linear(9216, hidden)
self.fc2 = nn.Linear(hidden, 10)
self.rpfc = nn.Linear(10, 10)
self.input_ch = InputChoice(2, 1)
def __init__(self):
super().__init__()
vc = nn.ValueChoice([(6, 3), (8, 5)])
self.conv = nn.Conv2d(3, vc[0], kernel_size=vc[1])
def __init__(self):
super().__init__()
self.linear = nn.LayerChoice([
nn.Linear(3, nn.ValueChoice([10, 20])),
nn.Linear(3, nn.ValueChoice([30, 40]))
])
def __init__(self):
super().__init__()
self.dropout_rate = nn.ValueChoice([0., 1.])
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, nn.ValueChoice([6, 8], label='shared'), 1)
self.conv2 = nn.Conv2d(3, nn.ValueChoice([6, 8], label='shared'), 1)
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, nn.ValueChoice([6, 8]), kernel_size=nn.ValueChoice([3, 5]))
def __init__(self,
num_labels: int = 1000,
base_widths: Tuple[int, ...] = (16, 16, 32, 64, 128, 256, 512,
1024),
width_multipliers: Tuple[float, ...] = (0.5, 0.625, 0.75, 1.0,
1.25, 1.5, 2.0),
expand_ratios: Tuple[int, ...] = (1, 2, 3, 4, 5, 6),
dropout_rate: float = 0.2,
bn_eps: float = 1e-3,
bn_momentum: float = 0.1):
super().__init__()
self.widths = [
nn.ValueChoice([
make_divisible(base_width * mult, 8)
for mult in width_multipliers
],
label=f'width_{i}')
for i, base_width in enumerate(base_widths)
]
self.expand_ratios = expand_ratios
blocks = [
# Stem
ConvBNReLU(3,
self.widths[0],
nn.ValueChoice([3, 5], label='ks_0'),
stride=2,
activation_layer=h_swish),
SeparableConv(self.widths[0],
self.widths[0],
activation_layer=nn.ReLU),
]
# counting for kernel sizes and expand ratios
self.layer_count = 2
blocks += [
# Body
self._make_stage(1, self.widths[0], self.widths[1], False, 2,
nn.ReLU),
self._make_stage(2, self.widths[1], self.widths[2], True, 2,
nn.ReLU),
self._make_stage(1, self.widths[2], self.widths[3], False, 2,
h_swish),
self._make_stage(1, self.widths[3], self.widths[4], True, 1,
h_swish),
self._make_stage(1, self.widths[4], self.widths[5], True, 2,
h_swish),
]
# Head
blocks += [
ConvBNReLU(self.widths[5],
self.widths[6],
1,
1,
activation_layer=h_swish),
nn.AdaptiveAvgPool2d(1),
ConvBNReLU(self.widths[6],
self.widths[7],
1,
1,
norm_layer=nn.Identity,
activation_layer=h_swish),
]
self.blocks = nn.Sequential(*blocks)
self.classifier = nn.Sequential(
nn.Dropout(dropout_rate),
nn.Linear(self.widths[7], num_labels),
)
reset_parameters(self, bn_momentum=bn_momentum, bn_eps=bn_eps)
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,
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):
super().__init__()
self.index = nn.ValueChoice([0, 1])
self.conv = MutableConv()
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,
search_embed_dim: Tuple[int, ...] = (192, 216, 240),
search_mlp_ratio: Tuple[float, ...] = (3.5, 4.0),
search_num_heads: Tuple[int, ...] = (3, 4),
search_depth: Tuple[int, ...] = (12, 13, 14),
img_size: int = 224,
patch_size: int = 16,
in_chans: int = 3,
num_classes: int = 1000,
qkv_bias: bool = False,
drop_rate: float = 0.,
attn_drop_rate: float = 0.,
drop_path_rate: float = 0.,
pre_norm: bool = True,
global_pool: bool = False,
abs_pos: bool = True,
qk_scale: Optional[float] = None,
rpe: bool = True,
):
super().__init__()
embed_dim = nn.ValueChoice(list(search_embed_dim), label="embed_dim")
fixed_embed_dim = nn.ModelParameterChoice(
list(search_embed_dim), label="embed_dim")
depth = nn.ValueChoice(list(search_depth), label="depth")
self.patch_embed = nn.Conv2d(
in_chans,
cast(int, embed_dim),
kernel_size=patch_size,
stride=patch_size)
self.patches_num = int((img_size // patch_size) ** 2)
self.global_pool = global_pool
self.cls_token = nn.Parameter(torch.zeros(1, 1, cast(int, fixed_embed_dim)))
trunc_normal_(self.cls_token, std=.02)
dpr = [
x.item() for x in torch.linspace(
0,
drop_path_rate,
max(search_depth))] # stochastic depth decay rule
self.abs_pos = abs_pos
if self.abs_pos:
self.pos_embed = nn.Parameter(torch.zeros(
1, self.patches_num + 1, cast(int, fixed_embed_dim)))
trunc_normal_(self.pos_embed, std=.02)
self.blocks = nn.Repeat(lambda index: nn.LayerChoice([
TransformerEncoderLayer(embed_dim=embed_dim,
fixed_embed_dim=fixed_embed_dim,
num_heads=num_heads, mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias, drop_rate=drop_rate,
attn_drop=attn_drop_rate,
drop_path=dpr[index],
rpe_length=img_size // patch_size,
qk_scale=qk_scale, rpe=rpe,
pre_norm=pre_norm,)
for mlp_ratio, num_heads in itertools.product(search_mlp_ratio, search_num_heads)
], label=f'layer{index}'), depth)
self.pre_norm = pre_norm
if self.pre_norm:
self.norm = nn.LayerNorm(cast(int, embed_dim))
self.head = nn.Linear(
cast(int, embed_dim),
num_classes) if num_classes > 0 else nn.Identity()
