You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
pytorch-image-models/timm/models/tresnet.py

332 lines
12 KiB

5 years ago
"""
TResNet: High Performance GPU-Dedicated Architecture
https://arxiv.org/pdf/2003.13630.pdf
Original model: https://github.com/mrT23/TResNet
"""
from collections import OrderedDict
5 years ago
import torch
import torch.nn as nn
from timm.layers import SpaceToDepthModule, BlurPool2d, InplaceAbn, ClassifierHead, SEModule
from ._builder import build_model_with_cfg
from ._registry import register_model
5 years ago
__all__ = ['TResNet'] # model_registry will add each entrypoint fn to this
5 years ago
def _cfg(url='', **kwargs):
return {
'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7),
'crop_pct': 0.875, 'interpolation': 'bilinear',
'mean': (0., 0., 0.), 'std': (1., 1., 1.),
'first_conv': 'body.conv1.0', 'classifier': 'head.fc',
5 years ago
**kwargs
}
default_cfgs = {
'tresnet_m': _cfg(
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_m_1k_miil_83_1-d236afcb.pth'),
'tresnet_m_miil_in21k': _cfg(
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_m_miil_in21k-901b6ed4.pth', num_classes=11221),
'tresnet_l': _cfg(
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_l_81_5-235b486c.pth'),
'tresnet_xl': _cfg(
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_xl_82_0-a2d51b00.pth'),
'tresnet_m_448': _cfg(
input_size=(3, 448, 448), pool_size=(14, 14),
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_m_448-bc359d10.pth'),
'tresnet_l_448': _cfg(
input_size=(3, 448, 448), pool_size=(14, 14),
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_l_448-940d0cd1.pth'),
'tresnet_xl_448': _cfg(
input_size=(3, 448, 448), pool_size=(14, 14),
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_xl_448-8c1815de.pth'),
'tresnet_v2_l': _cfg(
url='https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-tresnet/tresnet_l_v2_83_9-f36e4445.pth'),
5 years ago
}
def IABN2Float(module: nn.Module) -> nn.Module:
"""If `module` is IABN don't use half precision."""
if isinstance(module, InplaceAbn):
5 years ago
module.float()
for child in module.children():
IABN2Float(child)
5 years ago
return module
def conv2d_iabn(ni, nf, stride, kernel_size=3, groups=1, act_layer="leaky_relu", act_param=1e-2):
5 years ago
return nn.Sequential(
nn.Conv2d(
ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups, bias=False),
InplaceAbn(nf, act_layer=act_layer, act_param=act_param)
5 years ago
)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, aa_layer=None):
5 years ago
super(BasicBlock, self).__init__()
if stride == 1:
self.conv1 = conv2d_iabn(inplanes, planes, stride=1, act_param=1e-3)
5 years ago
else:
if aa_layer is None:
self.conv1 = conv2d_iabn(inplanes, planes, stride=2, act_param=1e-3)
5 years ago
else:
self.conv1 = nn.Sequential(
conv2d_iabn(inplanes, planes, stride=1, act_param=1e-3),
aa_layer(channels=planes, filt_size=3, stride=2))
5 years ago
self.conv2 = conv2d_iabn(planes, planes, stride=1, act_layer="identity")
5 years ago
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
rd_chs = max(planes * self.expansion // 4, 64)
self.se = SEModule(planes * self.expansion, rd_channels=rd_chs) if use_se else None
5 years ago
def forward(self, x):
if self.downsample is not None:
shortcut = self.downsample(x)
5 years ago
else:
shortcut = x
5 years ago
out = self.conv1(x)
out = self.conv2(out)
if self.se is not None:
out = self.se(out)
5 years ago
out = out + shortcut
5 years ago
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(
self, inplanes, planes, stride=1, downsample=None, use_se=True,
act_layer="leaky_relu", aa_layer=None):
5 years ago
super(Bottleneck, self).__init__()
self.conv1 = conv2d_iabn(
inplanes, planes, kernel_size=1, stride=1, act_layer=act_layer, act_param=1e-3)
5 years ago
if stride == 1:
self.conv2 = conv2d_iabn(
planes, planes, kernel_size=3, stride=1, act_layer=act_layer, act_param=1e-3)
5 years ago
else:
if aa_layer is None:
self.conv2 = conv2d_iabn(
planes, planes, kernel_size=3, stride=2, act_layer=act_layer, act_param=1e-3)
5 years ago
else:
self.conv2 = nn.Sequential(
conv2d_iabn(planes, planes, kernel_size=3, stride=1, act_layer=act_layer, act_param=1e-3),
aa_layer(channels=planes, filt_size=3, stride=2))
5 years ago
reduction_chs = max(planes * self.expansion // 8, 64)
self.se = SEModule(planes, rd_channels=reduction_chs) if use_se else None
self.conv3 = conv2d_iabn(
planes, planes * self.expansion, kernel_size=1, stride=1, act_layer="identity")
5 years ago
self.act = nn.ReLU(inplace=True)
5 years ago
self.downsample = downsample
self.stride = stride
def forward(self, x):
if self.downsample is not None:
shortcut = self.downsample(x)
5 years ago
else:
shortcut = x
5 years ago
out = self.conv1(x)
out = self.conv2(out)
if self.se is not None:
out = self.se(out)
5 years ago
out = self.conv3(out)
out = out + shortcut # no inplace
out = self.act(out)
5 years ago
return out
class TResNet(nn.Module):
def __init__(
self,
layers,
in_chans=3,
num_classes=1000,
width_factor=1.0,
v2=False,
global_pool='fast',
drop_rate=0.,
):
self.num_classes = num_classes
self.drop_rate = drop_rate
5 years ago
super(TResNet, self).__init__()
aa_layer = BlurPool2d
5 years ago
# TResnet stages
self.inplanes = int(64 * width_factor)
self.planes = int(64 * width_factor)
if v2:
self.inplanes = self.inplanes // 8 * 8
self.planes = self.planes // 8 * 8
conv1 = conv2d_iabn(in_chans * 16, self.planes, stride=1, kernel_size=3)
layer1 = self._make_layer(
Bottleneck if v2 else BasicBlock, self.planes, layers[0], stride=1, use_se=True, aa_layer=aa_layer)
layer2 = self._make_layer(
Bottleneck if v2 else BasicBlock, self.planes * 2, layers[1], stride=2, use_se=True, aa_layer=aa_layer)
layer3 = self._make_layer(
Bottleneck, self.planes * 4, layers[2], stride=2, use_se=True, aa_layer=aa_layer)
layer4 = self._make_layer(
Bottleneck, self.planes * 8, layers[3], stride=2, use_se=False, aa_layer=aa_layer)
5 years ago
# body
self.body = nn.Sequential(OrderedDict([
('SpaceToDepth', SpaceToDepthModule()),
5 years ago
('conv1', conv1),
('layer1', layer1),
('layer2', layer2),
('layer3', layer3),
('layer4', layer4)]))
self.feature_info = [
dict(num_chs=self.planes, reduction=2, module=''), # Not with S2D?
dict(num_chs=self.planes * (Bottleneck.expansion if v2 else 1), reduction=4, module='body.layer1'),
dict(num_chs=self.planes * 2 * (Bottleneck.expansion if v2 else 1), reduction=8, module='body.layer2'),
dict(num_chs=self.planes * 4 * Bottleneck.expansion, reduction=16, module='body.layer3'),
dict(num_chs=self.planes * 8 * Bottleneck.expansion, reduction=32, module='body.layer4'),
]
5 years ago
# head
self.num_features = (self.planes * 8) * Bottleneck.expansion
self.head = ClassifierHead(self.num_features, num_classes, pool_type=global_pool, drop_rate=drop_rate)
5 years ago
# model initialization
5 years ago
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu')
elif isinstance(m, nn.BatchNorm2d) or isinstance(m, InplaceAbn):
5 years ago
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
# residual connections special initialization
for m in self.modules():
if isinstance(m, BasicBlock):
m.conv2[1].weight = nn.Parameter(torch.zeros_like(m.conv2[1].weight)) # BN to zero
if isinstance(m, Bottleneck):
m.conv3[1].weight = nn.Parameter(torch.zeros_like(m.conv3[1].weight)) # BN to zero
if isinstance(m, nn.Linear):
m.weight.data.normal_(0, 0.01)
5 years ago
def _make_layer(self, block, planes, blocks, stride=1, use_se=True, aa_layer=None):
5 years ago
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
layers = []
if stride == 2:
# avg pooling before 1x1 conv
layers.append(nn.AvgPool2d(kernel_size=2, stride=2, ceil_mode=True, count_include_pad=False))
layers += [conv2d_iabn(
self.inplanes, planes * block.expansion, kernel_size=1, stride=1, act_layer="identity")]
5 years ago
downsample = nn.Sequential(*layers)
layers = []
layers.append(block(
self.inplanes, planes, stride, downsample, use_se=use_se, aa_layer=aa_layer))
5 years ago
self.inplanes = planes * block.expansion
for i in range(1, blocks):
layers.append(
block(self.inplanes, planes, use_se=use_se, aa_layer=aa_layer))
5 years ago
return nn.Sequential(*layers)
@torch.jit.ignore
def group_matcher(self, coarse=False):
matcher = dict(stem=r'^body\.conv1', blocks=r'^body\.layer(\d+)' if coarse else r'^body\.layer(\d+)\.(\d+)')
return matcher
@torch.jit.ignore
def set_grad_checkpointing(self, enable=True):
assert not enable, 'gradient checkpointing not supported'
@torch.jit.ignore
def get_classifier(self):
return self.head.fc
5 years ago
def reset_classifier(self, num_classes, global_pool='fast'):
self.head = ClassifierHead(
self.num_features, num_classes, pool_type=global_pool, drop_rate=self.drop_rate)
5 years ago
def forward_features(self, x):
return self.body(x)
def forward_head(self, x, pre_logits: bool = False):
return x if pre_logits else self.head(x)
def forward(self, x):
x = self.forward_features(x)
x = self.forward_head(x)
return x
5 years ago
def _create_tresnet(variant, pretrained=False, **kwargs):
return build_model_with_cfg(
TResNet, variant, pretrained,
feature_cfg=dict(out_indices=(1, 2, 3, 4), flatten_sequential=True),
**kwargs)
5 years ago
@register_model
def tresnet_m(pretrained=False, **kwargs):
model_kwargs = dict(layers=[3, 4, 11, 3], **kwargs)
return _create_tresnet('tresnet_m', pretrained=pretrained, **model_kwargs)
5 years ago
@register_model
def tresnet_m_miil_in21k(pretrained=False, **kwargs):
model_kwargs = dict(layers=[3, 4, 11, 3], **kwargs)
return _create_tresnet('tresnet_m_miil_in21k', pretrained=pretrained, **model_kwargs)
5 years ago
@register_model
def tresnet_l(pretrained=False, **kwargs):
model_kwargs = dict(layers=[4, 5, 18, 3], width_factor=1.2, **kwargs)
return _create_tresnet('tresnet_l', pretrained=pretrained, **model_kwargs)
5 years ago
@register_model
def tresnet_v2_l(pretrained=False, **kwargs):
model_kwargs = dict(layers=[3, 4, 23, 3], width_factor=1.0, v2=True, **kwargs)
return _create_tresnet('tresnet_v2_l', pretrained=pretrained, **model_kwargs)
5 years ago
@register_model
def tresnet_xl(pretrained=False, **kwargs):
model_kwargs = dict(layers=[4, 5, 24, 3], width_factor=1.3, **kwargs)
return _create_tresnet('tresnet_xl', pretrained=pretrained, **model_kwargs)
@register_model
def tresnet_m_448(pretrained=False, **kwargs):
model_kwargs = dict(layers=[3, 4, 11, 3], **kwargs)
return _create_tresnet('tresnet_m_448', pretrained=pretrained, **model_kwargs)
@register_model
def tresnet_l_448(pretrained=False, **kwargs):
model_kwargs = dict(layers=[4, 5, 18, 3], width_factor=1.2, **kwargs)
return _create_tresnet('tresnet_l_448', pretrained=pretrained, **model_kwargs)
@register_model
def tresnet_xl_448(pretrained=False, **kwargs):
model_kwargs = dict(layers=[4, 5, 24, 3], width_factor=1.3, **kwargs)
return _create_tresnet('tresnet_xl_448', pretrained=pretrained, **model_kwargs)