Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torch.nn import Module as Module | |
| from collections import OrderedDict | |
| from pipeline.models.tresnet.layers.anti_aliasing import AntiAliasDownsampleLayer | |
| from .layers.avg_pool import FastAvgPool2d | |
| from .layers.general_layers import SEModule, SpaceToDepthModule | |
| from inplace_abn import InPlaceABN, ABN | |
| import torch.nn.functional as F | |
| def InplacABN_to_ABN(module: nn.Module) -> nn.Module: | |
| # convert all InplaceABN layer to bit-accurate ABN layers. | |
| if isinstance(module, InPlaceABN): | |
| module_new = ABN(module.num_features, activation=module.activation, | |
| activation_param=module.activation_param) | |
| for key in module.state_dict(): | |
| module_new.state_dict()[key].copy_(module.state_dict()[key]) | |
| module_new.training = module.training | |
| module_new.weight.data = module_new.weight.abs() + module_new.eps | |
| return module_new | |
| for name, child in reversed(module._modules.items()): | |
| new_child = InplacABN_to_ABN(child) | |
| if new_child != child: | |
| module._modules[name] = new_child | |
| return module | |
| class bottleneck_head(nn.Module): | |
| def __init__(self, num_features, num_classes, bottleneck_features=200): | |
| super(bottleneck_head, self).__init__() | |
| self.embedding_generator = nn.ModuleList() | |
| self.embedding_generator.append(nn.Linear(num_features, bottleneck_features)) | |
| self.embedding_generator = nn.Sequential(*self.embedding_generator) | |
| self.FC = nn.Linear(bottleneck_features, num_classes) | |
| def forward(self, x): | |
| self.embedding = self.embedding_generator(x) | |
| logits = self.FC(self.embedding) | |
| return logits | |
| def conv2d(ni, nf, stride): | |
| return nn.Sequential( | |
| nn.Conv2d(ni, nf, kernel_size=3, stride=stride, padding=1, bias=False), | |
| nn.BatchNorm2d(nf), | |
| nn.ReLU(inplace=True) | |
| ) | |
| def conv2d_ABN(ni, nf, stride, activation="leaky_relu", kernel_size=3, activation_param=1e-2, groups=1): | |
| return nn.Sequential( | |
| nn.Conv2d(ni, nf, kernel_size=kernel_size, stride=stride, padding=kernel_size // 2, groups=groups, | |
| bias=False), | |
| InPlaceABN(num_features=nf, activation=activation, activation_param=activation_param) | |
| ) | |
| class BasicBlock(Module): | |
| expansion = 1 | |
| def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None): | |
| super(BasicBlock, self).__init__() | |
| if stride == 1: | |
| self.conv1 = conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3) | |
| else: | |
| if anti_alias_layer is None: | |
| self.conv1 = conv2d_ABN(inplanes, planes, stride=2, activation_param=1e-3) | |
| else: | |
| self.conv1 = nn.Sequential(conv2d_ABN(inplanes, planes, stride=1, activation_param=1e-3), | |
| anti_alias_layer(channels=planes, filt_size=3, stride=2)) | |
| self.conv2 = conv2d_ABN(planes, planes, stride=1, activation="identity") | |
| self.relu = nn.ReLU(inplace=True) | |
| self.downsample = downsample | |
| self.stride = stride | |
| reduce_layer_planes = max(planes * self.expansion // 4, 64) | |
| self.se = SEModule(planes * self.expansion, reduce_layer_planes) if use_se else None | |
| def forward(self, x): | |
| if self.downsample is not None: | |
| residual = self.downsample(x) | |
| else: | |
| residual = x | |
| out = self.conv1(x) | |
| out = self.conv2(out) | |
| if self.se is not None: out = self.se(out) | |
| out += residual | |
| out = self.relu(out) | |
| return out | |
| class Bottleneck(Module): | |
| expansion = 4 | |
| def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True, anti_alias_layer=None): | |
| super(Bottleneck, self).__init__() | |
| self.conv1 = conv2d_ABN(inplanes, planes, kernel_size=1, stride=1, activation="leaky_relu", | |
| activation_param=1e-3) | |
| if stride == 1: | |
| self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=1, activation="leaky_relu", | |
| activation_param=1e-3) | |
| else: | |
| if anti_alias_layer is None: | |
| self.conv2 = conv2d_ABN(planes, planes, kernel_size=3, stride=2, activation="leaky_relu", | |
| activation_param=1e-3) | |
| else: | |
| self.conv2 = nn.Sequential(conv2d_ABN(planes, planes, kernel_size=3, stride=1, | |
| activation="leaky_relu", activation_param=1e-3), | |
| anti_alias_layer(channels=planes, filt_size=3, stride=2)) | |
| self.conv3 = conv2d_ABN(planes, planes * self.expansion, kernel_size=1, stride=1, | |
| activation="identity") | |
| self.relu = nn.ReLU(inplace=True) | |
| self.downsample = downsample | |
| self.stride = stride | |
| reduce_layer_planes = max(planes * self.expansion // 8, 64) | |
| self.se = SEModule(planes, reduce_layer_planes) if use_se else None | |
| def forward(self, x): | |
| if self.downsample is not None: | |
| residual = self.downsample(x) | |
| else: | |
| residual = x | |
| out = self.conv1(x) | |
| out = self.conv2(out) | |
| if self.se is not None: out = self.se(out) | |
| out = self.conv3(out) | |
| out = out + residual # no inplace | |
| out = self.relu(out) | |
| return out | |
| class TResNet(Module): | |
| def __init__(self, layers, in_chans=3, num_classes=1000, width_factor=1.0, | |
| do_bottleneck_head=False,bottleneck_features=512): | |
| super(TResNet, self).__init__() | |
| # Loss function | |
| self.loss_func = F.binary_cross_entropy_with_logits | |
| # JIT layers | |
| space_to_depth = SpaceToDepthModule() | |
| anti_alias_layer = AntiAliasDownsampleLayer | |
| global_pool_layer = FastAvgPool2d(flatten=True) | |
| # TResnet stages | |
| self.inplanes = int(64 * width_factor) | |
| self.planes = int(64 * width_factor) | |
| conv1 = conv2d_ABN(in_chans * 16, self.planes, stride=1, kernel_size=3) | |
| layer1 = self._make_layer(BasicBlock, self.planes, layers[0], stride=1, use_se=True, | |
| anti_alias_layer=anti_alias_layer) # 56x56 | |
| layer2 = self._make_layer(BasicBlock, self.planes * 2, layers[1], stride=2, use_se=True, | |
| anti_alias_layer=anti_alias_layer) # 28x28 | |
| layer3 = self._make_layer(Bottleneck, self.planes * 4, layers[2], stride=2, use_se=True, | |
| anti_alias_layer=anti_alias_layer) # 14x14 | |
| layer4 = self._make_layer(Bottleneck, self.planes * 8, layers[3], stride=2, use_se=False, | |
| anti_alias_layer=anti_alias_layer) # 7x7 | |
| # body | |
| self.body = nn.Sequential(OrderedDict([ | |
| ('SpaceToDepth', space_to_depth), | |
| ('conv1', conv1), | |
| ('layer1', layer1), | |
| ('layer2', layer2), | |
| ('layer3', layer3), | |
| ('layer4', layer4)])) | |
| # head | |
| self.embeddings = [] | |
| self.global_pool = nn.Sequential(OrderedDict([('global_pool_layer', global_pool_layer)])) | |
| self.num_features = (self.planes * 8) * Bottleneck.expansion | |
| if do_bottleneck_head: | |
| fc = bottleneck_head(self.num_features, num_classes, | |
| bottleneck_features=bottleneck_features) | |
| else: | |
| fc = nn.Linear(self.num_features , num_classes) | |
| self.head = nn.Sequential(OrderedDict([('fc', fc)])) | |
| # model initilization | |
| 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): | |
| 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) | |
| def _make_layer(self, block, planes, blocks, stride=1, use_se=True, anti_alias_layer=None): | |
| 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_ABN(self.inplanes, planes * block.expansion, kernel_size=1, stride=1, | |
| activation="identity")] | |
| downsample = nn.Sequential(*layers) | |
| layers = [] | |
| layers.append(block(self.inplanes, planes, stride, downsample, use_se=use_se, | |
| anti_alias_layer=anti_alias_layer)) | |
| self.inplanes = planes * block.expansion | |
| for i in range(1, blocks): layers.append( | |
| block(self.inplanes, planes, use_se=use_se, anti_alias_layer=anti_alias_layer)) | |
| return nn.Sequential(*layers) | |
| def forward_train(self, x, target): | |
| x = self.body(x) | |
| self.embeddings = self.global_pool(x) | |
| logits = self.head(self.embeddings) | |
| loss = self.loss_func(logits, target, reduction="mean") | |
| return logits, loss | |
| def forward_test(self, x): | |
| x = self.body(x) | |
| self.embeddings = self.global_pool(x) | |
| logits = self.head(self.embeddings) | |
| return logits | |
| def forward(self, x, target=None): | |
| if target is not None: | |
| return self.forward_train(x, target) | |
| else: | |
| return self.forward_test(x) | |
| def TResnetM(num_classes): | |
| """Constructs a medium TResnet model. | |
| """ | |
| in_chans = 3 | |
| model = TResNet(layers=[3, 4, 11, 3], num_classes=num_classes, in_chans=in_chans) | |
| return model | |
| def TResnetL(num_classes): | |
| """Constructs a large TResnet model. | |
| """ | |
| in_chans = 3 | |
| do_bottleneck_head = False | |
| model = TResNet(layers=[4, 5, 18, 3], num_classes=num_classes, in_chans=in_chans, width_factor=1.2, | |
| do_bottleneck_head=do_bottleneck_head) | |
| return model | |
| def TResnetXL(num_classes): | |
| """Constructs a xlarge TResnet model. | |
| """ | |
| in_chans = 3 | |
| model = TResNet(layers=[4, 5, 24, 3], num_classes=num_classes, in_chans=in_chans, width_factor=1.3) | |
| return model | |