Class PoseNet_res34

script/dm/pose_model.py:229–260

# PoseNet (SE(3)) w/ resnet34 backnone. We found dropout layer is unnecessary, so we set droprate as 0 in reported results.
class PoseNet_res34(nn.Module):
    def __init__(self, droprate=0.5, pretrained=True,
        feat_dim=2048):
        super(PoseNet_res34, self).__init__()
        self.droprate = droprate

        # replace the last FC layer in feature extractor
        self.feature_extractor = models.resnet34(pretrained=True)
        self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1)
        fe_out_planes = self.feature_extractor.fc.in_features
        self.feature_extractor.fc = nn.Linear(fe_out_planes, feat_dim)
        self.fc_pose = nn.Linear(feat_dim, 12)

        # initialize
        if pretrained:
          init_modules = [self.feature_extractor.fc]
        else:
          init_modules = self.modules()

        for m in init_modules:
          if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
            nn.init.kaiming_normal_(m.weight.data)
            if m.bias is not None:
              nn.init.constant_(m.bias.data, 0)

    def forward(self, x):
        x = self.feature_extractor(x)
        x = F.relu(x)
        if self.droprate > 0:
          x = F.dropout(x, p=self.droprate)
        predict = self.fc_pose(x)
        return predict


# from MapNet paper CVPR 2018