RabbitRUI/ruispace

1import torch2from torch import nn3 4__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']5 6 7def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):8    """3x3 convolution with padding"""9    return nn.Conv2d(in_planes,10                     out_planes,11                     kernel_size=3,12                     stride=stride,13                     padding=dilation,14                     groups=groups,15                     bias=False,16                     dilation=dilation)17 18 19def conv1x1(in_planes, out_planes, stride=1):20    """1x1 convolution"""21    return nn.Conv2d(in_planes,22                     out_planes,23                     kernel_size=1,24                     stride=stride,25                     bias=False)26 27 28class IBasicBlock(nn.Module):29    expansion = 130    def __init__(self, inplanes, planes, stride=1, downsample=None,31                 groups=1, base_width=64, dilation=1):32        super(IBasicBlock, self).__init__()33        if groups != 1 or base_width != 64:34            raise ValueError('BasicBlock only supports groups=1 and base_width=64')35        if dilation > 1:36            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")37        self.bn1 = nn.BatchNorm2d(inplanes, eps=1e-05,)38        self.conv1 = conv3x3(inplanes, planes)39        self.bn2 = nn.BatchNorm2d(planes, eps=1e-05,)40        self.prelu = nn.PReLU(planes)41        self.conv2 = conv3x3(planes, planes, stride)42        self.bn3 = nn.BatchNorm2d(planes, eps=1e-05,)43        self.downsample = downsample44        self.stride = stride45 46    def forward(self, x):47        identity = x48        out = self.bn1(x)49        out = self.conv1(out)50        out = self.bn2(out)51        out = self.prelu(out)52        out = self.conv2(out)53        out = self.bn3(out)54        if self.downsample is not None:55            identity = self.downsample(x)56        out += identity57        return out58 59 60class IResNet(nn.Module):61    fc_scale = 7 * 762    def __init__(self,63                 block, layers, dropout=0, num_features=512, zero_init_residual=False,64                 groups=1, width_per_group=64, replace_stride_with_dilation=None, fp16=False):65        super(IResNet, self).__init__()66        self.fp16 = fp1667        self.inplanes = 6468        self.dilation = 169        if replace_stride_with_dilation is None:70            replace_stride_with_dilation = [False, False, False]71        if len(replace_stride_with_dilation) != 3:72            raise ValueError("replace_stride_with_dilation should be None "73                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))74        self.groups = groups75        self.base_width = width_per_group76        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)77        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)78        self.prelu = nn.PReLU(self.inplanes)79        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)80        self.layer2 = self._make_layer(block,81                                       128,82                                       layers[1],83                                       stride=2,84                                       dilate=replace_stride_with_dilation[0])85        self.layer3 = self._make_layer(block,86                                       256,87                                       layers[2],88                                       stride=2,89                                       dilate=replace_stride_with_dilation[1])90        self.layer4 = self._make_layer(block,91                                       512,92                                       layers[3],93                                       stride=2,94                                       dilate=replace_stride_with_dilation[2])95        self.bn2 = nn.BatchNorm2d(512 * block.expansion, eps=1e-05,)96        self.dropout = nn.Dropout(p=dropout, inplace=True)97        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)98        self.features = nn.BatchNorm1d(num_features, eps=1e-05)99        nn.init.constant_(self.features.weight, 1.0)100        self.features.weight.requires_grad = False101 102        for m in self.modules():103            if isinstance(m, nn.Conv2d):104                nn.init.normal_(m.weight, 0, 0.1)105            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):106                nn.init.constant_(m.weight, 1)107                nn.init.constant_(m.bias, 0)108 109        if zero_init_residual:110            for m in self.modules():111                if isinstance(m, IBasicBlock):112                    nn.init.constant_(m.bn2.weight, 0)113 114    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):115        downsample = None116        previous_dilation = self.dilation117        if dilate:118            self.dilation *= stride119            stride = 1120        if stride != 1 or self.inplanes != planes * block.expansion:121            downsample = nn.Sequential(122                conv1x1(self.inplanes, planes * block.expansion, stride),123                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, ),124            )125        layers = []126        layers.append(127            block(self.inplanes, planes, stride, downsample, self.groups,128                  self.base_width, previous_dilation))129        self.inplanes = planes * block.expansion130        for _ in range(1, blocks):131            layers.append(132                block(self.inplanes,133                      planes,134                      groups=self.groups,135                      base_width=self.base_width,136                      dilation=self.dilation))137 138        return nn.Sequential(*layers)139 140    def forward(self, x):141        with torch.cuda.amp.autocast(self.fp16):142            x = self.conv1(x)143            x = self.bn1(x)144            x = self.prelu(x)145            x = self.layer1(x)146            x = self.layer2(x)147            x = self.layer3(x)148            x = self.layer4(x)149            x = self.bn2(x)150            x = torch.flatten(x, 1)151            x = self.dropout(x)152        x = self.fc(x.float() if self.fp16 else x)153        x = self.features(x)154        return x155 156 157def _iresnet(arch, block, layers, pretrained, progress, **kwargs):158    model = IResNet(block, layers, **kwargs)159    if pretrained:160        raise ValueError()161    return model162 163 164def iresnet18(pretrained=False, progress=True, **kwargs):165    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,166                    progress, **kwargs)167 168 169def iresnet34(pretrained=False, progress=True, **kwargs):170    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,171                    progress, **kwargs)172 173 174def iresnet50(pretrained=False, progress=True, **kwargs):175    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,176                    progress, **kwargs)177 178 179def iresnet100(pretrained=False, progress=True, **kwargs):180    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,181                    progress, **kwargs)182 183 184def iresnet200(pretrained=False, progress=True, **kwargs):185    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,186                    progress, **kwargs)187 188