| 263 | # from MapNet paper CVPR 2018 |
| 264 | class PoseNet(nn.Module): |
| 265 | def __init__(self, feature_extractor, droprate=0.5, pretrained=True, |
| 266 | feat_dim=2048, filter_nans=False): |
| 267 | super(PoseNet, self).__init__() |
| 268 | self.droprate = droprate |
| 269 | |
| 270 | # replace the last FC layer in feature extractor |
| 271 | self.feature_extractor = models.resnet34(pretrained=True) |
| 272 | self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1) |
| 273 | fe_out_planes = self.feature_extractor.fc.in_features |
| 274 | self.feature_extractor.fc = nn.Linear(fe_out_planes, feat_dim) |
| 275 | |
| 276 | self.fc_xyz = nn.Linear(feat_dim, 3) |
| 277 | self.fc_wpqr = nn.Linear(feat_dim, 3) |
| 278 | if filter_nans: |
| 279 | self.fc_wpqr.register_backward_hook(hook=filter_hook) |
| 280 | # initialize |
| 281 | if pretrained: |
| 282 | init_modules = [self.feature_extractor.fc, self.fc_xyz, self.fc_wpqr] |
| 283 | else: |
| 284 | init_modules = self.modules() |
| 285 | |
| 286 | for m in init_modules: |
| 287 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): |
| 288 | nn.init.kaiming_normal_(m.weight.data) |
| 289 | if m.bias is not None: |
| 290 | nn.init.constant_(m.bias.data, 0) |
| 291 | |
| 292 | def forward(self, x): |
| 293 | x = self.feature_extractor(x) |