1 import torch 2 import torch.nn as nn 3 from .utils import load_state_dict_from_url 4 5 6 __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 7 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d', 8 'wide_resnet50_2', 'wide_resnet101_2'] 9 10 11 model_urls = { 12 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 13 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 14 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 15 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 16 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 17 'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth', 18 'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth', 19 'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth', 20 'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth', 21 } 22 23 # 封装 24 def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 25 """3x3 convolution with padding""" 26 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 27 padding=dilation, groups=groups, bias=False, dilation=dilation) 28 29 30 def conv1x1(in_planes, out_planes, stride=1): 31 """1x1 convolution""" 32 return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 33 34 #block定义 35 class BasicBlock(nn.Module): 36 expansion = 1 37 #定义 38 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, 39 base_width=64, dilation=1, norm_layer=None): 40 super(BasicBlock, self).__init__() 41 if norm_layer is None: 42 norm_layer = nn.BatchNorm2d 43 if groups != 1 or base_width != 64: 44 raise ValueError('BasicBlock only supports groups=1 and base_width=64') 45 if dilation > 1: 46 raise NotImplementedError("Dilation > 1 not supported in BasicBlock") 47 # Both self.conv1 and self.downsample layers downsample the input when stride != 1 48 self.conv1 = conv3x3(inplanes, planes, stride) 49 self.bn1 = norm_layer(planes) # 归一化 50 self.relu = nn.ReLU(inplace=True) 51 self.conv2 = conv3x3(planes, planes) 52 self.bn2 = norm_layer(planes) 53 self.downsample = downsample 54 self.stride = stride 55 #实现 56 def forward(self, x): 57 # 保存x 做残差 58 identity = x 59 60 out = self.conv1(x) 61 out = self.bn1(out) 62 out = self.relu(out) 63 64 out = self.conv2(out) 65 out = self.bn2(out) 66 67 # 下采样 68 if self.downsample is not None: 69 identity = self.downsample(x) 70 71 out += identity # 先和 X 融合，再做relu 72 out = self.relu(out) 73 74 return out 75 76 # 瓶颈（50层以上） 77 class Bottleneck(nn.Module): 78 # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2) 79 # while original implementation places the stride at the first 1x1 convolution(self.conv1) 80 # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385. 81 # This variant is also known as ResNet V1.5 and improves accuracy according to 82 # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch. 83 84 # 通道放大的倍数 85 expansion = 4 86 87 def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, 88 base_width=64, dilation=1, norm_layer=None): 89 super(Bottleneck, self).__init__() 90 if norm_layer is None: 91 norm_layer = nn.BatchNorm2d 92 width = int(planes * (base_width / 64.)) * groups 93 # Both self.conv2 and self.downsample layers downsample the input when stride != 1 94 self.conv1 = conv1x1(inplanes, width) 95 self.bn1 = norm_layer(width) 96 self.conv2 = conv3x3(width, width, stride, groups, dilation) 97 self.bn2 = norm_layer(width) 98 self.conv3 = conv1x1(width, planes * self.expansion) 99 self.bn3 = norm_layer(planes * self.expansion) 100 self.relu = nn.ReLU(inplace=True) 101 self.downsample = downsample 102 self.stride = stride 103 104 # 网络前向传播过程（调用过程） 105 def forward(self, x): 106 identity = x 107 108 out = self.conv1(x) 109 out = self.bn1(out) 110 out = self.relu(out) 111 112 out = self.conv2(out) 113 out = self.bn2(out) 114 out = self.relu(out) 115 116 out = self.conv3(out) 117 out = self.bn3(out) 118 119 if self.downsample is not None: 120 identity = self.downsample(x) 121 122 out += identity 123 out = self.relu(out) 124 125 return out 126 127 128 class ResNet(nn.Module): 129 130 def __init__(self, block, layers, num_classes=1000, zero_init_residual=False, 131 groups=1, width_per_group=64, replace_stride_with_dilation=None, 132 norm_layer=None): 133 super(ResNet, self).__init__() 134 if norm_layer is None: 135 norm_layer = nn.BatchNorm2d 136 self._norm_layer = norm_layer 137 138 self.inplanes = 64 139 self.dilation = 1 140 if replace_stride_with_dilation is None: 141 # each element in the tuple indicates if we should replace 142 # the 2x2 stride with a dilated convolution instead 143 replace_stride_with_dilation = [False, False, False] 144 if len(replace_stride_with_dilation) != 3: 145 raise ValueError("replace_stride_with_dilation should be None " 146 "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 147 self.groups = groups 148 self.base_width = width_per_group 149 self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, 150 bias=False) 151 self.bn1 = norm_layer(self.inplanes) 152 self.relu = nn.ReLU(inplace=True) 153 self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 154 #每一层stage 155 self.layer1 = self._make_layer(block, 64, layers[0]) 156 self.layer2 = self._make_layer(block, 128, layers[1], stride=2, 157 dilate=replace_stride_with_dilation[0]) 158 self.layer3 = self._make_layer(block, 256, layers[2], stride=2, 159 dilate=replace_stride_with_dilation[1]) 160 self.layer4 = self._make_layer(block, 512, layers[3], stride=2, 161 dilate=replace_stride_with_dilation[2]) 162 self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) 163 self.fc = nn.Linear(512 * block.expansion, num_classes) 164 165 #参数初始化 166 for m in self.modules(): 167 if isinstance(m, nn.Conv2d): 168 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 169 elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 170 nn.init.constant_(m.weight, 1) 171 nn.init.constant_(m.bias, 0) 172 173 # Zero-initialize the last BN in each residual branch, 174 # so that the residual branch starts with zeros, and each residual block behaves like an identity. 175 # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 176 if zero_init_residual: 177 for m in self.modules(): 178 if isinstance(m, Bottleneck): 179 nn.init.constant_(m.bn3.weight, 0) 180 elif isinstance(m, BasicBlock): 181 nn.init.constant_(m.bn2.weight, 0) 182 183 def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 184 norm_layer = self._norm_layer 185 downsample = None 186 previous_dilation = self.dilation 187 if dilate: 188 self.dilation *= stride 189 stride = 1 190 if stride != 1 or self.inplanes != planes * block.expansion: 191 downsample = nn.Sequential( 192 conv1x1(self.inplanes, planes * block.expansion, stride), 193 norm_layer(planes * block.expansion), 194 ) 195 196 layers = [] 197 layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 198 self.base_width, previous_dilation, norm_layer)) 199 self.inplanes = planes * block.expansion 200 for _ in range(1, blocks): 201 layers.append(block(self.inplanes, planes, groups=self.groups, 202 base_width=self.base_width, dilation=self.dilation, 203 norm_layer=norm_layer)) 204 205 return nn.Sequential(*layers) 206 207 def _forward_impl(self, x): 208 # See note [TorchScript super()] 209 x = self.conv1(x) 210 x = self.bn1(x) 211 x = self.relu(x) 212 x = self.maxpool(x) 213 214 x = self.layer1(x) 215 x = self.layer2(x) 216 x = self.layer3(x) 217 x = self.layer4(x) 218 219 x = self.avgpool(x) 220 x = torch.flatten(x, 1) 221 x = self.fc(x) 222 223 return x 224 225 def forward(self, x): 226 return self._forward_impl(x) 227 228 229 def _resnet(arch, block, layers, pretrained, progress, **kwargs): 230 model = ResNet(block, layers, **kwargs) 231 if pretrained: 232 state_dict = load_state_dict_from_url(model_urls[arch], 233 progress=progress) 234 model.load_state_dict(state_dict) 235 return model 236 237 238 def resnet18(pretrained=False, progress=True, **kwargs): 239 r"""ResNet-18 model from 240 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ 241 Args: 242 pretrained (bool): If True, returns a model pre-trained on ImageNet 243 progress (bool): If True, displays a progress bar of the download to stderr 244 """ 245 return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, 246 **kwargs) 247 248 249 def resnet34(pretrained=False, progress=True, **kwargs): 250 r"""ResNet-34 model from 251 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ 252 Args: 253 pretrained (bool): If True, returns a model pre-trained on ImageNet 254 progress (bool): If True, displays a progress bar of the download to stderr 255 """ 256 return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, 257 **kwargs) 258 259 260 def resnet50(pretrained=False, progress=True, **kwargs): 261 r"""ResNet-50 model from 262 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ 263 Args: 264 pretrained (bool): If True, returns a model pre-trained on ImageNet 265 progress (bool): If True, displays a progress bar of the download to stderr 266 """ 267 return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, 268 **kwargs) 269 270 271 def resnet101(pretrained=False, progress=True, **kwargs): 272 r"""ResNet-101 model from 273 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ 274 Args: 275 pretrained (bool): If True, returns a model pre-trained on ImageNet 276 progress (bool): If True, displays a progress bar of the download to stderr 277 """ 278 return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, 279 **kwargs) 280 281 282 def resnet152(pretrained=False, progress=True, **kwargs): 283 r"""ResNet-152 model from 284 `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ 285 Args: 286 pretrained (bool): If True, returns a model pre-trained on ImageNet 287 progress (bool): If True, displays a progress bar of the download to stderr 288 """ 289 return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, 290 **kwargs) 291 292 293 def resnext50_32x4d(pretrained=False, progress=True, **kwargs): 294 r"""ResNeXt-50 32x4d model from 295 `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_ 296 Args: 297 pretrained (bool): If True, returns a model pre-trained on ImageNet 298 progress (bool): If True, displays a progress bar of the download to stderr 299 """ 300 kwargs['groups'] = 32 301 kwargs['width_per_group'] = 4 302 return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], 303 pretrained, progress, **kwargs) 304 305 306 def resnext101_32x8d(pretrained=False, progress=True, **kwargs): 307 r"""ResNeXt-101 32x8d model from 308 `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_ 309 Args: 310 pretrained (bool): If True, returns a model pre-trained on ImageNet 311 progress (bool): If True, displays a progress bar of the download to stderr 312 """ 313 kwargs['groups'] = 32 314 kwargs['width_per_group'] = 8 315 return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], 316 pretrained, progress, **kwargs) 317 318 319 def wide_resnet50_2(pretrained=False, progress=True, **kwargs): 320 r"""Wide ResNet-50-2 model from 321 `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_ 322 The model is the same as ResNet except for the bottleneck number of channels 323 which is twice larger in every block. The number of channels in outer 1x1 324 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 325 channels, and in Wide ResNet-50-2 has 2048-1024-2048. 326 Args: 327 pretrained (bool): If True, returns a model pre-trained on ImageNet 328 progress (bool): If True, displays a progress bar of the download to stderr 329 """ 330 kwargs['width_per_group'] = 64 * 2 331 return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], 332 pretrained, progress, **kwargs) 333 334 335 def wide_resnet101_2(pretrained=False, progress=True, **kwargs): 336 r"""Wide ResNet-101-2 model from 337 `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_ 338 The model is the same as ResNet except for the bottleneck number of channels 339 which is twice larger in every block. The number of channels in outer 1x1 340 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 341 channels, and in Wide ResNet-50-2 has 2048-1024-2048. 342 Args: 343 pretrained (bool): If True, returns a model pre-trained on ImageNet 344 progress (bool): If True, displays a progress bar of the download to stderr 345 """ 346 kwargs['width_per_group'] = 64 * 2 347 return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], 348 pretrained, progress, **kwargs)