def __init__(
self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
use_se=False,
dilation=1,
norm_layer=ABN,
norm_act="relu",
antialias=False,
):
super(Bottleneck, self).__init__()
antialias = antialias and stride == 2
width = int(math.floor(planes * (base_width / 64)) * groups)
outplanes = planes * self.expansion
self.conv1 = conv1x1(inplanes, width)
self.bn1 = norm_layer(width, activation=norm_act)
conv2_stride = 1 if antialias else stride
self.conv2 = conv3x3(width, width, conv2_stride, groups, dilation)
self.bn2 = norm_layer(width, activation=norm_act)
self.conv3 = conv1x1(width, outplanes)
self.bn3 = norm_layer(outplanes, activation="identity")
self.se_module = SEModule(outplanes, planes // 4) if use_se else None
self.final_act = activation_from_name(norm_act)
self.downsample = downsample
self.blurpool = BlurPool()
self.antialias = antialias
def __init__(
self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
use_se=False,
dilation=1,
norm_layer=ABN,
norm_act="relu",
antialias=False,
):
super(BasicBlock, self).__init__()
antialias = antialias and stride == 2
assert groups == 1, "BasicBlock only supports groups of 1"
assert base_width == 64, "BasicBlock doest not support changing base width"
outplanes = planes * self.expansion
conv1_stride = 1 if antialias else stride
self.conv1 = conv3x3(inplanes, planes, conv1_stride, dilation)
self.bn1 = norm_layer(planes, activation=norm_act)
self.conv2 = conv3x3(planes, outplanes)
self.bn2 = norm_layer(outplanes, activation="identity")
self.se_module = SEModule(outplanes, planes // 4) if use_se else None
self.final_act = activation_from_name(norm_act)
self.downsample = downsample
self.blurpool = BlurPool()
self.antialias = antialias
def __init__(
self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
attn_type=None,
dilation=1,
norm_layer=ABN,
norm_act="relu",
antialias=False,
keep_prob=1, # for drop connect
):
super(Bottleneck, self).__init__()
antialias = antialias and stride == 2
width = int(math.floor(planes * (base_width / 64)) * groups)
outplanes = planes * self.expansion
self.conv1 = conv1x1(inplanes, width)
self.bn1 = norm_layer(width, activation=norm_act)
conv2_stride = 1 if antialias else stride
self.conv2 = conv3x3(width, width, conv2_stride, groups, dilation)
self.bn2 = norm_layer(width, activation=norm_act)
self.conv3 = conv1x1(width, outplanes)
self.bn3 = norm_layer(outplanes, activation="identity")
self.se_module = get_attn(attn_type)(outplanes, planes // 4)
self.final_act = activation_from_name(norm_act)
self.downsample = downsample
self.blurpool = BlurPool(
channels=width) if antialias else nn.Identity()
self.antialias = antialias
self.drop_connect = DropConnect(
keep_prob) if keep_prob < 1 else nn.Identity()
def __init__(
self,
in_chs,
mid_chs,
out_chs,
dw_kernel_size=3,
dw_str2_kernel_size=3,
stride=1,
attn_type=None,
keep_prob=1, # drop connect param
norm_layer=ABN,
norm_act="relu",
force_residual=False,
force_expansion=False, # always have expansion
):
super().__init__()
self.has_residual = in_chs == out_chs and stride == 1
self.force_residual = force_residual
if force_residual:
self.blurpool = BlurPool(channels=in_chs) if stride == 2 else nn.Identity()
self.in_chs = in_chs
if in_chs != mid_chs or force_expansion:
self.expansion = nn.Sequential(norm_layer(in_chs, activation=norm_act), conv1x1(in_chs, mid_chs))
else:
self.expansion = nn.Identity()
self.bn2 = norm_layer(mid_chs, activation=norm_act)
dw_kernel_size = dw_str2_kernel_size if stride == 2 else dw_kernel_size
self.conv_dw = nn.Conv2d(
mid_chs, mid_chs, dw_kernel_size, stride=stride, groups=mid_chs, bias=False, padding=dw_kernel_size // 2,
)
# some models like MobileNet use mid_chs here instead of in_channels. But I don't care for now
self.se = get_attn(attn_type)(mid_chs, in_chs // 4, norm_act)
self.bn3 = norm_layer(mid_chs, activation=norm_act) # Note it's NOT identity for PreAct
self.conv_pw1 = conv1x1(mid_chs, out_chs)
self.drop_connect = DropConnect(keep_prob) if keep_prob < 1 else nn.Identity()
def __init__(
self,
inplanes,
planes,
stride=1,
downsample=None,
groups=1,
base_width=64,
attn_type=None,
dilation=1,
norm_layer=ABN,
norm_act="relu",
antialias=False,
keep_prob=1,
):
super(BasicBlock, self).__init__()
antialias = antialias and stride == 2
assert groups == 1, "BasicBlock only supports groups of 1"
assert base_width == 64, "BasicBlock doest not support changing base width"
outplanes = planes * self.expansion
conv1_stride = 1 if antialias else stride
self.conv1 = conv3x3(inplanes, planes, conv1_stride, groups, dilation)
self.bn1 = norm_layer(planes, activation=norm_act)
self.conv2 = conv3x3(planes, outplanes)
self.bn2 = norm_layer(outplanes, activation="identity")
self.se_module = get_attn(attn_type)(outplanes, planes // 4)
self.final_act = activation_from_name(norm_act)
self.downsample = downsample
self.blurpool = BlurPool(
channels=planes) if antialias else nn.Identity()
self.antialias = antialias
self.drop_connect = DropConnect(
keep_prob) if keep_prob < 1 else nn.Identity()
def __init__(
self,
in_chs,
mid_chs,
out_chs,
stride=1,
norm_layer=ABN,
norm_act="relu",
keep_prob=1, # for drop connect
antialias=False,
):
super().__init__()
self.has_residual = in_chs == out_chs and stride == 1
self.stride = stride
if self.stride == 2: # only use Res2Net for stride == 1
self.blocks = nn.Sequential(
norm_layer(in_chs, activation=norm_act),
conv3x3(in_chs, mid_chs, stride=1 if antialias else 2),
BlurPool(channels=mid_chs) if antialias else nn.Identity(),
norm_layer(mid_chs, activation=norm_act),
conv3x3(mid_chs, out_chs),
)
else:
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.block_1 = nn.Sequential(
conv3x3(in_chs // 4, in_chs // 4), norm_layer(in_chs // 4, activation=norm_act),
)
self.block_2 = nn.Sequential(
conv3x3(in_chs // 4, in_chs // 4), norm_layer(in_chs // 4, activation=norm_act),
)
self.block_3 = nn.Sequential(
conv3x3(in_chs // 4, in_chs // 4), norm_layer(in_chs // 4, activation=norm_act),
)
self.last_conv = conv3x3(in_chs, out_chs) # expand in last conv in block
def __init__(
self,
in_chs,
mid_chs,
out_chs,
stride=1,
groups=1,
groups_width=None,
norm_layer=ABN,
norm_act="relu",
force_residual=False, # force residual in stride=2 blocks
# keep_prob=1, # for drop connect
):
super().__init__()
groups = mid_chs // groups_width if groups_width else groups
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.conv1 = conv1x1(in_chs, mid_chs)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, mid_chs, stride=stride, groups=groups)
self.bn3 = norm_layer(mid_chs, activation=norm_act)
# last conv is not followed by bn, but anyway bias here makes it slightly worse (on Imagenet)
self.conv3 = conv1x1(mid_chs, out_chs)
self.has_residual = in_chs == out_chs and stride == 1
self.force_residual = force_residual
if force_residual:
self.blurpool = BlurPool(channels=in_chs)
self.in_chs = in_chs
def __init__(
self,
in_chs,
mid_chs,
out_chs,
dw_kernel_size=3,
dw_str2_kernel_size=3,
stride=1,
attn_type=None,
keep_prob=1, # drop connect param
norm_layer=ABN,
norm_act="relu",
):
super().__init__()
self.blurpool = BlurPool(channels=in_chs) if stride == 2 else nn.Identity()
self.in_chs = in_chs
self.pw1 = nn.Sequential(norm_layer(in_chs, activation=norm_act), conv1x1(in_chs, mid_chs))
self.pw2 = nn.Sequential(norm_layer(mid_chs, activation=norm_act), conv1x1(mid_chs, out_chs))
dw_kernel_size = dw_str2_kernel_size if stride == 2 else dw_kernel_size
conv_dw = nn.Conv2d(
out_chs, out_chs, dw_kernel_size, stride=stride, groups=out_chs, bias=False, padding=dw_kernel_size // 2,
)
self.dw = nn.Sequential(norm_layer(out_chs, activation=norm_act), conv_dw)
self.se = get_attn(attn_type)(mid_chs, in_chs // 4, norm_act)
self.drop_connect = DropConnect(keep_prob) if keep_prob < 1 else nn.Identity()
def _make_layer(self,
planes,
blocks,
stride=1,
dilation=1,
use_se=None,
norm_layer=None,
norm_act=None,
antialias=None):
downsample = None
if stride != 1 or self.inplanes != planes * self.expansion:
downsample_layers = []
if antialias and stride == 2: # using OrderedDict to preserve ordering and allow loading
downsample_layers += [("blur", BlurPool())]
downsample_layers += [
("0",
conv1x1(self.inplanes,
planes * self.expansion,
stride=1 if antialias else stride)),
("1", norm_layer(planes * self.expansion,
activation="identity")),
]
downsample = nn.Sequential(OrderedDict(downsample_layers))
layers = [
self.block(
self.inplanes,
planes,
stride,
downsample,
self.groups,
self.base_width,
use_se,
dilation,
norm_layer,
norm_act,
antialias,
)
]
self.inplanes = planes * self.expansion
for _ in range(1, blocks):
layers.append(
self.block(
self.inplanes,
planes,
1,
None,
self.groups,
self.base_width,
use_se,
dilation,
norm_layer,
norm_act,
antialias,
))
return nn.Sequential(*layers)
def __init__(
self,
in_chs,
mid_chs,
out_chs,
stride=1,
groups=1,
groups_width=None,
norm_layer=ABN,
norm_act="relu",
keep_prob=1, # for drop connect
dim_reduction="stride & expand", # "expand -> stride", "stride & expand"
force_residual=False, # always have residual
):
super().__init__()
self.has_residual = in_chs == out_chs and stride == 1
self.force_residual = force_residual
if force_residual:
self.blurpool = BlurPool(channels=in_chs)
self.in_chs = in_chs
groups = in_chs // groups_width if groups_width else groups
if dim_reduction == "expand -> stride":
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.conv1 = conv3x3(in_chs, mid_chs)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, out_chs, stride=stride)
elif dim_reduction == "s2d":
if stride == 2:
# BN before S2D to make sure different pixels from one channel are normalized the same way
self.bn1 = nn.Sequential(norm_layer(in_chs, activation=norm_act), SpaceToDepth(block_size=2))
self.conv1 = conv3x3(in_chs * 4, mid_chs)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, out_chs)
else: # same as stride & expand
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.conv1 = conv3x3(in_chs, mid_chs, stride=stride)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, out_chs)
# elif dim_reduction == "stride -> expand":
# # it's ~20% faster to have stride first. maybe accuracy drop isn't that big
# # TODO: test MixConv type of block here. I expect it to have the same speed and N params
# # while performance should increase
# self.conv1 = conv3x3(in_chs, in_chs, stride=stride)
# self.bn1 = norm_layer(in_chs, activation=norm_act)
# self.conv2 = conv3x3(in_chs, out_chs)
elif dim_reduction == "stride & expand": # only this one is supported for now
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.conv1 = conv3x3(in_chs, mid_chs, stride=stride)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, out_chs)
elif dim_reduction == "mixconv stride & expand":
self.bn1 = norm_layer(in_chs, activation=norm_act)
self.conv1 = conv3x3(in_chs, mid_chs, stride=stride)
self.bn2 = norm_layer(mid_chs, activation=norm_act)
self.conv2 = conv3x3(mid_chs, out_chs)
else:
raise ValueError(f"{dim_reduction} is not valid dim reduction in PreAct BasicBlock")
def _make_layers(self, cfg):
layers = []
in_channels = self.in_channels
for v in cfg:
if v == "M":
if self.antialias:
layers += [
nn.MaxPool2d(kernel_size=2, stride=1),
BlurPool()
]
else:
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
layers += [
conv2d,
self.norm_layer(v, activation=self.norm_act)
]
in_channels = v
return nn.Sequential(*layers)
def __init__(
self,
block=None,
layers=None,
pretrained=None, # not used. here for proper signature
num_classes=1000,
in_channels=3,
use_se=False,
groups=1,
base_width=64,
deep_stem=False,
dilated=False,
norm_layer='abn',
norm_act='relu',
antialias=False,
encoder=False,
drop_rate=0.0,
global_pool='avg',
init_bn0=True):
stem_width = 64
if norm_layer.lower() == 'abn':
norm_act = 'relu'
norm_layer = bn_from_name(norm_layer)
self.inplanes = stem_width
self.num_classes = num_classes
self.groups = groups
self.base_width = base_width
self.drop_rate = drop_rate
self.block = block
self.expansion = block.expansion
self.dilated = dilated
self.norm_act = norm_act
super(ResNet, self).__init__()
if deep_stem:
self.conv1 = nn.Sequential(
conv3x3(in_channels, stem_width // 2, 2),
norm_layer(stem_width // 2, activation=norm_act),
conv3x3(stem_width // 2, stem_width // 2, 2),
norm_layer(stem_width // 2, activation=norm_act),
conv3x3(stem_width // 2, stem_width))
else:
self.conv1 = nn.Conv2d(in_channels,
stem_width,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = norm_layer(stem_width, activation=norm_act)
if deep_stem:
self.maxpool = nn.Sequential() # don't need it
elif antialias:
self.maxpool = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=1, padding=1), BlurPool())
else:
# for se resnets fist maxpool is slightly different
self.maxpool = nn.MaxPool2d(kernel_size=3,
stride=2,
padding=0 if use_se else 1,
ceil_mode=True if use_se else False)
# Output stride is 8 with dilated and 32 without
stride_3_4 = 1 if self.dilated else 2
dilation_3 = 2 if self.dilated else 1
dilation_4 = 4 if self.dilated else 1
largs = dict(use_se=use_se,
norm_layer=norm_layer,
norm_act=norm_act,
antialias=antialias)
self.layer1 = self._make_layer(64, layers[0], stride=1, **largs)
self.layer2 = self._make_layer(128, layers[1], stride=2, **largs)
self.layer3 = self._make_layer(256,
layers[2],
stride=stride_3_4,
dilation=dilation_3,
**largs)
self.layer4 = self._make_layer(512,
layers[3],
stride=stride_3_4,
dilation=dilation_4,
**largs)
self.global_pool = GlobalPool2d(global_pool)
self.num_features = 512 * self.expansion
self.encoder = encoder
if not encoder:
self.last_linear = nn.Linear(
self.num_features * self.global_pool.feat_mult(), num_classes)
else:
self.forward = self.encoder_features
self._initialize_weights(init_bn0)