Use the right version of CCT

Author: federicobrancasi

Tests/Models/RunCCT/CCT/CCT/cct.py

@@ -0,0 +1,5 @@
+def register_model(func):
+    """
+    Fallback wrapper in case timm isn't installed
+    """
+    return func

Tests/Models/RunCCT/CCT/CCT/registry.py

@@ -0,0 +1,39 @@
+import torch.nn as nn
+
+
+class Embedder(nn.Module):
+
+    def __init__(self,
+                 word_embedding_dim = 300,
+                 vocab_size = 100000,
+                 padding_idx = 1,
+                 pretrained_weight = None,
+                 embed_freeze = False,
+                 *args,
+                 **kwargs):
+        super(Embedder, self).__init__()
+        self.embeddings = nn.Embedding.from_pretrained(pretrained_weight, freeze=embed_freeze) \
+            if pretrained_weight is not None else \
+            nn.Embedding(vocab_size, word_embedding_dim, padding_idx=padding_idx)
+        self.embeddings.weight.requires_grad = not embed_freeze
+
+    def forward_mask(self, mask):
+        bsz, seq_len = mask.shape
+        new_mask = mask.view(bsz, seq_len, 1)
+        new_mask = new_mask.sum(-1)
+        new_mask = (new_mask > 0)
+        return new_mask
+
+    def forward(self, x, mask = None):
+        embed = self.embeddings(x)
+        embed = embed if mask is None else embed * self.forward_mask(mask).unsqueeze(-1).float()
+        return embed, mask
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std = .02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        else:
+            nn.init.normal_(m.weight)

Tests/Models/RunCCT/CCT/CCT/utils/__init__.py

@@ -0,0 +1,45 @@
+import logging
+import math
+
+import torch
+import torch.nn.functional as F
+
+_logger = logging.getLogger('train')
+
+
+def resize_pos_embed(posemb, posemb_new, num_tokens = 1):
+    # Copied from `timm` by Ross Wightman:
+    # github.com/rwightman/pytorch-image-models
+    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
+    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
+    ntok_new = posemb_new.shape[1]
+    if num_tokens:
+        posemb_tok, posemb_grid = posemb[:, :num_tokens], posemb[0, num_tokens:]
+        ntok_new -= num_tokens
+    else:
+        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
+    gs_old = int(math.sqrt(len(posemb_grid)))
+    gs_new = int(math.sqrt(ntok_new))
+    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
+    posemb_grid = F.interpolate(posemb_grid, size = (gs_new, gs_new), mode = 'bilinear')
+    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
+    posemb = torch.cat([posemb_tok, posemb_grid], dim = 1)
+    return posemb
+
+
+def pe_check(model, state_dict, pe_key = 'classifier.positional_emb'):
+    if pe_key is not None and pe_key in state_dict.keys() and pe_key in model.state_dict().keys():
+        if model.state_dict()[pe_key].shape != state_dict[pe_key].shape:
+            state_dict[pe_key] = resize_pos_embed(state_dict[pe_key],
+                                                  model.state_dict()[pe_key],
+                                                  num_tokens = model.classifier.num_tokens)
+    return state_dict
+
+
+def fc_check(model, state_dict, fc_key = 'classifier.fc'):
+    for key in [f'{fc_key}.weight', f'{fc_key}.bias']:
+        if key is not None and key in state_dict.keys() and key in model.state_dict().keys():
+            if model.state_dict()[key].shape != state_dict[key].shape:
+                _logger.warning(f'Removing {key}, number of classes has changed.')
+                state_dict[key] = model.state_dict()[key]
+    return state_dict

Tests/Models/RunCCT/CCT/CCT/utils/embedder.py

@@ -0,0 +1,39 @@
+# Thanks to rwightman's timm package
+# github.com:rwightman/pytorch-image-models
+
+import torch
+import torch.nn as nn
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """
+    Obtained from: github.com:rwightman/pytorch-image-models
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype = x.dtype, device = x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """
+    Obtained from: github.com:rwightman/pytorch-image-models
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob = None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)

Tests/Models/RunCCT/CCT/CCT/utils/helpers.py

@@ -0,0 +1,116 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Tokenizer(nn.Module):
+
+    def __init__(self,
+                 kernel_size,
+                 stride,
+                 padding,
+                 pooling_kernel_size = 3,
+                 pooling_stride = 2,
+                 pooling_padding = 1,
+                 n_conv_layers = 1,
+                 n_input_channels = 3,
+                 n_output_channels = 64,
+                 in_planes = 64,
+                 activation = None,
+                 max_pool = True,
+                 conv_bias = False):
+        super(Tokenizer, self).__init__()
+
+        n_filter_list = [n_input_channels] + \
+                        [in_planes for _ in range(n_conv_layers - 1)] + \
+                        [n_output_channels]
+
+        self.conv_layers = nn.Sequential(*[
+            nn.Sequential(
+                nn.Conv2d(n_filter_list[i],
+                          n_filter_list[i + 1],
+                          kernel_size = (kernel_size, kernel_size),
+                          stride = (stride, stride),
+                          padding = (padding, padding),
+                          bias = conv_bias),
+                nn.Identity() if activation is None else activation(),
+                nn.MaxPool2d(kernel_size = pooling_kernel_size, stride = pooling_stride, padding = pooling_padding
+                            ) if max_pool else nn.Identity()) for i in range(n_conv_layers)
+        ])
+
+        self.flattener = nn.Flatten(2, 3)
+        self.apply(self.init_weight)
+
+    def sequence_length(self, n_channels = 3, height = 224, width = 224):
+        return self.forward(torch.zeros((1, n_channels, height, width))).shape[1]
+
+    def forward(self, x):
+        return self.flattener(self.conv_layers(x)).transpose(-2, -1)
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight)
+
+
+class TextTokenizer(nn.Module):
+
+    def __init__(self,
+                 kernel_size,
+                 stride,
+                 padding,
+                 pooling_kernel_size = 3,
+                 pooling_stride = 2,
+                 pooling_padding = 1,
+                 embedding_dim = 300,
+                 n_output_channels = 128,
+                 activation = None,
+                 max_pool = True,
+                 *args,
+                 **kwargs):
+        super(TextTokenizer, self).__init__()
+
+        self.max_pool = max_pool
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(1,
+                      n_output_channels,
+                      kernel_size = (kernel_size, embedding_dim),
+                      stride = (stride, 1),
+                      padding = (padding, 0),
+                      bias = False),
+            nn.Identity() if activation is None else activation(),
+            nn.MaxPool2d(
+                kernel_size = (pooling_kernel_size, 1), stride = (pooling_stride,
+                                                                  1), padding = (pooling_padding,
+                                                                                 0)) if max_pool else nn.Identity())
+
+        self.apply(self.init_weight)
+
+    def seq_len(self, seq_len = 32, embed_dim = 300):
+        return self.forward(torch.zeros((1, seq_len, embed_dim)))[0].shape[1]
+
+    def forward_mask(self, mask):
+        new_mask = mask.unsqueeze(1).float()
+        cnn_weight = torch.ones((1, 1, self.conv_layers[0].kernel_size[0]), device = mask.device, dtype = torch.float)
+        new_mask = F.conv1d(new_mask, cnn_weight, None, self.conv_layers[0].stride[0], self.conv_layers[0].padding[0],
+                            1, 1)
+        if self.max_pool:
+            new_mask = F.max_pool1d(new_mask, self.conv_layers[2].kernel_size[0], self.conv_layers[2].stride[0],
+                                    self.conv_layers[2].padding[0], 1, False, False)
+        new_mask = new_mask.squeeze(1)
+        new_mask = (new_mask > 0)
+        return new_mask
+
+    def forward(self, x, mask = None):
+        x = x.unsqueeze(1)
+        x = self.conv_layers(x)
+        x = x.transpose(1, 3).squeeze(1)
+        if mask is not None:
+            mask = self.forward_mask(mask).unsqueeze(-1).float()
+            x = x * mask
+        return x, mask
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight)