in ppuda/deepnets1m/net.py [0:0]
def __init__(self,
C,
num_classes,
genotype,
n_cells,
ks=3,
is_imagenet_input=True,
stem_pool=False,
stem_type=0,
is_vit=None,
norm='bn-track',
preproc=True,
C_mult=2,
fc_layers=0,
fc_dim=0,
glob_avg=True,
auxiliary=False,
compress_params=False
):
super(Network, self).__init__()
self.genotype = genotype
self._auxiliary = auxiliary
self.drop_path_prob = 0
self.expected_image_sz = 224 if is_imagenet_input else 32
self._is_vit = sum([n[0] == 'msa' for n in genotype.normal + genotype.reduce]) > 0 if is_vit is None else is_vit
steps = len(genotype.normal_concat) # number of inputs to the concatenation operation
if steps > 1 or C_mult > 1:
assert preproc, 'preprocessing layers must be used in this case'
self._stem_type = stem_type
assert stem_type in [0, 1], ('either 0 (simple) or 1 (imagenet-style) stem must be chosen', stem_type)
C_prev_prev = C_prev = C_curr = C
# Define the stem
if self._is_vit:
# Visual Transformer stem
self.stem0 = OPS['conv_stride'](3, C, 16 if is_imagenet_input else 3, None, norm)
self.pos_enc = PosEnc(C, 14 if is_imagenet_input else 11)
elif stem_type == 0:
# Simple stem
C_stem = int(C * (3 if (preproc and not is_imagenet_input) else 1))
self.stem = nn.Sequential(
nn.Conv2d(3, C_stem, ks, stride=4 if is_imagenet_input else 1, padding=ks // 2, bias=False),
get_norm_layer(norm, C_stem),
nn.MaxPool2d(3, stride=2, padding=1, ceil_mode=False) if stem_pool else nn.Identity(),
)
C_prev_prev = C_prev = C_stem
else:
# ImageNet-style stem
self.stem0 = nn.Sequential(
nn.Conv2d(3, C // 2, kernel_size=ks, stride=2 if is_imagenet_input else 1, padding=ks // 2, bias=False),
get_norm_layer(norm, C // 2),
nn.ReLU(inplace=True),
nn.Conv2d(C // 2, C, kernel_size=3, stride=2 if is_imagenet_input else 1, padding=1, bias=False),
get_norm_layer(norm, C)
)
self.stem1 = nn.Sequential(
nn.ReLU(inplace=True),
nn.Conv2d(C, C, 3, stride=2, padding=1, bias=False),
get_norm_layer(norm, C)
)
self._n_cells = n_cells
self.cells = nn.ModuleList()
is_reduction = lambda cell_ind: cell_ind in [n_cells // 3, 2 * n_cells // 3] and cell_ind > 0
self._auxiliary_cell_ind = 2 * n_cells // 3
reduction_prev = stem_type == 1
for cell_ind in range(n_cells):
if is_reduction(cell_ind):
C_curr *= C_mult
reduction = True
else:
reduction = False
reduction_next = is_reduction(cell_ind + 1)
cell = Cell(genotype,
C_prev_prev,
C_prev,
C_in=C_curr if preproc else C_prev,
C_out=C_curr * (C_mult if reduction_next and steps == 1 and not preproc else 1),
reduction=reduction,
reduction_prev=reduction_prev,
norm=norm,
is_vit=self._is_vit,
preproc=preproc)
self.cells.append(cell)
reduction_prev = reduction
C_prev_prev, C_prev = C_prev, cell.multiplier * C_curr
if auxiliary and cell_ind == self._auxiliary_cell_ind:
if is_imagenet_input:
self.auxiliary_head = AuxiliaryHeadImageNet(C_prev, num_classes, norm=norm)
else:
self.auxiliary_head = AuxiliaryHeadCIFAR(C_prev, num_classes, norm=norm,
pool_sz=2 if (stem_type == 1 or stem_pool) else 5)
self._glob_avg = glob_avg
if glob_avg:
self.global_pooling = nn.AdaptiveAvgPool2d(1)
else:
if is_imagenet_input:
s = 7 if (stem_type == 1 or stem_pool) else 14
else:
s = 4 if (stem_type == 1 or stem_pool) else 8
C_prev *= s ** 2
fc = [nn.Linear(C_prev, fc_dim if fc_layers > 1 else num_classes)]
for i in range(fc_layers - 1):
assert fc_dim > 0, fc_dim
fc.append(nn.ReLU(inplace=True))
fc.append(nn.Dropout(p=0.5, inplace=False))
fc.append(nn.Linear(in_features=fc_dim, out_features=fc_dim if i < fc_layers - 2 else num_classes))
self.classifier = nn.Sequential(*fc)
if compress_params:
for p in self.parameters():
p.data = p.data.bool()