A slight accelerated version of DFNet, we experimentally found this version's performance is similar to original DFNet but inferences faster
| 172 | return feature_maps, predict |
| 173 | |
| 174 | class DFNet_s(nn.Module): |
| 175 | ''' A slight accelerated version of DFNet, we experimentally found this version's performance is similar to original DFNet but inferences faster ''' |
| 176 | default_conf = { |
| 177 | 'hypercolumn_layers': ["conv1_2"], |
| 178 | 'output_dim': 128, |
| 179 | } |
| 180 | mean = [0.485, 0.456, 0.406] |
| 181 | std = [0.229, 0.224, 0.225] |
| 182 | |
| 183 | def __init__(self, feat_dim=12, places365_model_path=''): |
| 184 | super(DFNet_s, self).__init__() |
| 185 | |
| 186 | self.layer_to_index = {k: v for v, k in enumerate(vgg16_layers.keys())} |
| 187 | self.hypercolumn_indices = [self.layer_to_index[n] for n in self.default_conf['hypercolumn_layers']] # [2, 14, 28] |
| 188 | |
| 189 | # Initialize architecture |
| 190 | vgg16 = models.vgg16(pretrained=True) |
| 191 | |
| 192 | self.encoder = nn.Sequential(*list(vgg16.features.children())) |
| 193 | |
| 194 | self.scales = [] |
| 195 | current_scale = 0 |
| 196 | for i, layer in enumerate(self.encoder): |
| 197 | if isinstance(layer, torch.nn.MaxPool2d): |
| 198 | current_scale += 1 |
| 199 | if i in self.hypercolumn_indices: |
| 200 | self.scales.append(2**current_scale) |
| 201 | |
| 202 | ## adaptation layers, see off branches from fig.3 in S2DNet paper |
| 203 | self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim']) |
| 204 | |
| 205 | # pose regression layers |
| 206 | self.avgpool = nn.AdaptiveAvgPool2d(1) |
| 207 | self.fc_pose = nn.Linear(512, feat_dim) |
| 208 | |
| 209 | def forward(self, x, return_feature=False, isSingleStream=False, return_pose=True, upsampleH=240, upsampleW=427): |
| 210 | ''' |
| 211 | inference DFNet_s. It can regress camera pose as well as extract intermediate layer features. |
| 212 | :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream |
| 213 | :param return_feature: whether to return features as output |
| 214 | :param isSingleStream: whether it's an single stream inference or siamese network inference |
| 215 | :param upsampleH: feature upsample size H |
| 216 | :param upsampleW: feature upsample size W |
| 217 | :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None |
| 218 | :return predict: [2B, 12] or [B, 12] |
| 219 | ''' |
| 220 | |
| 221 | # normalize input data |
| 222 | mean, std = x.new_tensor(self.mean), x.new_tensor(self.std) |
| 223 | x = (x - mean[:, None, None]) / std[:, None, None] |
| 224 | |
| 225 | ### encoder ### |
| 226 | feature_maps = [] |
| 227 | for i in range(len(self.encoder)): |
| 228 | x = self.encoder[i](x) |
| 229 | |
| 230 | if i in self.hypercolumn_indices: |
| 231 | feature = x.clone() |