third_party/SOLD2/sold2/model/nets/descriptor_decoder.py

import torch
import torch.nn as nn


class SuperpointDescriptor(nn.Module):
    """ Descriptor decoder based on the SuperPoint arcihtecture. """
    def __init__(self, input_feat_dim=128):
        super(SuperpointDescriptor, self).__init__()
        self.relu = torch.nn.ReLU(inplace=True)
        self.convPa = torch.nn.Conv2d(input_feat_dim, 256, kernel_size=3,
                                      stride=1, padding=1)        
        self.convPb = torch.nn.Conv2d(256, 128, kernel_size=1,
                                      stride=1, padding=0)

    def forward(self, input_features):
        feat = self.relu(self.convPa(input_features))
        semi = self.convPb(feat)

        return semi