| 109 | |
| 110 | |
| 111 | class ResnetDilated(nn.Module): |
| 112 | def __init__(self, orig_resnet, dilate_scale=8): |
| 113 | super(ResnetDilated, self).__init__() |
| 114 | from functools import partial |
| 115 | |
| 116 | if dilate_scale == 8: |
| 117 | orig_resnet.layer3.apply( |
| 118 | partial(self._nostride_dilate, dilate=2)) |
| 119 | orig_resnet.layer4.apply( |
| 120 | partial(self._nostride_dilate, dilate=4)) |
| 121 | elif dilate_scale == 16: |
| 122 | orig_resnet.layer4.apply( |
| 123 | partial(self._nostride_dilate, dilate=2)) |
| 124 | |
| 125 | # take pretrained resnet, except AvgPool and FC |
| 126 | self.conv1 = orig_resnet.conv1 |
| 127 | self.bn1 = orig_resnet.bn1 |
| 128 | self.relu1 = orig_resnet.relu1 |
| 129 | self.conv2 = orig_resnet.conv2 |
| 130 | self.bn2 = orig_resnet.bn2 |
| 131 | self.relu2 = orig_resnet.relu2 |
| 132 | self.conv3 = orig_resnet.conv3 |
| 133 | self.bn3 = orig_resnet.bn3 |
| 134 | self.relu3 = orig_resnet.relu3 |
| 135 | self.maxpool = orig_resnet.maxpool |
| 136 | self.layer1 = orig_resnet.layer1 |
| 137 | self.layer2 = orig_resnet.layer2 |
| 138 | self.layer3 = orig_resnet.layer3 |
| 139 | self.layer4 = orig_resnet.layer4 |
| 140 | |
| 141 | def _nostride_dilate(self, m, dilate): |
| 142 | classname = m.__class__.__name__ |
| 143 | if classname.find('Conv') != -1: |
| 144 | # the convolution with stride |
| 145 | if m.stride == (2, 2): |
| 146 | m.stride = (1, 1) |
| 147 | if m.kernel_size == (3, 3): |
| 148 | m.dilation = (dilate//2, dilate//2) |
| 149 | m.padding = (dilate//2, dilate//2) |
| 150 | # other convoluions |
| 151 | else: |
| 152 | if m.kernel_size == (3, 3): |
| 153 | m.dilation = (dilate, dilate) |
| 154 | m.padding = (dilate, dilate) |
| 155 | |
| 156 | def forward(self, x, return_feature_maps=False): |
| 157 | conv_out = [] |
| 158 | |
| 159 | x = self.relu1(self.bn1(self.conv1(x))) |
| 160 | x = self.relu2(self.bn2(self.conv2(x))) |
| 161 | x = self.relu3(self.bn3(self.conv3(x))) |
| 162 | x = self.maxpool(x) |
| 163 | |
| 164 | x = self.layer1(x); conv_out.append(x); |
| 165 | x = self.layer2(x); conv_out.append(x); |
| 166 | x = self.layer3(x); conv_out.append(x); |
| 167 | x = self.layer4(x); conv_out.append(x); |
| 168 | |