naver/SuperFeatures

1# Copyright (C) 2021-2022 Naver Corporation. All rights reserved.
2# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
3
4import os
5import torch
6from torch import nn
7import torchvision
8
9from how import layers
10
11from lit import LocalfeatureIntegrationTransformer
12
13from how.networks.how_net import HOWNet
14
15class FIReNet(HOWNet):
16
17    def __init__(self, features, attention, lit, dim_reduction, meta, runtime):
18        super().__init__(features, attention, None, dim_reduction, meta, runtime)
19        self.lit = lit
20        self.return_global = False
21        
22    def copy_excluding_dim_reduction(self):
23        """Return a copy of this network without the dim_reduction layer"""
24        meta = {**self.meta, "outputdim": self.meta['backbone_dim']}
25        return self.__class__(self.features, self.attention, self.lit, None, meta, self.runtime)
26
27    def copy_with_runtime(self, runtime):
28        """Return a copy of this network with a different runtime dict"""
29        return self.__class__(self.features, self.attention, self.lit, self.dim_reduction, self.meta, runtime)
30
31    def parameter_groups(self):
32        """Return torch parameter groups"""
33        layers = [self.features, self.attention, self.smoothing, self.lit]
34        parameters = [{'params': x.parameters()} for x in layers if x is not None]
35        if self.dim_reduction:
36            # Do not update dimensionality reduction layer
37            parameters.append({'params': self.dim_reduction.parameters(), 'lr': 0.0})
38        return parameters
39
40    def get_superfeatures(self, x, *, scales):
41        """
42        return a list of tuple (features, attentionmpas) where each is a list containing requested scales
43        features is a tensor BxDxNx1
44        attentionmaps is a tensor BxNxHxW
45        """
46        feats = []
47        attns = []
48        strengths = []
49        for s in scales:
50            xs = nn.functional.interpolate(x, scale_factor=s, mode='bilinear', align_corners=False)
51            o = self.features(xs)
52            o, attn = self.lit(o)
53            strength = self.attention(o)
54            if self.smoothing:
55                o = self.smoothing(o)
56            if self.dim_reduction:
57                o = self.dim_reduction(o)
58            feats.append(o)
59            attns.append(attn)
60            strengths.append(strength)
61        return feats, attns, strengths
62        
63    def forward(self, x):
64        return self.get_superfeatures(x, scales=self.runtime['training_scales'])
65        
66    
67def init_network(architecture, pretrained, skip_layer, dim_reduction, lit, runtime):
68    """Initialize FIRe network
69    :param str architecture: Network backbone architecture (e.g. resnet18)
70    :param str pretrained: url of the pretrained model (None for using random initialization)
71    :param int skip_layer: How many layers of blocks should be skipped (from the end)
72    :param dict dim_reduction: Options for the dimensionality reduction layer
73    :param dict lit: Options for the lit layer
74    :param dict runtime: Runtime options to be stored in the network
75    :return FIRe: Initialized network
76    """
77    # Take convolutional layers as features, always ends with ReLU to make last activations non-negative
78    net_in = getattr(torchvision.models, architecture)(pretrained=False) # use trained weights including the LIT module instead 
79    if architecture.startswith('alexnet') or architecture.startswith('vgg'):
80        features = list(net_in.features.children())[:-1]
81    elif architecture.startswith('resnet'):
82        features = list(net_in.children())[:-2]
83    elif architecture.startswith('densenet'):
84        features = list(net_in.features.children()) + [nn.ReLU(inplace=True)]
85    elif architecture.startswith('squeezenet'):
86        features = list(net_in.features.children())
87    else:
88        raise ValueError('Unsupported or unknown architecture: {}!'.format(architecture))
89
90    if skip_layer > 0:
91        features = features[:-skip_layer]
92    backbone_dim = 2048 // (2 ** skip_layer)
93
94    att_layer = layers.attention.L2Attention()
95
96    lit_layer = LocalfeatureIntegrationTransformer(**lit, input_dim=backbone_dim)
97
98    reduction_layer = None
99    if dim_reduction:
100        reduction_layer = layers.dim_reduction.ConvDimReduction(**dim_reduction, input_dim=lit['dim'])
101
102    meta = {
103        "architecture": architecture,
104        "backbone_dim": lit['dim'],
105        "outputdim": reduction_layer.out_channels if dim_reduction else lit['dim'],
106        "corercf_size": 32 // (2 ** skip_layer),
107    }
108    net = FIReNet(nn.Sequential(*features), att_layer, lit_layer, reduction_layer, meta, runtime)
109    
110    if pretrained is not None:
111        assert os.path.isfile(pretrained), pretrained
112        ckpt = torch.load(pretrained, map_location='cpu')
113        missing, unexpected = net.load_state_dict(ckpt['state_dict'], strict=False)
114        assert all(['dim_reduction' in a for a in missing]), "Loading did not go well"
115        assert all(['fc' in a for a in unexpected]), "Loading did not go well"
116    return net
117