def __init__(self, **kwargs):
super(AnyNet, self).__init__()
self.logger = get_logger(__name__)
if kwargs:
self._construct(
stem_w=kwargs["stem_w"],
ds=kwargs["ds"],
ws=kwargs["ws"],
ss=kwargs["ss"],
bms=kwargs["bms"],
gws=kwargs["gws"],
se_r=kwargs["se_r"],
nc=kwargs["nc"],
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
# Note that there is no bias due to BN
fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(mean=0.0, std=math.sqrt(2.0 / fan_out))
elif (isinstance(m, SyncBatchNorm2d) or isinstance(m, nn.BatchNorm2d)):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
n = m.weight.size(1)
m.weight.data.normal_(0, 1.0 / float(n))
m.bias.data.zero_()
def efficientnet(width_coefficient=None, depth_coefficient=None,
dropout_rate=0.2, drop_connect_rate=0.3, override_block=None):
"""Creates a efficientnet model."""
blocks_args = [
'r1_k3_s11_e1_i32_o16_se0.25', 'r2_k3_s22_e6_i16_o24_se0.25',
'r2_k5_s22_e6_i24_o40_se0.25', 'r3_k3_s22_e6_i40_o80_se0.25',
'r3_k5_s11_e6_i80_o112_se0.25', 'r4_k5_s22_e6_i112_o192_se0.25',
'r1_k3_s11_e6_i192_o320_se0.25',
]
if override_block is not None:
assert isinstance(override_block, dict)
for k, v in override_block.items():
blocks_args[int(k)] = v
logger = get_logger(__name__)
logger.info('overrided blocks_args: {}'.format(blocks_args))
global_params = GlobalParams(dropout_rate=dropout_rate,
drop_connect_rate=drop_connect_rate,
data_format='channels_last',
num_classes=1000,
width_coefficient=width_coefficient,
depth_coefficient=depth_coefficient,
depth_divisor=8,
min_depth=None)
decoder = BlockDecoder()
return decoder.decode(blocks_args), global_params
def __init__(self, key_metric, defect_classes, recall_thres, tpr_thres):
super(CustomEvaluator, self).__init__()
self.key_metric = key_metric
self.recall_thres = recall_thres
self.tpr_thres = tpr_thres
self.defect_classes = defect_classes
self.logger = get_logger(__name__)
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
num_channels, fuse_method, multi_scale_output=True):
super(HighResolutionModule, self).__init__()
self.logger = get_logger(__name__)
self._check_branches(
num_branches, blocks, num_blocks, num_inchannels, num_channels)
self.num_inchannels = num_inchannels
self.fuse_method = fuse_method
self.num_branches = num_branches
self.multi_scale_output = multi_scale_output
self.branches = self._make_branches(
num_branches, blocks, num_blocks, num_channels)
self.fuse_layers = self._make_fuse_layers()
self.relu = nn.ReLU(False)
def round_filters(filters, global_params):
"""Round number of filters based on depth multiplier."""
orig_f = filters
multiplier = global_params.width_coefficient
divisor = global_params.depth_divisor
min_depth = global_params.min_depth
if not multiplier:
return filters
filters *= multiplier
min_depth = min_depth or divisor
new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_filters < 0.9 * filters:
new_filters += divisor
logger = get_logger(__name__)
logger.info('round_filter input={} output={}'.format(orig_f, new_filters))
return int(new_filters)
def get_model_params(model_name, override_params=None, override_block=None):
"""Get the block args and global params for a given model."""
if model_name.startswith('efficientnet'):
width_coefficient, depth_coefficient, _, dropout_rate = (efficientnet_params(model_name))
blocks_args, global_params = efficientnet(width_coefficient, depth_coefficient,
dropout_rate, override_block=override_block)
else:
raise NotImplementedError('model name is not pre-defined: %s' % model_name)
if override_params is not None:
# ValueError will be raised here if override_params has fields not included
# in global_params.
global_params = global_params._replace(**override_params)
logger = get_logger(__name__)
logger.info(blocks_args)
logger.info(global_params)
return blocks_args, global_params
def setup_env(self):
# dist
self.dist = EasyDict()
self.dist.rank, self.dist.world_size = link.get_rank(
), link.get_world_size()
self.prototype_info.world_size = self.dist.world_size
# directories
self.path = EasyDict()
self.path.root_path = os.path.dirname(self.config_file)
self.path.save_path = os.path.join(self.path.root_path, 'checkpoints')
self.path.event_path = os.path.join(self.path.root_path, 'events')
self.path.result_path = os.path.join(self.path.root_path, 'results')
makedir(self.path.save_path)
makedir(self.path.event_path)
makedir(self.path.result_path)
# tb_logger
if self.dist.rank == 0:
self.tb_logger = SummaryWriter(self.path.event_path)
# logger
create_logger(os.path.join(self.path.root_path, 'log.txt'))
self.logger = get_logger(__name__)
self.logger.info(f'config: {pprint.pformat(self.config)}')
if 'SLURM_NODELIST' in os.environ:
self.logger.info(f"hostnames: {os.environ['SLURM_NODELIST']}")
# load pretrain checkpoint
if hasattr(self.config.saver, 'pretrain'):
self.state = torch.load(self.config.saver.pretrain.path, 'cpu')
self.logger.info(
f"Recovering from {self.config.saver.pretrain.path}, keys={list(self.state.keys())}"
)
if hasattr(self.config.saver.pretrain, 'ignore'):
self.state = modify_state(self.state,
self.config.saver.pretrain.ignore)
else:
self.state = {}
self.state['last_iter'] = 0
# others
torch.backends.cudnn.benchmark = True
def __init__(self,
blocks_args=None,
global_params=None,
use_fc_bn=False,
fc_bn_init_scale=1.0,
bn=None):
super(EfficientNet, self).__init__()
global BN
BN = get_bn(bn)
if not isinstance(blocks_args, list):
raise ValueError('blocks_args should be a list.')
self.logger = get_logger(__name__)
self._global_params = global_params
self._blocks_args = blocks_args
self.use_fc_bn = use_fc_bn
self.fc_bn_init_scale = fc_bn_init_scale
self._build()
def _setup_env(self):
# distribution information
self.dist = EasyDict()
self.dist.rank, self.dist.world_size = link.get_rank(
), link.get_world_size()
# directories
self.path = EasyDict()
self.path.root_path = self.work_dir
self.path.save_path = os.path.join(self.path.root_path, 'checkpoints')
self.path.event_path = os.path.join(self.path.root_path, 'events')
self.path.result_path = os.path.join(self.path.root_path, 'results')
makedir(self.path.save_path)
makedir(self.path.event_path)
makedir(self.path.result_path)
# create tensorboard logger
if self.dist.rank == 0:
self.tb_logger = SummaryWriter(self.path.event_path)
# create logger
create_logger(os.path.join(self.path.root_path, 'log.txt'))
self.logger = get_logger(__name__)
self.logger.info(f'config: {pprint.pformat(self.config)}')
self.logger.info(f"hostnames: {os.environ['SLURM_NODELIST']}")
# others
torch.backends.cudnn.benchmark = True
def __init__(self,
block,
layers,
num_classes=1000,
deep_stem=False,
avg_down=False,
bypass_last_bn=False,
bn=None):
r"""
Arguments:
- layers (:obj:`list` of 4 ints): how many layers in each stage
- num_classes (:obj:`int`): number of classification classes
- deep_stem (:obj:`bool`): whether to use deep_stem as the first conv
- avg_down (:obj:`bool`): whether to use avg_down when spatial downsample
- bypass_last_bn (:obj:`bool`): whether use bypass_last_bn
- bn (:obj:`dict`): definition of batchnorm
"""
super(PreactResNet, self).__init__()
logger = get_logger(__name__)
global BN, bypass_bn_weight_list
BN = get_bn(bn)
bypass_bn_weight_list = []
self.inplanes = 64
self.deep_stem = deep_stem
self.avg_down = avg_down
self.logger = get_logger(__name__)
if self.deep_stem:
self.conv1 = nn.Sequential(
nn.Conv2d(3,
32,
kernel_size=3,
stride=2,
padding=1,
bias=False),
BN(32),
nn.ReLU(inplace=True),
nn.Conv2d(32,
32,
kernel_size=3,
stride=1,
padding=1,
bias=False),
BN(32),
nn.ReLU(inplace=True),
nn.Conv2d(32,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False),
)
else:
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = BN(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.final_bn = BN(512 * block.expansion)
self.final_relu = nn.ReLU(inplace=True)
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif (isinstance(m, SyncBatchNorm2d)
or isinstance(m, nn.BatchNorm2d)):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
n = m.weight.size(1)
m.weight.data.normal_(0, 1.0 / float(n))
m.bias.data.zero_()
if bypass_last_bn:
for param in bypass_bn_weight_list:
param.data.zero_()
logger.info('bypass {} bn.weight in BottleneckBlocks'.format(
len(bypass_bn_weight_list)))
def __init__(self,
block,
layers,
inplanes=64,
num_classes=1000,
zero_init_residual=False,
groups=1,
width_per_group=64,
replace_stride_with_dilation=None,
norm_layer=None,
deep_stem=False,
avg_down=False,
freeze_layer=False,
bn=None):
super(ResNet, self).__init__()
global BN
self.logger = get_logger(__name__)
if norm_layer is None:
BN = get_bn(bn)
norm_layer = BN
else:
norm_layer = get_norm_layer(norm_layer)
self._norm_layer = norm_layer
self.inplanes = inplanes
self.dilation = 1
self.deep_stem = deep_stem
self.avg_down = avg_down
self.num_classes = num_classes
self.freeze_layer = freeze_layer
if replace_stride_with_dilation is None:
# each element in the tuple indicates if we should replace
# the 2x2 stride with a dilated convolution instead
replace_stride_with_dilation = [False, False, False]
if len(replace_stride_with_dilation) != 3:
raise ValueError("replace_stride_with_dilation should be None "
"or a 3-element tuple, got {}".format(
replace_stride_with_dilation))
self.groups = groups
self.base_width = width_per_group
if self.deep_stem:
self.conv1 = nn.Sequential(
nn.Conv2d(3,
inplanes // 2,
kernel_size=3,
stride=2,
padding=1,
bias=False),
norm_layer(inplanes // 2),
nn.ReLU(inplace=True),
nn.Conv2d(inplanes // 2,
inplanes // 2,
kernel_size=3,
stride=1,
padding=1,
bias=False),
norm_layer(inplanes // 2),
nn.ReLU(inplace=True),
nn.Conv2d(inplanes // 2,
inplanes,
kernel_size=3,
stride=1,
padding=1,
bias=False),
)
else:
self.conv1 = nn.Conv2d(3,
inplanes,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = norm_layer(self.inplanes)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block,
128,
layers[1],
stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block,
256,
layers[2],
stride=2,
dilate=replace_stride_with_dilation[1])
self.layer4 = self._make_layer(block,
512,
layers[3],
stride=2,
dilate=replace_stride_with_dilation[2])
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
elif isinstance(m,
(nn.BatchNorm2d, nn.GroupNorm, SyncBatchNorm2d)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
# Zero-initialize the last BN in each residual branch,
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
nn.init.constant_(m.bn3.weight, 0)
elif isinstance(m, BasicBlock):
nn.init.constant_(m.bn2.weight, 0)
def __init__(self,
num_classes=1000,
scale=1.0,
inverted_residual_setting=None,
round_nearest=8,
block=InvertedResidual,
dropout=0.2,
bn=None,
num_experts=1,
final_condconv=False,
fc_condconv=False,
combine_kernel=False):
super(MobileNetV2CondConv, self).__init__()
global BN
BN = get_bn(bn)
self.logger = get_logger(__name__)
self.fc_condconv = fc_condconv
self.logger.info('Number of experts is {}'.format(num_experts))
self.logger.info(
'Replace finalconv with CondConv: {}'.format(final_condconv))
self.logger.info('Replace fc with CondConv: {}'.format(fc_condconv))
self.logger.info(
'Combine kernels to implement CondConv: {}'.format(combine_kernel))
if block is None:
block = InvertedResidual
input_channel = 32
last_channel = 1280
if inverted_residual_setting is None:
inverted_residual_setting = [
# t, c, n, s
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
# only check the first element, assuming user knows t,c,n,s are required
if len(inverted_residual_setting) == 0 or len(
inverted_residual_setting[0]) != 4:
raise ValueError("inverted_residual_setting should be non-empty "
"or a 4-element list, got {}".format(
inverted_residual_setting))
# building first layer
input_channel = _make_divisible(input_channel * scale, round_nearest)
self.last_channel = _make_divisible(last_channel * max(1.0, scale),
round_nearest)
features = [ConvBNReLU(3, input_channel, stride=2)]
# building inverted residual blocks
for t, c, n, s in inverted_residual_setting:
output_channel = _make_divisible(c * scale, round_nearest)
for i in range(n):
stride = s if i == 0 else 1
features.append(
block(input_channel,
output_channel,
stride,
expand_ratio=t,
num_experts=num_experts,
combine_kernel=combine_kernel))
input_channel = output_channel
# building last several layers
if final_condconv:
features.append(
CondConvBNReLU(input_channel,
self.last_channel,
kernel_size=1,
num_experts=num_experts,
combine_kernel=combine_kernel))
else:
features.append(
ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
# make it nn.Sequential
self.features = nn.Sequential(*features)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
# building classifier
if fc_condconv:
# change kernel_size to the size of feature maps
self.dropout = nn.Dropout(0.2)
self.classifier = CondConv2d(self.last_channel,
num_classes,
kernel_size=1,
bias=False,
num_experts=num_experts,
combine_kernel=combine_kernel)
self.classifier_router = BasicRouter(self.last_channel,
num_experts)
else:
self.classifier = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(self.last_channel, num_classes),
)
# weight initialization
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d) or isinstance(
m, link.nn.SyncBatchNorm2d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.zeros_(m.bias)