Ross Wightman
3 years ago
parent
9a38416fbd
commit
5a8e1e643e
@ -0,0 +1,91 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
from .padding import get_padding
|
||||
from .conv2d_same import conv2d_same
|
||||
|
||||
|
||||
def get_weight(module):
|
||||
std, mean = torch.std_mean(module.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
|
||||
weight = (module.weight - mean) / (std + module.eps)
|
||||
return weight
|
||||
|
||||
|
||||
class StdConv2d(nn.Conv2d):
|
||||
"""Conv2d with Weight Standardization. Used for BiT ResNet-V2 models.
|
||||
|
||||
Paper: `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization` -
|
||||
https://arxiv.org/abs/1903.10520v2
|
||||
"""
|
||||
def __init__(
|
||||
self, in_channel, out_channels, kernel_size, stride=1,
|
||||
padding=None, dilation=1, groups=1, bias=False, eps=1e-5):
|
||||
if padding is None:
|
||||
padding = get_padding(kernel_size, stride, dilation)
|
||||
super().__init__(
|
||||
in_channel, out_channels, kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
self.eps = eps
|
||||
|
||||
def get_weight(self):
|
||||
std, mean = torch.std_mean(self.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
|
||||
weight = (self.weight - mean) / (std + self.eps)
|
||||
return weight
|
||||
|
||||
def forward(self, x):
|
||||
x = F.conv2d(x, self.get_weight(), self.bias, self.stride, self.padding, self.dilation, self.groups)
|
||||
return x
|
||||
|
||||
|
||||
class StdConv2dSame(nn.Conv2d):
|
||||
"""Conv2d with Weight Standardization. TF compatible SAME padding. Used for ViT Hybrid model.
|
||||
|
||||
Paper: `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization` -
|
||||
https://arxiv.org/abs/1903.10520v2
|
||||
"""
|
||||
def __init__(
|
||||
self, in_channel, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=False, eps=1e-5):
|
||||
super().__init__(
|
||||
in_channel, out_channels, kernel_size, stride=stride,
|
||||
padding=0, dilation=dilation, groups=groups, bias=bias)
|
||||
self.eps = eps
|
||||
|
||||
def get_weight(self):
|
||||
std, mean = torch.std_mean(self.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
|
||||
weight = (self.weight - mean) / (std + self.eps)
|
||||
return weight
|
||||
|
||||
def forward(self, x):
|
||||
x = conv2d_same(x, self.get_weight(), self.bias, self.stride, self.padding, self.dilation, self.groups)
|
||||
return x
|
||||
|
||||
|
||||
class ScaledStdConv2d(nn.Conv2d):
|
||||
"""Conv2d layer with Scaled Weight Standardization.
|
||||
|
||||
Paper: `Characterizing signal propagation to close the performance gap in unnormalized ResNets` -
|
||||
https://arxiv.org/abs/2101.08692
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels, out_channels, kernel_size,
|
||||
stride=1, padding=None, dilation=1, groups=1, bias=True, gain=True, gamma=1.0, eps=1e-5):
|
||||
if padding is None:
|
||||
padding = get_padding(kernel_size, stride, dilation)
|
||||
super().__init__(
|
||||
in_channels, out_channels, kernel_size, stride=stride,
|
||||
padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
self.gain = nn.Parameter(torch.ones(self.out_channels, 1, 1, 1)) if gain else None
|
||||
self.gamma = gamma * self.weight[0].numel() ** 0.5 # gamma * sqrt(fan-in)
|
||||
self.eps = eps
|
||||
|
||||
def get_weight(self):
|
||||
std, mean = torch.std_mean(self.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
|
||||
weight = (self.weight - mean) / (self.gamma * std + self.eps)
|
||||
if self.gain is not None:
|
||||
weight = weight * self.gain
|
||||
return weight
|
||||
|
||||
def forward(self, x):
|
||||
return F.conv2d(x, self.get_weight(), self.bias, self.stride, self.padding, self.dilation, self.groups)
|
||||
@ -0,0 +1,441 @@
|
||||
""" Normalizer Free RegNet / ResNet (pre-activation) Models
|
||||
|
||||
Paper: `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
|
||||
- https://arxiv.org/abs/2101.08692
|
||||
|
||||
Hacked together by / copyright Ross Wightman, 2021.
|
||||
"""
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from collections import OrderedDict
|
||||
from typing import Tuple, Optional
|
||||
from functools import partial
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
|
||||
from .helpers import build_model_with_cfg
|
||||
from .registry import register_model
|
||||
from .layers import ClassifierHead, DropPath, AvgPool2dSame, ScaledStdConv2d, get_act_layer, get_attn, make_divisible
|
||||
|
||||
|
||||
def _dcfg(url='', **kwargs):
|
||||
return {
|
||||
'url': url,
|
||||
'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
|
||||
'crop_pct': 0.875, 'interpolation': 'bicubic',
|
||||
'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
|
||||
'first_conv': 'stem.conv', 'classifier': 'head.fc',
|
||||
**kwargs
|
||||
}
|
||||
|
||||
# FIXME finish
|
||||
default_cfgs = {
|
||||
'nf_regnet_b0': _dcfg(url=''),
|
||||
'nf_regnet_b1': _dcfg(url='', input_size=(3, 240, 240)),
|
||||
'nf_regnet_b2': _dcfg(url='', input_size=(3, 256, 256)),
|
||||
'nf_regnet_b3': _dcfg(url='', input_size=(3, 272, 272)),
|
||||
'nf_regnet_b4': _dcfg(url='', input_size=(3, 320, 320)),
|
||||
'nf_regnet_b5': _dcfg(url='', input_size=(3, 384, 384)),
|
||||
|
||||
'nf_resnet26d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_resnet50d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_resnet101d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
|
||||
'nf_seresnet26d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_seresnet50d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_seresnet101d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
|
||||
'nf_ecaresnet26d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_ecaresnet50d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
'nf_ecaresnet101d': _dcfg(url='', first_conv='stem.conv1'),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class NfCfg:
|
||||
depths: Tuple[int, int, int, int]
|
||||
channels: Tuple[int, int, int, int]
|
||||
alpha: float = 0.2
|
||||
stem_type: str = '3x3'
|
||||
stem_chs: Optional[int] = None
|
||||
group_size: Optional[int] = 8
|
||||
attn_layer: Optional[str] = 'se'
|
||||
attn_kwargs: dict = field(default_factory=lambda: dict(reduction_ratio=0.5, divisor=8))
|
||||
attn_gain: float = 2.0 # NF correction gain to apply if attn layer is used
|
||||
width_factor: float = 0.75
|
||||
bottle_ratio: float = 2.25
|
||||
efficient: bool = True # enables EfficientNet-like options that are used in paper 'nf_regnet_b*' models
|
||||
num_features: int = 1280 # num out_channels for final conv (when enabled in efficient mode)
|
||||
ch_div: int = 8 # round channels % 8 == 0 to keep tensor-core use optimal
|
||||
skipinit: bool = False
|
||||
act_layer: str = 'silu'
|
||||
|
||||
|
||||
model_cfgs = dict(
|
||||
# EffNet influenced RegNet defs
|
||||
nf_regnet_b0=NfCfg(depths=(1, 3, 6, 6), channels=(48, 104, 208, 440), num_features=1280),
|
||||
nf_regnet_b1=NfCfg(depths=(2, 4, 7, 7), channels=(48, 104, 208, 440), num_features=1280),
|
||||
nf_regnet_b2=NfCfg(depths=(2, 4, 8, 8), channels=(56, 112, 232, 488), num_features=1416),
|
||||
nf_regnet_b3=NfCfg(depths=(2, 5, 9, 9), channels=(56, 128, 248, 528), num_features=1536),
|
||||
nf_regnet_b4=NfCfg(depths=(2, 6, 11, 11), channels=(64, 144, 288, 616), num_features=1792),
|
||||
nf_regnet_b5=NfCfg(depths=(3, 7, 14, 14), channels=(80, 168, 336, 704), num_features=2048),
|
||||
|
||||
# ResNet (preact, D style deep stem/avg down) defs
|
||||
nf_resnet26d=NfCfg(
|
||||
depths=(2, 2, 2, 2), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer=None,),
|
||||
nf_resnet50d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer=None),
|
||||
nf_resnet101d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer=None),
|
||||
|
||||
|
||||
nf_seresnet26d=NfCfg(
|
||||
depths=(2, 2, 2, 2), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='se', attn_kwargs=dict(reduction_ratio=0.25)),
|
||||
nf_seresnet50d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='se', attn_kwargs=dict(reduction_ratio=0.25)),
|
||||
nf_seresnet101d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='se', attn_kwargs=dict(reduction_ratio=0.25)),
|
||||
|
||||
|
||||
nf_ecaresnet26d=NfCfg(
|
||||
depths=(2, 2, 2, 2), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='eca', attn_kwargs=dict()),
|
||||
nf_ecaresnet50d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='eca', attn_kwargs=dict()),
|
||||
nf_ecaresnet101d=NfCfg(
|
||||
depths=(3, 4, 6, 3), channels=(256, 512, 1024, 2048),
|
||||
stem_type='deep', stem_chs=64, width_factor=1.0, bottle_ratio=0.25, efficient=False, group_size=None,
|
||||
act_layer='relu', attn_layer='eca', attn_kwargs=dict()),
|
||||
|
||||
)
|
||||
|
||||
# class NormFreeSiLU(nn.Module):
|
||||
# _K = 1. / 0.5595
|
||||
# def __init__(self, inplace=False):
|
||||
# super().__init__()
|
||||
# self.inplace = inplace
|
||||
#
|
||||
# def forward(self, x):
|
||||
# return F.silu(x, inplace=self.inplace) * self._K
|
||||
#
|
||||
#
|
||||
# class NormFreeReLU(nn.Module):
|
||||
# _K = (0.5 * (1. - 1. / math.pi)) ** -0.5
|
||||
#
|
||||
# def __init__(self, inplace=False):
|
||||
# super().__init__()
|
||||
# self.inplace = inplace
|
||||
#
|
||||
# def forward(self, x):
|
||||
# return F.relu(x, inplace=self.inplace) * self._K
|
||||
|
||||
|
||||
class DownsampleAvg(nn.Module):
|
||||
def __init__(
|
||||
self, in_chs, out_chs, stride=1, dilation=1, first_dilation=None, conv_layer=ScaledStdConv2d):
|
||||
""" AvgPool Downsampling as in 'D' ResNet variants. Support for dilation."""
|
||||
super(DownsampleAvg, self).__init__()
|
||||
avg_stride = stride if dilation == 1 else 1
|
||||
if stride > 1 or dilation > 1:
|
||||
avg_pool_fn = AvgPool2dSame if avg_stride == 1 and dilation > 1 else nn.AvgPool2d
|
||||
self.pool = avg_pool_fn(2, avg_stride, ceil_mode=True, count_include_pad=False)
|
||||
else:
|
||||
self.pool = nn.Identity()
|
||||
self.conv = conv_layer(in_chs, out_chs, 1, stride=1)
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv(self.pool(x))
|
||||
|
||||
|
||||
class NormalizationFreeBlock(nn.Module):
|
||||
"""Normalization-free pre-activation block.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, in_chs, out_chs=None, stride=1, dilation=1, first_dilation=None,
|
||||
alpha=1.0, beta=1.0, bottle_ratio=0.25, efficient=True, ch_div=1, group_size=None,
|
||||
attn_layer=None, attn_gain=2.0, act_layer=None, conv_layer=None, drop_path_rate=0., skipinit=False):
|
||||
super().__init__()
|
||||
first_dilation = first_dilation or dilation
|
||||
out_chs = out_chs or in_chs
|
||||
# EfficientNet-like models scale bottleneck from in_chs, otherwise scale from out_chs like ResNet
|
||||
mid_chs = make_divisible(in_chs * bottle_ratio if efficient else out_chs * bottle_ratio, ch_div)
|
||||
groups = 1
|
||||
if group_size is not None:
|
||||
# NOTE: not correcting the mid_chs % group_size, fix model def if broken. I want % ch_div == 0 to stand.
|
||||
groups = mid_chs // group_size
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.attn_gain = attn_gain
|
||||
|
||||
if in_chs != out_chs or stride != 1 or dilation != first_dilation:
|
||||
self.downsample = DownsampleAvg(
|
||||
in_chs, out_chs, stride=stride, dilation=dilation, first_dilation=first_dilation, conv_layer=conv_layer)
|
||||
else:
|
||||
self.downsample = None
|
||||
|
||||
self.act1 = act_layer()
|
||||
self.conv1 = conv_layer(in_chs, mid_chs, 1)
|
||||
self.act2 = act_layer(inplace=True)
|
||||
self.conv2 = conv_layer(mid_chs, mid_chs, 3, stride=stride, dilation=first_dilation, groups=groups)
|
||||
if attn_layer is not None:
|
||||
self.attn = attn_layer(mid_chs)
|
||||
else:
|
||||
self.attn = None
|
||||
self.act3 = act_layer()
|
||||
self.conv3 = conv_layer(mid_chs, out_chs, 1)
|
||||
self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
|
||||
self.skipinit_gain = nn.Parameter(torch.tensor(0.)) if skipinit else None
|
||||
|
||||
def forward(self, x):
|
||||
out = self.act1(x) * self.beta
|
||||
|
||||
# shortcut branch
|
||||
shortcut = x
|
||||
if self.downsample is not None:
|
||||
shortcut = self.downsample(out)
|
||||
|
||||
# residual branch
|
||||
out = self.conv1(out)
|
||||
out = self.conv2(self.act2(out))
|
||||
if self.attn is not None:
|
||||
out = self.attn_gain * self.attn(out)
|
||||
out = self.conv3(self.act3(out))
|
||||
out = self.drop_path(out)
|
||||
if self.skipinit_gain is None:
|
||||
out = out * self.alpha + shortcut
|
||||
else:
|
||||
# this really slows things down for some reason, TBD
|
||||
out = out * self.alpha * self.skipinit_gain + shortcut
|
||||
return out
|
||||
|
||||
|
||||
def create_stem(in_chs, out_chs, stem_type='', conv_layer=None):
|
||||
stem = OrderedDict()
|
||||
assert stem_type in ('', 'deep', '3x3', '7x7')
|
||||
if 'deep' in stem_type:
|
||||
# 3 deep 3x3 conv stack as in ResNet V1D models
|
||||
mid_chs = out_chs // 2
|
||||
stem['conv1'] = conv_layer(in_chs, mid_chs, kernel_size=3, stride=2)
|
||||
stem['conv2'] = conv_layer(mid_chs, mid_chs, kernel_size=3, stride=1)
|
||||
stem['conv3'] = conv_layer(mid_chs, out_chs, kernel_size=3, stride=1)
|
||||
elif '3x3' in stem_type:
|
||||
# 3x3 stem conv as in RegNet
|
||||
stem['conv'] = conv_layer(in_chs, out_chs, kernel_size=3, stride=2)
|
||||
else:
|
||||
# 7x7 stem conv as in ResNet
|
||||
stem['conv'] = conv_layer(in_chs, out_chs, kernel_size=7, stride=2)
|
||||
|
||||
return nn.Sequential(stem)
|
||||
|
||||
|
||||
_nonlin_gamma = dict(
|
||||
silu=.5595,
|
||||
relu=(0.5 * (1. - 1. / math.pi)) ** 0.5,
|
||||
identity=1.0
|
||||
)
|
||||
|
||||
|
||||
class NormalizerFreeNet(nn.Module):
|
||||
""" Normalizer-free ResNets and RegNets
|
||||
|
||||
As described in `Characterizing signal propagation to close the performance gap in unnormalized ResNets`
|
||||
- https://arxiv.org/abs/2101.08692
|
||||
|
||||
This model aims to cover both the NFRegNet-Bx models as detailed in the paper's code snippets and
|
||||
the (preact) ResNet models described earlier in the paper.
|
||||
|
||||
There are a few differences:
|
||||
* channels are rounded to be divisible by 8 by default (keep TC happy), this changes param counts
|
||||
* activation correcting gamma constants are moved into the ScaledStdConv as it has less performance
|
||||
impact in PyTorch when done with the weight scaling there. This likely wasn't a concern in the JAX impl.
|
||||
* skipinit is disabled by default, it seems to have a rather drastic impact on GPU memory use and throughput
|
||||
for what it is/does. Approx 8-10% throughput loss.
|
||||
"""
|
||||
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
|
||||
act_layer = get_act_layer(cfg.act_layer)
|
||||
assert cfg.act_layer in _nonlin_gamma, f"Please add non-linearity constants for activation ({cfg.act_layer})."
|
||||
conv_layer = partial(ScaledStdConv2d, bias=True, gain=True, gamma=_nonlin_gamma[cfg.act_layer])
|
||||
attn_layer = partial(get_attn(cfg.attn_layer), **cfg.attn_kwargs) if cfg.attn_layer else None
|
||||
|
||||
self.feature_info = [] # FIXME fill out feature info
|
||||
|
||||
stem_chs = cfg.stem_chs or cfg.channels[0]
|
||||
stem_chs = make_divisible(stem_chs * cfg.width_factor, cfg.ch_div)
|
||||
self.stem = create_stem(in_chans, stem_chs, cfg.stem_type, conv_layer=conv_layer)
|
||||
|
||||
prev_chs = stem_chs
|
||||
dpr = [x.tolist() for x in torch.linspace(0, drop_path_rate, sum(cfg.depths)).split(cfg.depths)]
|
||||
net_stride = 2
|
||||
dilation = 1
|
||||
expected_var = 1.0
|
||||
stages = []
|
||||
for stage_idx, stage_depth in enumerate(cfg.depths):
|
||||
if net_stride >= output_stride:
|
||||
dilation *= 2
|
||||
stride = 1
|
||||
else:
|
||||
stride = 2
|
||||
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 += [NormalizationFreeBlock(
|
||||
in_chs=prev_chs, out_chs=out_chs,
|
||||
alpha=cfg.alpha,
|
||||
beta=1. / expected_var ** 0.5, # NOTE: beta used as multiplier in block
|
||||
stride=stride if block_idx == 0 else 1,
|
||||
dilation=dilation,
|
||||
first_dilation=first_dilation,
|
||||
group_size=cfg.group_size,
|
||||
bottle_ratio=1. if cfg.efficient and first_block else cfg.bottle_ratio,
|
||||
efficient=cfg.efficient,
|
||||
ch_div=cfg.ch_div,
|
||||
attn_layer=attn_layer,
|
||||
attn_gain=cfg.attn_gain,
|
||||
act_layer=act_layer,
|
||||
conv_layer=conv_layer,
|
||||
drop_path_rate=dpr[stage_idx][block_idx],
|
||||
skipinit=cfg.skipinit,
|
||||
)]
|
||||
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
|
||||
stages += [nn.Sequential(*blocks)]
|
||||
self.stages = nn.Sequential(*stages)
|
||||
|
||||
if cfg.efficient and 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)
|
||||
else:
|
||||
self.num_features = prev_chs
|
||||
self.final_conv = nn.Identity()
|
||||
self.final_act = act_layer()
|
||||
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):
|
||||
nn.init.zeros_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Conv2d):
|
||||
# as per discussion with paper authors, original in haiku is
|
||||
# hk.initializers.VarianceScaling(1.0, 'fan_in', 'normal')' w/ zero'd bias
|
||||
nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='linear')
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def get_classifier(self):
|
||||
return self.head.fc
|
||||
|
||||
def reset_classifier(self, num_classes, global_pool='avg'):
|
||||
self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=self.drop_rate)
|
||||
|
||||
def forward_features(self, x):
|
||||
x = self.stem(x)
|
||||
x = self.stages(x)
|
||||
x = self.final_conv(x)
|
||||
x = self.final_act(x)
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
x = self.forward_features(x)
|
||||
x = self.head(x)
|
||||
return x
|
||||
|
||||
|
||||
def _create_normfreenet(variant, pretrained=False, **kwargs):
|
||||
feature_cfg = dict(flatten_sequential=True)
|
||||
feature_cfg['feature_cls'] = 'hook' # pre-act models need hooks to grab feat from act1 in bottleneck blocks
|
||||
|
||||
return build_model_with_cfg(
|
||||
NormalizerFreeNet, variant, pretrained, model_cfg=model_cfgs[variant], default_cfg=default_cfgs[variant],
|
||||
feature_cfg=feature_cfg, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b0(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b0', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b1(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b1', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b2(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b2', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b3(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b3', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b4(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b4', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_regnet_b5(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_regnet_b5', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_resnet26d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_resnet26d', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_resnet50d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_resnet50d', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_seresnet26d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_seresnet26d', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_seresnet50d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_seresnet50d', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_ecaresnet26d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_ecaresnet26d', pretrained=pretrained, **kwargs)
|
||||
|
||||
|
||||
@register_model
|
||||
def nf_ecaresnet50d(pretrained=False, **kwargs):
|
||||
return _create_normfreenet('nf_ecaresnet50d', pretrained=pretrained, **kwargs)
|
||||
Loading…
Reference in new issue