| 173 | |
| 174 | # pyramid pooling, deep supervision |
| 175 | class PPM(nn.Module): |
| 176 | def __init__(self, num_class=150, fc_dim=4096, |
| 177 | use_softmax=False, pool_scales=(1, 2, 3, 6)): |
| 178 | super(PPM, self).__init__() |
| 179 | self.use_softmax = use_softmax |
| 180 | |
| 181 | self.ppm = [] |
| 182 | for scale in pool_scales: |
| 183 | self.ppm.append(nn.Sequential( |
| 184 | nn.AdaptiveAvgPool2d(scale), |
| 185 | nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False), |
| 186 | BatchNorm2d(512), |
| 187 | nn.ReLU(inplace=True) |
| 188 | )) |
| 189 | self.ppm = nn.ModuleList(self.ppm) |
| 190 | self.conv_last = nn.Sequential( |
| 191 | nn.Conv2d(fc_dim+len(pool_scales)*512, 512, |
| 192 | kernel_size=3, padding=1, bias=False), |
| 193 | BatchNorm2d(512), |
| 194 | nn.ReLU(inplace=True), |
| 195 | nn.Dropout2d(0.1), |
| 196 | nn.Conv2d(512, num_class, kernel_size=1) |
| 197 | ) |
| 198 | |
| 199 | def forward(self, conv_out, segSize=None): |
| 200 | conv4 = conv_out[-2] |
| 201 | conv5 = conv_out[-1] |
| 202 | ###### |
| 203 | x1 = nn.functional.interpolate(conv4, size=segSize, mode='bilinear', align_corners=False) |
| 204 | ###### |
| 205 | input_size = conv5.size() |
| 206 | ppm_out = [conv5] |
| 207 | for pool_scale in self.ppm: |
| 208 | ppm_out.append(nn.functional.interpolate( |
| 209 | pool_scale(conv5), |
| 210 | (input_size[2], input_size[3]), |
| 211 | mode='bilinear', align_corners=False)) |
| 212 | ppm_out = torch.cat(ppm_out, 1) |
| 213 | x = self.conv_last(ppm_out) |
| 214 | x = nn.functional.interpolate( |
| 215 | x, size=segSize, mode='bilinear', align_corners=False) |
| 216 | x = nn.functional.softmax(x, dim=1) |
| 217 | return x, x1 |
| 218 | |
| 219 | |