|
| 1 | +import torch.nn as nn |
| 2 | +import torch.nn.functional as F |
| 3 | + |
| 4 | + |
| 5 | +def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): |
| 6 | + """3x3 convolution with padding""" |
| 7 | + return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, |
| 8 | + padding=dilation, groups=groups, bias=False, dilation=dilation) |
| 9 | + |
| 10 | + |
| 11 | +class BasicBlock(nn.Module): |
| 12 | + expansion = 1 |
| 13 | + |
| 14 | + def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, |
| 15 | + base_width=64, dilation=1, norm_layer=None, dcn=None): |
| 16 | + super(BasicBlock, self).__init__() |
| 17 | + if norm_layer is None: |
| 18 | + norm_layer = nn.BatchNorm2d |
| 19 | + if groups != 1 or base_width != 64: |
| 20 | + raise ValueError('BasicBlock only supports groups=1 and base_width=64') |
| 21 | + if dilation > 1: |
| 22 | + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") |
| 23 | + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 |
| 24 | + self.conv1 = conv3x3(inplanes, planes, stride) |
| 25 | + self.bn1 = norm_layer(planes) |
| 26 | + self.relu = nn.ReLU(inplace=True) |
| 27 | + self.conv2 = conv3x3(planes, planes) |
| 28 | + self.bn2 = norm_layer(planes) |
| 29 | + self.downsample = downsample |
| 30 | + self.stride = stride |
| 31 | + |
| 32 | + def forward(self, x): |
| 33 | + identity = x |
| 34 | + |
| 35 | + out = self.conv1(x) |
| 36 | + out = self.bn1(out) |
| 37 | + out = self.relu(out) |
| 38 | + |
| 39 | + out = self.conv2(out) |
| 40 | + out = self.bn2(out) |
| 41 | + |
| 42 | + if self.downsample is not None: |
| 43 | + identity = self.downsample(x) |
| 44 | + |
| 45 | + out += identity |
| 46 | + out = self.relu(out) |
| 47 | + |
| 48 | + return out |
| 49 | + |
| 50 | + |
| 51 | +class Bottleneck(nn.Module): |
| 52 | + expansion = 4 |
| 53 | + |
| 54 | + def __init__(self, inplanes, planes, stride=1, |
| 55 | + downsample=None, norm_layer=nn.BatchNorm2d, dcn=None): |
| 56 | + super(Bottleneck, self).__init__() |
| 57 | + self.dcn = dcn |
| 58 | + self.with_dcn = dcn is not None |
| 59 | + |
| 60 | + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) |
| 61 | + self.bn1 = norm_layer(planes, momentum=0.1) |
| 62 | + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, |
| 63 | + padding=1, bias=False) |
| 64 | + |
| 65 | + self.bn2 = norm_layer(planes, momentum=0.1) |
| 66 | + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) |
| 67 | + self.bn3 = norm_layer(planes * 4, momentum=0.1) |
| 68 | + self.downsample = downsample |
| 69 | + self.stride = stride |
| 70 | + |
| 71 | + def forward(self, x): |
| 72 | + residual = x |
| 73 | + |
| 74 | + out = F.relu(self.bn1(self.conv1(x)), inplace=True) |
| 75 | + if not self.with_dcn: |
| 76 | + out = F.relu(self.bn2(self.conv2(out)), inplace=True) |
| 77 | + elif self.with_modulated_dcn: |
| 78 | + offset_mask = self.conv2_offset(out) |
| 79 | + offset = offset_mask[:, :18 * self.deformable_groups, :, :] |
| 80 | + mask = offset_mask[:, -9 * self.deformable_groups:, :, :] |
| 81 | + mask = mask.sigmoid() |
| 82 | + out = F.relu(self.bn2(self.conv2(out, offset, mask))) |
| 83 | + else: |
| 84 | + offset = self.conv2_offset(out) |
| 85 | + out = F.relu(self.bn2(self.conv2(out, offset)), inplace=True) |
| 86 | + |
| 87 | + out = self.conv3(out) |
| 88 | + out = self.bn3(out) |
| 89 | + |
| 90 | + if self.downsample is not None: |
| 91 | + residual = self.downsample(x) |
| 92 | + |
| 93 | + out += residual |
| 94 | + out = F.relu(out) |
| 95 | + |
| 96 | + return out |
| 97 | + |
| 98 | + |
| 99 | +class ResNet(nn.Module): |
| 100 | + """ ResNet """ |
| 101 | + |
| 102 | + def __init__(self, architecture, norm_layer=nn.BatchNorm2d, dcn=None, stage_with_dcn=(False, False, False, False)): |
| 103 | + super(ResNet, self).__init__() |
| 104 | + self._norm_layer = norm_layer |
| 105 | + assert architecture in ["resnet18", "resnet34", "resnet50", "resnet101", 'resnet152'] |
| 106 | + layers = { |
| 107 | + 'resnet18': [2, 2, 2, 2], |
| 108 | + 'resnet34': [3, 4, 6, 3], |
| 109 | + 'resnet50': [3, 4, 6, 3], |
| 110 | + 'resnet101': [3, 4, 23, 3], |
| 111 | + 'resnet152': [3, 8, 36, 3], |
| 112 | + } |
| 113 | + self.inplanes = 64 |
| 114 | + if architecture == "resnet18" or architecture == 'resnet34': |
| 115 | + self.block = BasicBlock |
| 116 | + else: |
| 117 | + self.block = Bottleneck |
| 118 | + self.layers = layers[architecture] |
| 119 | + |
| 120 | + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, |
| 121 | + stride=2, padding=3, bias=False) |
| 122 | + self.bn1 = norm_layer(64, eps=1e-5, momentum=0.1, affine=True) |
| 123 | + self.relu = nn.ReLU(inplace=True) |
| 124 | + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
| 125 | + |
| 126 | + stage_dcn = [dcn if with_dcn else None for with_dcn in stage_with_dcn] |
| 127 | + |
| 128 | + self.layer1 = self.make_layer( |
| 129 | + self.block, 64, self.layers[0], dcn=stage_dcn[0]) |
| 130 | + self.layer2 = self.make_layer( |
| 131 | + self.block, 128, self.layers[1], stride=2, dcn=stage_dcn[1]) |
| 132 | + self.layer3 = self.make_layer( |
| 133 | + self.block, 256, self.layers[2], stride=2, dcn=stage_dcn[2]) |
| 134 | + |
| 135 | + self.layer4 = self.make_layer( |
| 136 | + self.block, 512, self.layers[3], stride=2, dcn=stage_dcn[3]) |
| 137 | + |
| 138 | + def forward(self, x): |
| 139 | + x = self.maxpool(self.relu(self.bn1(self.conv1(x)))) # 64 * h/4 * w/4 |
| 140 | + x = self.layer1(x) # 256 * h/4 * w/4 |
| 141 | + x = self.layer2(x) # 512 * h/8 * w/8 |
| 142 | + x = self.layer3(x) # 1024 * h/16 * w/16 |
| 143 | + x = self.layer4(x) # 2048 * h/32 * w/32 |
| 144 | + return x |
| 145 | + |
| 146 | + def stages(self): |
| 147 | + return [self.layer1, self.layer2, self.layer3, self.layer4] |
| 148 | + |
| 149 | + def make_layer(self, block, planes, blocks, stride=1, dcn=None): |
| 150 | + downsample = None |
| 151 | + if stride != 1 or self.inplanes != planes * block.expansion: |
| 152 | + downsample = nn.Sequential( |
| 153 | + nn.Conv2d(self.inplanes, planes * block.expansion, |
| 154 | + kernel_size=1, stride=stride, bias=False), |
| 155 | + self._norm_layer(planes * block.expansion), |
| 156 | + ) |
| 157 | + |
| 158 | + layers = [] |
| 159 | + layers.append(block(self.inplanes, planes, stride, downsample, |
| 160 | + norm_layer=self._norm_layer, dcn=dcn)) |
| 161 | + self.inplanes = planes * block.expansion |
| 162 | + for i in range(1, blocks): |
| 163 | + layers.append(block(self.inplanes, planes, |
| 164 | + norm_layer=self._norm_layer, dcn=dcn)) |
| 165 | + |
| 166 | + return nn.Sequential(*layers) |
0 commit comments