| | import torch |
| | import torch.nn as nn |
| | import torch.utils.model_zoo as model_zoo |
| | import os |
| |
|
| | model_urls = { |
| | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', |
| | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', |
| | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', |
| | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', |
| | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', |
| | 'resnet18stem': 'https://download.openmmlab.com/pretrain/third_party/resnet18_v1c-b5776b93.pth', |
| | 'resnet50stem': 'https://download.openmmlab.com/pretrain/third_party/resnet50_v1c-2cccc1ad.pth', |
| | 'resnet101stem': 'https://download.openmmlab.com/pretrain/third_party/resnet101_v1c-e67eebb6.pth', |
| | } |
| |
|
| | def conv3x3(in_planes, outplanes, stride=1): |
| | |
| | return nn.Conv2d(in_planes, outplanes, kernel_size=3, stride=stride, padding=1, bias=False) |
| |
|
| | class BasicBlock(nn.Module): |
| | """ |
| | Basic Block for Resnet |
| | """ |
| | expansion = 1 |
| |
|
| | def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None): |
| | super(BasicBlock, self).__init__() |
| | self.conv1 = conv3x3(inplanes, planes, stride) |
| | self.bn1 = nn.BatchNorm2d(planes) |
| | self.relu = nn.ReLU(inplace=True) |
| | self.conv2 = conv3x3(planes, planes) |
| | self.bn2 = nn.BatchNorm2d(planes) |
| | self.downsample = downsample |
| | self.stride = stride |
| |
|
| | def forward(self, x): |
| | residual = x |
| |
|
| | out = self.conv1(x) |
| | out = self.bn1(out) |
| | out = self.relu(out) |
| |
|
| | out = self.conv2(out) |
| | out = self.bn2(out) |
| |
|
| | if self.downsample is not None: |
| | residual = self.downsample(x) |
| |
|
| | out += residual |
| | out = self.relu(out) |
| |
|
| | return out |
| |
|
| | class Bottleneck(nn.Module): |
| | expansion = 4 |
| | def __init__(self, in_planes, planes, stride=1, dilation=1, downsample=None): |
| | super(Bottleneck, self).__init__() |
| | self.conv1 = nn.Conv2d(in_planes, planes, 1, bias=False) |
| | self.bn1 = nn.BatchNorm2d(planes) |
| | self.conv2 = nn.Conv2d(planes, planes, 3, stride=stride, padding=dilation, |
| | dilation=dilation, bias=False) |
| | self.bn2 = nn.BatchNorm2d(planes) |
| | self.conv3 = nn.Conv2d(planes, planes*self.expansion, 1, bias=False) |
| | self.bn3 = nn.BatchNorm2d(planes*self.expansion) |
| |
|
| | self.relu = nn.ReLU(inplace=False) |
| | self.relu_inplace = nn.ReLU(inplace=True) |
| |
|
| | self.downsample = downsample |
| | self.dilation = dilation |
| | self.stride = stride |
| |
|
| | def forward(self, x): |
| | residual = x |
| |
|
| | out = self.relu(self.bn1(self.conv1(x))) |
| | out = self.relu(self.bn2(self.conv2(out))) |
| | out = self.bn3(self.conv3(out)) |
| |
|
| | if self.downsample is not None: |
| | residual = self.downsample(x) |
| | |
| | out = out + residual |
| | out = self.relu_inplace(out) |
| |
|
| | return out |
| |
|
| | class Resnet(nn.Module): |
| | def __init__(self, block, layers, out_stride=8, use_stem=False, stem_channels=64, in_channels=3): |
| | self.inplanes = 64 |
| | super(Resnet, self).__init__() |
| | outstride_to_strides_and_dilations = { |
| | 8: ((1, 2, 1, 1), (1, 1, 2, 4)), |
| | 16: ((1, 2, 2, 1), (1, 1, 1, 2)), |
| | 32: ((1, 2, 2, 2), (1, 1, 1, 1)), |
| | } |
| | stride_list, dilation_list = outstride_to_strides_and_dilations[out_stride] |
| |
|
| | self.use_stem = use_stem |
| | if use_stem: |
| | self.stem = nn.Sequential( |
| | conv3x3(in_channels, stem_channels//2, stride=2), |
| | nn.BatchNorm2d(stem_channels//2), |
| | nn.ReLU(inplace=False), |
| |
|
| | conv3x3(stem_channels//2, stem_channels//2), |
| | nn.BatchNorm2d(stem_channels//2), |
| | nn.ReLU(inplace=False), |
| |
|
| | conv3x3(stem_channels//2, stem_channels), |
| | nn.BatchNorm2d(stem_channels), |
| | nn.ReLU(inplace=False) |
| | ) |
| | else: |
| | self.conv1 = nn.Conv2d(in_channels, stem_channels, kernel_size=7, stride=2, padding=3, bias=False) |
| | self.bn1 = nn.BatchNorm2d(stem_channels) |
| | self.relu = nn.ReLU(inplace=False) |
| | |
| | |
| | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
| |
|
| | self.layer1 = self._make_layer(block, 64, blocks=layers[0], stride=stride_list[0], dilation=dilation_list[0]) |
| | self.layer2 = self._make_layer(block, 128, blocks=layers[1], stride=stride_list[1], dilation=dilation_list[1]) |
| | self.layer3 = self._make_layer(block, 256, blocks=layers[2], stride=stride_list[2], dilation=dilation_list[2]) |
| | self.layer4 = self._make_layer(block, 512, blocks=layers[3], stride=stride_list[3], dilation=dilation_list[3]) |
| |
|
| | def _make_layer(self, block, planes, blocks, stride=1, dilation=1, contract_dilation=True): |
| | downsample = None |
| | dilations = [dilation] * blocks |
| |
|
| | if contract_dilation and dilation > 1: dilations[0] = dilation // 2 |
| | if stride != 1 or self.inplanes != planes * block.expansion: |
| | downsample = nn.Sequential( |
| | nn.Conv2d(self.inplanes, planes*block.expansion, kernel_size=1, stride=stride, bias=False), |
| | nn.BatchNorm2d(planes*block.expansion) |
| | ) |
| |
|
| | layers = [] |
| | layers.append(block(self.inplanes, planes, stride, dilation=dilations[0], downsample=downsample)) |
| | self.inplanes = planes * block.expansion |
| |
|
| | for i in range(1, blocks): |
| | layers.append(block(self.inplanes, planes, dilation=dilations[i])) |
| | return nn.Sequential(*layers) |
| | |
| | def forward(self, x): |
| | if self.use_stem: |
| | x = self.stem(x) |
| | else: |
| | x = self.relu(self.bn1(self.conv1(x))) |
| |
|
| | x = self.maxpool(x) |
| |
|
| | x1 = self.layer1(x) |
| | x2 = self.layer2(x1) |
| | x3 = self.layer3(x2) |
| | x4 = self.layer4(x3) |
| | |
| | outs = [x1, x2, x3, x4] |
| |
|
| | return tuple(outs) |
| |
|
| | def get_resnet18(pretrained=True): |
| | model = Resnet(BasicBlock, [2, 2, 2, 2], out_stride=32) |
| | if pretrained: |
| | checkpoint = model_zoo.load_url(model_urls['resnet18']) |
| | if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] |
| | else: state_dict = checkpoint |
| | model.load_state_dict(state_dict, strict=False) |
| | return model |
| |
|
| | def get_resnet50_OS8(pretrained=True): |
| | model = Resnet(Bottleneck, [3, 4, 6, 3], out_stride=8, use_stem=True) |
| | if pretrained: |
| | checkpoint = model_zoo.load_url(model_urls['resnet50stem']) |
| | if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] |
| | else: state_dict = checkpoint |
| | model.load_state_dict(state_dict, strict=False) |
| | return model |
| |
|
| | def get_resnet50_OS32(pretrained=True): |
| | model = Resnet(Bottleneck, [3, 4, 6, 3], out_stride=32, use_stem=False) |
| | if pretrained: |
| | checkpoint = model_zoo.load_url(model_urls['resnet50']) |
| | if 'state_dict' in checkpoint: state_dict = checkpoint['state_dict'] |
| | else: state_dict = checkpoint |
| | model.load_state_dict(state_dict, strict=False) |
| | return model |
| |
|
| | if __name__ == "__main__": |
| | model = get_resnet50_OS32() |
| | x = torch.randn(4, 3, 256, 256) |
| | x = model(x)[-1] |
| | print(x.shape) |
