Spaces:
Running
on
L40S
Running
on
L40S
File size: 4,721 Bytes
38e20ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
# The implementation is adopted from TFace,made pubicly available under the Apache-2.0 license at
# https://github.com/Tencent/TFace/blob/master/recognition/torchkit/backbone/model_resnet.py
import torch.nn as nn
from torch.nn import BatchNorm1d, BatchNorm2d, Conv2d, Dropout, Linear, MaxPool2d, Module, ReLU, Sequential
from .common import initialize_weights
def conv3x3(in_planes, out_planes, stride=1):
""" 3x3 convolution with padding
"""
return Conv2d(
in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
def conv1x1(in_planes, out_planes, stride=1):
""" 1x1 convolution
"""
return Conv2d(
in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
class Bottleneck(Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = conv1x1(inplanes, planes)
self.bn1 = BatchNorm2d(planes)
self.conv2 = conv3x3(planes, planes, stride)
self.bn2 = BatchNorm2d(planes)
self.conv3 = conv1x1(planes, planes * self.expansion)
self.bn3 = BatchNorm2d(planes * self.expansion)
self.relu = ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
out = self.conv3(out)
out = self.bn3(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class ResNet(Module):
""" ResNet backbone
"""
def __init__(self, input_size, block, layers, zero_init_residual=True):
""" Args:
input_size: input_size of backbone
block: block function
layers: layers in each block
"""
super(ResNet, self).__init__()
assert input_size[0] in [112, 224], \
'input_size should be [112, 112] or [224, 224]'
self.inplanes = 64
self.conv1 = Conv2d(
3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = BatchNorm2d(64)
self.relu = ReLU(inplace=True)
self.maxpool = MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.bn_o1 = BatchNorm2d(2048)
self.dropout = Dropout()
if input_size[0] == 112:
self.fc = Linear(2048 * 4 * 4, 512)
else:
self.fc = Linear(2048 * 7 * 7, 512)
self.bn_o2 = BatchNorm1d(512)
initialize_weights(self.modules)
if zero_init_residual:
for m in self.modules():
if isinstance(m, Bottleneck):
nn.init.constant_(m.bn3.weight, 0)
def _make_layer(self, block, planes, blocks, stride=1):
downsample = None
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = Sequential(
conv1x1(self.inplanes, planes * block.expansion, stride),
BatchNorm2d(planes * block.expansion),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes))
return Sequential(*layers)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.bn_o1(x)
x = self.dropout(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
x = self.bn_o2(x)
return x
def ResNet_50(input_size, **kwargs):
""" Constructs a ResNet-50 model.
"""
model = ResNet(input_size, Bottleneck, [3, 4, 6, 3], **kwargs)
return model
def ResNet_101(input_size, **kwargs):
""" Constructs a ResNet-101 model.
"""
model = ResNet(input_size, Bottleneck, [3, 4, 23, 3], **kwargs)
return model
def ResNet_152(input_size, **kwargs):
""" Constructs a ResNet-152 model.
"""
model = ResNet(input_size, Bottleneck, [3, 8, 36, 3], **kwargs)
return model
|