Class Decoder

DeDoDe/decoder.py:6–18

class Decoder(nn.Module):
    def __init__(self, layers, *args, super_resolution = False, num_prototypes = 1, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.layers = layers
        self.scales = self.layers.keys()
        self.super_resolution = super_resolution
        self.num_prototypes = num_prototypes
    def forward(self, features, context = None, scale = None):
        if context is not None:
            features = torch.cat((features, context), dim = 1)
        stuff = self.layers[scale](features)
        logits, context = stuff[:,:self.num_prototypes], stuff[:,self.num_prototypes:]
        return logits, context

class ConvRefiner(nn.Module):
    def __init__(