def generate_nas_layer(self,
C_cur,
C_p,
C_pp,
reduction_p,
reduction_cur,
genotype,
cell_num=3,
bn_affine=True):
cells = nn.ModuleList()
# reduction_idx = (cell_num + 1) // 2 - 1
# the first cell(block) is downsample
# reduction_idx = 0
if self.res_stem:
reduction_idx = 0
else:
reduction_idx = cell_num - 1
for i in range(cell_num):
if i == reduction_idx and reduction_cur:
C_cur *= 2
reduction = True
else:
reduction = False
cell = AugmentCell(genotype, C_pp, C_p, C_cur, reduction_p,
reduction, bn_affine)
reduction_p = reduction
cells.append(cell)
C_cur_out = C_cur * len(cell.concat)
C_pp, C_p = C_p, C_cur_out
return cells
def __init__(self,
input_size,
C_in,
C,
n_classes,
n_layers,
auxiliary,
genotype,
stem_multiplier=3,
SSC=False):
"""
Args:
input_size: size of height and width (assuming height = width)
C_in: # of input channels
C: # of starting model channels
"""
super().__init__()
self.C_in = C_in
self.C = C
self.n_classes = n_classes
self.n_layers = n_layers
self.genotype = genotype
# aux head position
self.aux_pos = 2 * n_layers // 3 if auxiliary else -1
C_cur = stem_multiplier * C
self.stem = nn.Sequential(nn.Conv2d(C_in, C_cur, 3, 1, 1, bias=False),
nn.BatchNorm2d(C_cur))
C_pp, C_p, C_cur = C_cur, C_cur, C
self.cells = nn.ModuleList()
reduction_p = False
for i in range(n_layers):
if i in [n_layers // 3, 2 * n_layers // 3]:
C_cur *= 2
reduction = True
else:
reduction = False
cell = AugmentCell(genotype, i, C_pp, C_p, C_cur, reduction_p,
reduction, SSC)
reduction_p = reduction
self.cells.append(cell)
C_cur_out = C_cur * len(cell.concat)
C_pp, C_p = C_p, C_cur_out
if i == self.aux_pos:
# [!] this auxiliary head is ignored in computing parameter size
# by the name 'aux_head'
self.aux_head = AuxiliaryHead(input_size // 4, C_p, n_classes)
self.gap = nn.AdaptiveAvgPool2d(1)
self.linear = nn.Linear(C_p, n_classes)
def __init__(self, input_size, C_in, C, n_classes, n_layers, auxiliary, genotype,
stem_multiplier=3):
"""
Args:
input_size: size of height and width (assuming height = width)
C_in: # of input channels
C: # of starting model channels
"""
super().__init__()
self.C_in = C_in
self.C = C
self.n_classes = n_classes
self.n_layers = n_layers
self.genotype = genotype
C_cur = 32
self.stem = nn.Sequential(
nn.BatchNorm2d(C_in),
nn.Conv2d(C_in, C_cur, 5, 2, 2, bias=False),
nn.BatchNorm2d(C_cur),
nn.ReLU(),
nn.Conv2d(C_cur, C_cur, 3, 2, 1, bias=False),
nn.BatchNorm2d(C_cur),
nn.ReLU(),
)
C_pp, C_p, C_cur = C_cur, C_cur, C_cur
self.cells = nn.ModuleList()
reduction_p = False
for i in range(n_layers):
if i in [1*n_layers//6, 3*n_layers//6, 5*n_layers//6]:
C_cur *= 2
reduction = True
else:
reduction = False
cell = AugmentCell(genotype, C_pp, C_p, C_cur, reduction_p, reduction)
reduction_p = reduction
self.cells.append(cell)
C_cur_out = C_cur * len(cell.concat)
C_pp, C_p = C_p, C_cur_out
self.gap = nn.Sequential(
nn.Conv2d(C_p, 512, 3, 2, 1, bias=False),
nn.BatchNorm2d(512),
nn.AdaptiveAvgPool2d(1)
)
self.linear = nn.Linear(512, n_classes)
def __init__(self,
input_size,
C_in,
C,
n_classes,
n_layers,
auxiliary,
genotype,
stem_multiplier=3):
super().__init__()
self.C_in = C_in
self.C = C
self.n_classes = n_classes
self.n_layers = n_layers
self.genotype = genotype
self.aux_pos = 2 * n_layers // 3 if auxiliary else -1
C_cur = stem_multiplier * C
self.stem = nn.Sequential(nn.Conv2d(C_in, C_cur, 3, 1, 1, bias=False),
nn.BatchNorm2d(C_cur))
C_pp, C_p, C_cur = C_cur, C_cur, C
self.cells = nn.ModuleList()
reduction_p = False
for i in range(n_layers):
if i in [n_layers // 3, 2 * n_layers // 3]:
C_cur *= 2
reduction = True
else:
reduction = False
cell = AugmentCell(genotype, C_pp, C_p, C_cur, reduction_p,
reduction)
reduction_p = reduction
self.cells.append(cell)
C_cur_out = C_cur * len(cell.concat)
C_pp, C_p = C_p, C_cur_out
if i == self.aux_pos:
self.aux_head = AuxiliaryHead(input_size // 4, C_p, n_classes)
self.gap = nn.AdaptiveAvgPool2d(1)
self.linear = nn.Linear(C_p, n_classes)