| 228 | # PoseNet (SE(3)) w/ resnet34 backnone. We found dropout layer is unnecessary, so we set droprate as 0 in reported results. |
| 229 | class PoseNet_res34(nn.Module): |
| 230 | def __init__(self, droprate=0.5, pretrained=True, |
| 231 | feat_dim=2048): |
| 232 | super(PoseNet_res34, self).__init__() |
| 233 | self.droprate = droprate |
| 234 | |
| 235 | # replace the last FC layer in feature extractor |
| 236 | self.feature_extractor = models.resnet34(pretrained=True) |
| 237 | self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1) |
| 238 | fe_out_planes = self.feature_extractor.fc.in_features |
| 239 | self.feature_extractor.fc = nn.Linear(fe_out_planes, feat_dim) |
| 240 | self.fc_pose = nn.Linear(feat_dim, 12) |
| 241 | |
| 242 | # initialize |
| 243 | if pretrained: |
| 244 | init_modules = [self.feature_extractor.fc] |
| 245 | else: |
| 246 | init_modules = self.modules() |
| 247 | |
| 248 | for m in init_modules: |
| 249 | if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): |
| 250 | nn.init.kaiming_normal_(m.weight.data) |
| 251 | if m.bias is not None: |
| 252 | nn.init.constant_(m.bias.data, 0) |
| 253 | |
| 254 | def forward(self, x): |
| 255 | x = self.feature_extractor(x) |