pytorch-image-models

1
from typing import Optional
2

3
import torch
4
from torch import nn
5
from torch import nn, Tensor
6
from torch.nn.modules.transformer import _get_activation_fn
7

8

9
def add_ml_decoder_head(model):
10
    if hasattr(model, 'global_pool') and hasattr(model, 'fc'):  # most CNN models, like Resnet50
11
        model.global_pool = nn.Identity()
12
        del model.fc
13
        num_classes = model.num_classes
14
        num_features = model.num_features
15
        model.fc = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
16
    elif hasattr(model, 'global_pool') and hasattr(model, 'classifier'):  # EfficientNet
17
        model.global_pool = nn.Identity()
18
        del model.classifier
19
        num_classes = model.num_classes
20
        num_features = model.num_features
21
        model.classifier = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
22
    elif 'RegNet' in model._get_name() or 'TResNet' in model._get_name():  # hasattr(model, 'head')
23
        del model.head
24
        num_classes = model.num_classes
25
        num_features = model.num_features
26
        model.head = MLDecoder(num_classes=num_classes, initial_num_features=num_features)
27
    else:
28
        print("Model code-writing is not aligned currently with ml-decoder")
29
        exit(-1)
30
    if hasattr(model, 'drop_rate'):  # Ml-Decoder has inner dropout
31
        model.drop_rate = 0
32
    return model
33

34

35
class TransformerDecoderLayerOptimal(nn.Module):
36
    def __init__(self, d_model, nhead=8, dim_feedforward=2048, dropout=0.1, activation="relu",
37
                 layer_norm_eps=1e-5) -> None:
38
        super(TransformerDecoderLayerOptimal, self).__init__()
39
        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
40
        self.dropout = nn.Dropout(dropout)
41
        self.dropout1 = nn.Dropout(dropout)
42
        self.dropout2 = nn.Dropout(dropout)
43
        self.dropout3 = nn.Dropout(dropout)
44

45
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
46

47
        # Implementation of Feedforward model
48
        self.linear1 = nn.Linear(d_model, dim_feedforward)
49
        self.linear2 = nn.Linear(dim_feedforward, d_model)
50

51
        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
52
        self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps)
53

54
        self.activation = _get_activation_fn(activation)
55

56
    def __setstate__(self, state):
57
        if 'activation' not in state:
58
            state['activation'] = torch.nn.functional.relu
59
        super(TransformerDecoderLayerOptimal, self).__setstate__(state)
60

61
    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
62
                memory_mask: Optional[Tensor] = None,
63
                tgt_key_padding_mask: Optional[Tensor] = None,
64
                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:
65
        tgt = tgt + self.dropout1(tgt)
66
        tgt = self.norm1(tgt)
67
        tgt2 = self.multihead_attn(tgt, memory, memory)[0]
68
        tgt = tgt + self.dropout2(tgt2)
69
        tgt = self.norm2(tgt)
70
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
71
        tgt = tgt + self.dropout3(tgt2)
72
        tgt = self.norm3(tgt)
73
        return tgt
74

75

76
# @torch.jit.script
77
# class ExtrapClasses(object):
78
#     def __init__(self, num_queries: int, group_size: int):
79
#         self.num_queries = num_queries
80
#         self.group_size = group_size
81
#
82
#     def __call__(self, h: torch.Tensor, class_embed_w: torch.Tensor, class_embed_b: torch.Tensor, out_extrap:
83
#     torch.Tensor):
84
#         # h = h.unsqueeze(-1).expand(-1, -1, -1, self.group_size)
85
#         h = h[..., None].repeat(1, 1, 1, self.group_size) # torch.Size([bs, 5, 768, groups])
86
#         w = class_embed_w.view((self.num_queries, h.shape[2], self.group_size))
87
#         out = (h * w).sum(dim=2) + class_embed_b
88
#         out = out.view((h.shape[0], self.group_size * self.num_queries))
89
#         return out
90

91
@torch.jit.script
92
class GroupFC(object):
93
    def __init__(self, embed_len_decoder: int):
94
        self.embed_len_decoder = embed_len_decoder
95

96
    def __call__(self, h: torch.Tensor, duplicate_pooling: torch.Tensor, out_extrap: torch.Tensor):
97
        for i in range(self.embed_len_decoder):
98
            h_i = h[:, i, :]
99
            w_i = duplicate_pooling[i, :, :]
100
            out_extrap[:, i, :] = torch.matmul(h_i, w_i)
101

102

103
class MLDecoder(nn.Module):
104
    def __init__(self, num_classes, num_of_groups=-1, decoder_embedding=768, initial_num_features=2048):
105
        super(MLDecoder, self).__init__()
106
        embed_len_decoder = 100 if num_of_groups < 0 else num_of_groups
107
        if embed_len_decoder > num_classes:
108
            embed_len_decoder = num_classes
109

110
        # switching to 768 initial embeddings
111
        decoder_embedding = 768 if decoder_embedding < 0 else decoder_embedding
112
        self.embed_standart = nn.Linear(initial_num_features, decoder_embedding)
113

114
        # decoder
115
        decoder_dropout = 0.1
116
        num_layers_decoder = 1
117
        dim_feedforward = 2048
118
        layer_decode = TransformerDecoderLayerOptimal(d_model=decoder_embedding,
119
                                                      dim_feedforward=dim_feedforward, dropout=decoder_dropout)
120
        self.decoder = nn.TransformerDecoder(layer_decode, num_layers=num_layers_decoder)
121

122
        # non-learnable queries
123
        self.query_embed = nn.Embedding(embed_len_decoder, decoder_embedding)
124
        self.query_embed.requires_grad_(False)
125

126
        # group fully-connected
127
        self.num_classes = num_classes
128
        self.duplicate_factor = int(num_classes / embed_len_decoder + 0.999)
129
        self.duplicate_pooling = torch.nn.Parameter(
130
            torch.Tensor(embed_len_decoder, decoder_embedding, self.duplicate_factor))
131
        self.duplicate_pooling_bias = torch.nn.Parameter(torch.Tensor(num_classes))
132
        torch.nn.init.xavier_normal_(self.duplicate_pooling)
133
        torch.nn.init.constant_(self.duplicate_pooling_bias, 0)
134
        self.group_fc = GroupFC(embed_len_decoder)
135

136
    def forward(self, x):
137
        if len(x.shape) == 4:  # [bs,2048, 7,7]
138
            embedding_spatial = x.flatten(2).transpose(1, 2)
139
        else:  # [bs, 197,468]
140
            embedding_spatial = x
141
        embedding_spatial_786 = self.embed_standart(embedding_spatial)
142
        embedding_spatial_786 = torch.nn.functional.relu(embedding_spatial_786, inplace=True)
143

144
        bs = embedding_spatial_786.shape[0]
145
        query_embed = self.query_embed.weight
146
        # tgt = query_embed.unsqueeze(1).repeat(1, bs, 1)
147
        tgt = query_embed.unsqueeze(1).expand(-1, bs, -1)  # no allocation of memory with expand
148
        h = self.decoder(tgt, embedding_spatial_786.transpose(0, 1))  # [embed_len_decoder, batch, 768]
149
        h = h.transpose(0, 1)
150

151
        out_extrap = torch.zeros(h.shape[0], h.shape[1], self.duplicate_factor, device=h.device, dtype=h.dtype)
152
        self.group_fc(h, self.duplicate_pooling, out_extrap)
153
        h_out = out_extrap.flatten(1)[:, :self.num_classes]
154
        h_out += self.duplicate_pooling_bias
155
        logits = h_out
156
        return logits
157