Spaces:
Paused
Paused
| import torch | |
| from torch import nn | |
| assert torch.__version__ >= "1.8.1" | |
| from torch.utils.checkpoint import checkpoint_sequential | |
| __all__ = ['iresnet2060'] | |
| def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): | |
| """3x3 convolution with padding""" | |
| return nn.Conv2d(in_planes, | |
| out_planes, | |
| kernel_size=3, | |
| stride=stride, | |
| padding=dilation, | |
| groups=groups, | |
| bias=False, | |
| dilation=dilation) | |
| def conv1x1(in_planes, out_planes, stride=1): | |
| """1x1 convolution""" | |
| return nn.Conv2d(in_planes, | |
| out_planes, | |
| kernel_size=1, | |
| stride=stride, | |
| bias=False) | |
| class IBasicBlock(nn.Module): | |
| expansion = 1 | |
| def __init__(self, inplanes, planes, stride=1, downsample=None, | |
| groups=1, base_width=64, dilation=1): | |
| super(IBasicBlock, self).__init__() | |
| if groups != 1 or base_width != 64: | |
| raise ValueError('BasicBlock only supports groups=1 and base_width=64') | |
| if dilation > 1: | |
| raise NotImplementedError("Dilation > 1 not supported in BasicBlock") | |
| self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05, ) | |
| self.conv1 = conv3x3(inplanes, planes) | |
| self.bn2 = nn.BatchNorm2d(planes, eps=1e-05, ) | |
| self.prelu = nn.PReLU(planes) | |
| self.conv2 = conv3x3(planes, planes, stride) | |
| self.bn3 = nn.BatchNorm2d(planes, eps=1e-05, ) | |
| self.downsample = downsample | |
| self.stride = stride | |
| def forward(self, x): | |
| identity = x | |
| out = self.bn1(x) | |
| out = self.conv1(out) | |
| out = self.bn2(out) | |
| out = self.prelu(out) | |
| out = self.conv2(out) | |
| out = self.bn3(out) | |
| if self.downsample is not None: | |
| identity = self.downsample(x) | |
| out += identity | |
| return out | |
| class IResNet(nn.Module): | |
| fc_scale = 7 * 7 | |
| def __init__(self, | |
| block, layers, dropout=0, num_features=512, zero_init_residual=False, | |
| groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False): | |
| super(IResNet, self).__init__() | |
| self.fp16 = fp16 | |
| self.inplanes = 64 | |
| self.dilation = 1 | |
| if replace_stride_with_dilation is None: | |
| replace_stride_with_dilation = [False, False, False] | |
| if len(replace_stride_with_dilation) != 3: | |
| raise ValueError("replace_stride_with_dilation should be None " | |
| "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) | |
| self.groups = groups | |
| self.base_width = width_per_group | |
| self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False) | |
| self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05) | |
| self.prelu = nn.PReLU(self.inplanes) | |
| self.layer1 = self._make_layer(block, 64, layers[0], stride=2) | |
| self.layer2 = self._make_layer(block, | |
| 128, | |
| layers[1], | |
| stride=2, | |
| dilate=replace_stride_with_dilation[0]) | |
| self.layer3 = self._make_layer(block, | |
| 256, | |
| layers[2], | |
| stride=2, | |
| dilate=replace_stride_with_dilation[1]) | |
| self.layer4 = self._make_layer(block, | |
| 512, | |
| layers[3], | |
| stride=2, | |
| dilate=replace_stride_with_dilation[2]) | |
| self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05, ) | |
| self.dropout = nn.Dropout(p=dropout, inplace=True) | |
| self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features) | |
| self.features = nn.BatchNorm1d(num_features, eps=1e-05) | |
| nn.init.constant_(self.features.weight, 1.0) | |
| self.features.weight.requires_grad = False | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv2d): | |
| nn.init.normal_(m.weight, 0, 0.1) | |
| elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): | |
| nn.init.constant_(m.weight, 1) | |
| nn.init.constant_(m.bias, 0) | |
| if zero_init_residual: | |
| for m in self.modules(): | |
| if isinstance(m, IBasicBlock): | |
| nn.init.constant_(m.bn2.weight, 0) | |
| def _make_layer(self, block, planes, blocks, stride=1, dilate=False): | |
| downsample = None | |
| previous_dilation = self.dilation | |
| if dilate: | |
| self.dilation *= stride | |
| stride = 1 | |
| if stride != 1 or self.inplanes != planes * block.expansion: | |
| downsample = nn.Sequential( | |
| conv1x1(self.inplanes, planes * block.expansion, stride), | |
| nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ), | |
| ) | |
| layers = [] | |
| layers.append( | |
| block(self.inplanes, planes, stride, downsample, self.groups, | |
| self.base_width, previous_dilation)) | |
| self.inplanes = planes * block.expansion | |
| for _ in range(1, blocks): | |
| layers.append( | |
| block(self.inplanes, | |
| planes, | |
| groups=self.groups, | |
| base_width=self.base_width, | |
| dilation=self.dilation)) | |
| return nn.Sequential(*layers) | |
| def checkpoint(self, func, num_seg, x): | |
| if self.training: | |
| return checkpoint_sequential(func, num_seg, x) | |
| else: | |
| return func(x) | |
| def forward(self, x): | |
| with torch.cuda.amp.autocast(self.fp16): | |
| x = self.conv1(x) | |
| x = self.bn1(x) | |
| x = self.prelu(x) | |
| x = self.layer1(x) | |
| x = self.checkpoint(self.layer2, 20, x) | |
| x = self.checkpoint(self.layer3, 100, x) | |
| x = self.layer4(x) | |
| x = self.bn2(x) | |
| x = torch.flatten(x, 1) | |
| x = self.dropout(x) | |
| x = self.fc(x.float() if self.fp16 else x) | |
| x = self.features(x) | |
| return x | |
| def _iresnet(arch, block, layers, pretrained, progress, **kwargs): | |
| model = IResNet(block, layers, **kwargs) | |
| if pretrained: | |
| raise ValueError() | |
| return model | |
| def iresnet2060(pretrained=False, progress=True, **kwargs): | |
| return _iresnet('iresnet2060', IBasicBlock, [3, 128, 1024 - 128, 3], pretrained, progress, **kwargs) | |