def __init__(self,
pretrain=True,
act_layer='leaky_relu',
acon_layer='none'):
super(CSPDarkNet53Encoder, self).__init__()
if acon_layer == 'acon':
base = _create_cspnet('cspdarknet53',
pretrained=pretrain,
pretrained_strict=False,
act_layer=get_act_layer(act_layer),
block_fn=DarkBlock_ACON)
elif acon_layer == 'meta_acon':
base = _create_cspnet('cspdarknet53',
pretrained=pretrain,
pretrained_strict=False,
act_layer=get_act_layer(act_layer),
block_fn=DarkBlock_MetaACON)
elif acon_layer == 'none':
base = timm.create_model('cspdarknet53',
pretrained=pretrain,
act_layer=get_act_layer(act_layer))
self.stem = base.stem
self.layer0 = base.stages[0]
self.layer1 = base.stages[1]
self.layer2 = base.stages[2]
self.layer3 = base.stages[3]
self.layer4 = base.stages[4]
del base
self.depth = 5
self.out_channels = (32, 64, 128, 256, 512, 1024)
self.in_channels = 3
def __init__(self, config, feature_info):
super(BiFpn, self).__init__()
self.num_levels = config.num_levels
norm_layer = config.norm_layer or nn.BatchNorm2d
if config.norm_kwargs:
norm_layer = partial(norm_layer, **config.norm_kwargs)
act_layer = get_act_layer(config.act_type) or _ACT_LAYER
fpn_config = config.fpn_config or get_fpn_config(
config.fpn_name,
min_level=config.min_level,
max_level=config.max_level)
self.resample = nn.ModuleDict()
for level in range(config.num_levels):
if level < len(feature_info):
in_chs = feature_info[level]['num_chs']
reduction = feature_info[level]['reduction']
else:
# Adds a coarser level by downsampling the last feature map
reduction_ratio = 2
self.resample[str(level)] = ResampleFeatureMap(
in_channels=in_chs,
out_channels=config.fpn_channels,
pad_type=config.pad_type,
downsample=config.downsample_type,
upsample=config.upsample_type,
norm_layer=norm_layer,
reduction_ratio=reduction_ratio,
apply_bn=config.apply_resample_bn,
conv_after_downsample=config.conv_after_downsample,
redundant_bias=config.redundant_bias,
)
in_chs = config.fpn_channels
reduction = int(reduction * reduction_ratio)
feature_info.append(dict(num_chs=in_chs, reduction=reduction))
self.cell = SequentialList()
for rep in range(config.fpn_cell_repeats):
logging.debug('building cell {}'.format(rep))
fpn_layer = BiFpnLayer(
feature_info=feature_info,
fpn_config=fpn_config,
fpn_channels=config.fpn_channels,
num_levels=config.num_levels,
pad_type=config.pad_type,
downsample=config.downsample_type,
upsample=config.upsample_type,
norm_layer=norm_layer,
act_layer=act_layer,
separable_conv=config.separable_conv,
apply_resample_bn=config.apply_resample_bn,
conv_after_downsample=config.conv_after_downsample,
conv_bn_relu_pattern=config.conv_bn_relu_pattern,
redundant_bias=config.redundant_bias,
)
self.cell.add_module(str(rep), fpn_layer)
feature_info = fpn_layer.feature_info
def __init__(self, config, num_outputs):
super(HeadNet, self).__init__()
self.num_levels = config.num_levels
self.bn_level_first = getattr(config, 'head_bn_level_first', False)
norm_layer = config.norm_layer or nn.BatchNorm2d
if config.norm_kwargs:
norm_layer = partial(norm_layer, **config.norm_kwargs)
act_type = config.head_act_type if getattr(config, 'head_act_type',
None) else config.act_type
act_layer = get_act_layer(act_type) or _ACT_LAYER
# Build convolution repeats
conv_fn = SeparableConv2d if config.separable_conv else ConvBnAct2d
conv_kwargs = dict(in_channels=config.fpn_channels,
out_channels=config.fpn_channels,
kernel_size=3,
padding=config.pad_type,
bias=config.redundant_bias,
act_layer=None,
norm_layer=None)
self.conv_rep = nn.ModuleList(
[conv_fn(**conv_kwargs) for _ in range(config.box_class_repeats)])
# Build batchnorm repeats. There is a unique batchnorm per feature level for each repeat.
# This can be organized with repeats first or feature levels first in module lists, the original models
# and weights were setup with repeats first, levels first is required for efficient torchscript usage.
self.bn_rep = nn.ModuleList()
if self.bn_level_first:
for _ in range(self.num_levels):
self.bn_rep.append(
nn.ModuleList([
norm_layer(config.fpn_channels)
for _ in range(config.box_class_repeats)
]))
else:
for _ in range(config.box_class_repeats):
self.bn_rep.append(
nn.ModuleList([
nn.Sequential(
OrderedDict([('bn',
norm_layer(config.fpn_channels))]))
for _ in range(self.num_levels)
]))
self.act = act_layer(inplace=True)
# Prediction (output) layer. Has bias with special init reqs, see init fn.
num_anchors = len(config.aspect_ratios) * config.num_scales
predict_kwargs = dict(in_channels=config.fpn_channels,
out_channels=num_outputs * num_anchors,
kernel_size=3,
padding=config.pad_type,
bias=True,
norm_layer=None,
act_layer=None)
self.predict = conv_fn(**predict_kwargs)
def __init__(self, config, num_outputs):
super(HeadNet, self).__init__()
norm_layer = config.norm_layer or nn.BatchNorm2d
norm_kwargs = config.norm_kwargs or {}
act_layer = get_act_layer(config.act_type) or _ACT_LAYER
self.config = config
num_anchors = len(config.aspect_ratios) * config.num_scales
self.conv_rep = nn.ModuleList()
self.bn_rep = nn.ModuleList()
conv_kwargs = dict(in_channels=config.fpn_channels,
out_channels=config.fpn_channels,
kernel_size=3,
padding=self.config.pad_type,
bias=config.redundant_bias,
act_layer=None,
norm_layer=None)
for i in range(config.box_class_repeats):
conv = SeparableConv2d(
**conv_kwargs) if config.separable_conv else ConvBnAct2d(
**conv_kwargs)
self.conv_rep.append(conv)
bn_levels = []
for _ in range(config.num_levels):
bn_seq = nn.Sequential()
bn_seq.add_module(
'bn', norm_layer(config.fpn_channels, **norm_kwargs))
bn_levels.append(bn_seq)
self.bn_rep.append(nn.ModuleList(bn_levels))
self.act = act_layer(inplace=True)
predict_kwargs = dict(in_channels=config.fpn_channels,
out_channels=num_outputs * num_anchors,
kernel_size=3,
padding=self.config.pad_type,
bias=True,
norm_layer=None,
act_layer=None)
if config.separable_conv:
self.predict = SeparableConv2d(**predict_kwargs)
else:
self.predict = ConvBnAct2d(**predict_kwargs)
def __init__(self, config, norm_kwargs=None, pretrained_backbone=True, alternate_init=False):
super(EfficientDet, self).__init__()
norm_kwargs = norm_kwargs or dict(eps=.001, momentum=.01)
self.backbone = create_model(
config.backbone_name, features_only=True, out_indices=(2, 3, 4),
pretrained=pretrained_backbone, **config.backbone_args)
feature_info = get_feature_info(self.backbone)
act_layer = get_act_layer(config.act_type)
self.fpn = BiFpn(config, feature_info, norm_kwargs=norm_kwargs, act_layer=act_layer)
self.class_net = HeadNet(config, num_outputs=config.num_classes, norm_kwargs=norm_kwargs, act_layer=act_layer)
self.box_net = HeadNet(config, num_outputs=4, norm_kwargs=norm_kwargs, act_layer=act_layer)
for n, m in self.named_modules():
if 'backbone' not in n:
if alternate_init:
_init_weight_alt(m, n)
else:
_init_weight(m, n)
def __init__(self, cfg: NfCfg, num_classes=1000, in_chans=3, global_pool='avg', output_stride=32,
drop_rate=0., drop_path_rate=0.):
super().__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
assert cfg.act_layer in _nonlin_gamma, f"Please add non-linearity constants for activation ({cfg.act_layer})."
conv_layer = ScaledStdConv2dSame if cfg.same_padding else ScaledStdConv2d
if cfg.gamma_in_act:
act_layer = act_with_gamma(cfg.act_layer, gamma=_nonlin_gamma[cfg.act_layer])
conv_layer = partial(conv_layer, eps=1e-4) # DM weights better with higher eps
else:
act_layer = get_act_layer(cfg.act_layer)
conv_layer = partial(conv_layer, gamma=_nonlin_gamma[cfg.act_layer])
attn_layer = partial(get_attn(cfg.attn_layer), **cfg.attn_kwargs) if cfg.attn_layer else None
stem_chs = make_divisible((cfg.stem_chs or cfg.channels[0]) * cfg.width_factor, cfg.ch_div)
self.stem, stem_stride, stem_feat = create_stem(
in_chans, stem_chs, cfg.stem_type, conv_layer=conv_layer, act_layer=act_layer)
self.feature_info = [stem_feat]
drop_path_rates = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths)).split(cfg.depths)]
prev_chs = stem_chs
net_stride = stem_stride
dilation = 1
expected_var = 1.0
stages = []
for stage_idx, stage_depth in enumerate(cfg.depths):
stride = 1 if stage_idx == 0 and stem_stride > 2 else 2
if net_stride >= output_stride and stride > 1:
dilation *= stride
stride = 1
net_stride *= stride
first_dilation = 1 if dilation in (1, 2) else 2
blocks = []
for block_idx in range(cfg.depths[stage_idx]):
first_block = block_idx == 0 and stage_idx == 0
out_chs = make_divisible(cfg.channels[stage_idx] * cfg.width_factor, cfg.ch_div)
blocks += [NormFreeBlock(
in_chs=prev_chs, out_chs=out_chs,
alpha=cfg.alpha,
beta=1. / expected_var ** 0.5,
stride=stride if block_idx == 0 else 1,
dilation=dilation,
first_dilation=first_dilation,
group_size=cfg.group_size,
bottle_ratio=1. if cfg.reg and first_block else cfg.bottle_ratio,
ch_div=cfg.ch_div,
reg=cfg.reg,
extra_conv=cfg.extra_conv,
skipinit=cfg.skipinit,
attn_layer=attn_layer,
attn_gain=cfg.attn_gain,
act_layer=act_layer,
conv_layer=conv_layer,
drop_path_rate=drop_path_rates[stage_idx][block_idx],
)]
if block_idx == 0:
expected_var = 1. # expected var is reset after first block of each stage
expected_var += cfg.alpha ** 2 # Even if reset occurs, increment expected variance
first_dilation = dilation
prev_chs = out_chs
self.feature_info += [dict(num_chs=prev_chs, reduction=net_stride, module=f'stages.{stage_idx}')]
stages += [nn.Sequential(*blocks)]
self.stages = nn.Sequential(*stages)
if cfg.num_features:
# The paper NFRegNet models have an EfficientNet-like final head convolution.
self.num_features = make_divisible(cfg.width_factor * cfg.num_features, cfg.ch_div)
self.final_conv = conv_layer(prev_chs, self.num_features, 1)
self.feature_info[-1] = dict(num_chs=self.num_features, reduction=net_stride, module=f'final_conv')
else:
self.num_features = prev_chs
self.final_conv = nn.Identity()
self.final_act = act_layer(inplace=cfg.num_features > 0)
self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=self.drop_rate)
for n, m in self.named_modules():
if 'fc' in n and isinstance(m, nn.Linear):
if cfg.zero_init_fc:
nn.init.zeros_(m.weight)
else:
nn.init.normal_(m.weight, 0., .01)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='linear')
if m.bias is not None:
nn.init.zeros_(m.bias)