Upload model files

models/R2D2_embedding.py

@@ -0,0 +1,46 @@
+import torch.nn as nn
+import torch
+
+# Embedding network used in Meta-learning with differentiable closed-form solvers
+# (Bertinetto et al., in submission to NIPS 2018).
+# They call the ridge rigressor version as "Ridge Regression Differentiable Discriminator (R2D2)."
+  
+# Note that they use a peculiar ordering of functions, namely conv-BN-pooling-lrelu,
+# as opposed to the conventional one (conv-BN-lrelu-pooling).
+  
+def R2D2_conv_block(in_channels, out_channels, retain_activation=True, keep_prob=1.0):
+    block = nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1),
+        nn.BatchNorm2d(out_channels),
+        nn.MaxPool2d(2)
+    )
+    if retain_activation:
+        block.add_module("LeakyReLU", nn.LeakyReLU(0.1))
+
+    if keep_prob < 1.0:
+        block.add_module("Dropout", nn.Dropout(p=1 - keep_prob, inplace=False))
+
+    return block
+
+class R2D2Embedding(nn.Module):
+    def __init__(self, x_dim=3, h1_dim=96, h2_dim=192, h3_dim=384, z_dim=512, \
+                 retain_last_activation=False):
+        super(R2D2Embedding, self).__init__()
+
+        self.block1 = R2D2_conv_block(x_dim, h1_dim)
+        self.block2 = R2D2_conv_block(h1_dim, h2_dim)
+        self.block3 = R2D2_conv_block(h2_dim, h3_dim, keep_prob=0.9)
+        # In the last conv block, we disable activation function to boost the classification accuracy.
+        # This trick was proposed by Gidaris et al. (CVPR 2018).
+        # With this trick, the accuracy goes up from 50% to 51%.
+        # Although the authors of R2D2 did not mention this trick in the paper,
+        # we were unable to reproduce the result of Bertinetto et al. without resorting to this trick.
+        self.block4 = R2D2_conv_block(h3_dim, z_dim, retain_activation=retain_last_activation, keep_prob=0.7)
+  
+    def forward(self, x):
+        b1 = self.block1(x)
+        b2 = self.block2(b1)
+        b3 = self.block3(b2)
+        b4 = self.block4(b3)
+        # Flatten and concatenate the output of the 3rd and 4th conv blocks as proposed in R2D2 paper.
+        return torch.cat((b3.view(b3.size(0), -1), b4.view(b4.size(0), -1)), 1)

models/ResNet12_embedding.py

@@ -0,0 +1,266 @@
+import timm
+import numpy as np
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from models.dropblock import DropBlock
+from models.vit import ViT
+from torchvision import models
+import random
+
+# This ResNet network was designed following the practice of the following papers:
+# TADAM: Task dependent adaptive metric for improved few-shot learning (Oreshkin et al., in NIPS 2018) and
+# A Simple Neural Attentive Meta-Learner (Mishra et al., in ICLR 2018).
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, drop_rate=0.0, drop_block=False, block_size=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.LeakyReLU(0.1)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv3x3(planes, planes)
+        self.bn3 = nn.BatchNorm2d(planes)
+        self.maxpool = nn.MaxPool2d(stride)
+        self.downsample = downsample
+        self.stride = stride
+        self.drop_rate = drop_rate
+        self.num_batches_tracked = 0
+        self.drop_block = drop_block
+        self.block_size = block_size
+        self.DropBlock = DropBlock(block_size=self.block_size)
+
+    def forward(self, x):
+        self.num_batches_tracked += 1
+
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        out += residual
+        out = self.relu(out)
+        out = self.maxpool(out)
+
+        if self.drop_rate > 0:
+            if self.drop_block == True:
+                feat_size = out.size()[2]
+                keep_rate = max(1.0 - self.drop_rate / (20*2000) *
+                                (self.num_batches_tracked), 1.0 - self.drop_rate)
+                gamma = (1 - keep_rate) / self.block_size**2 * \
+                    feat_size**2 / (feat_size - self.block_size + 1)**2
+                out = self.DropBlock(out, gamma=gamma)
+            else:
+                out = F.dropout(out, p=self.drop_rate,
+                                training=self.training, inplace=True)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, keep_prob=1.0, avg_pool=False, drop_rate=0.0, dropblock_size=5):
+        self.inplanes = 3
+        super(ResNet, self).__init__()
+
+        self.layer1 = self._make_layer(
+            block, 64, stride=2, drop_rate=drop_rate)
+        self.layer2 = self._make_layer(
+            block, 160, stride=2, drop_rate=drop_rate)
+        self.layer3 = self._make_layer(
+            block, 320, stride=2, drop_rate=drop_rate, drop_block=True, block_size=dropblock_size)
+        self.layer4 = self._make_layer(
+            block, 640, stride=2, drop_rate=drop_rate, drop_block=True, block_size=dropblock_size)
+        if avg_pool:
+            self.avgpool = nn.AvgPool2d(5, stride=1)
+        self.keep_prob = keep_prob
+        self.keep_avg_pool = avg_pool
+        self.dropout = nn.Dropout(p=1 - self.keep_prob, inplace=False)
+        self.drop_rate = drop_rate
+        self.linear1_1 = nn.Linear(2560, 512)
+        self.linear1_2 = nn.Linear(512, 64)
+        self.linear2_1 = nn.Linear(2560, 512)
+        self.linear2_2 = nn.Linear(512, 64)
+        self.linear3_1 = nn.Linear(2560, 512)
+        self.linear3_2 = nn.Linear(512, 64)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode='fan_out', nonlinearity='leaky_relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, planes, stride=1, drop_rate=0.0, drop_block=False, block_size=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=1, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride,
+                            downsample, drop_rate, drop_block, block_size))
+        self.inplanes = planes * block.expansion
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        if self.keep_avg_pool:
+            x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+
+        # x1 = F.relu(self.linear1_1(x))
+        # x1 = self.linear1_2(x1)
+        # x2 = F.relu(self.linear2_1(x))
+        # x2 = self.linear2_2(x2)
+        # x3 = F.relu(self.linear3_1(x))
+        # x3 = self.linear3_2(x3)
+        # return [x1, x2, x3]
+        return x
+        # return x
+
+
+
+class custom_model(nn.Module):
+
+    def __init__(self,num_layer=5):
+        super(custom_model, self).__init__()
+        # self.classifier = timm.create_model('densenet121', pretrained=True)
+        self.classifier = timm.create_model('tf_efficientnet_b7_ns', pretrained=True)
+        # self.classifier = nn.Sequential(*list(classifier.children())[:-1])
+        self.num_layer = num_layer
+        # self.classifier = models.resnet34(pretrained=True, progress=True)
+        # self.classifier = ViT(
+        #     image_size = 32,
+        #     patch_size = 3,
+        #     dim = 512,
+        #     depth = 6,
+        #     heads = 8,
+        #     mlp_dim = 1000,
+        #     dropout = 0.1,
+        #     emb_dropout = 0.1
+        # )
+
+        # self.bn1 = nn.BatchNorm1d(num_features=1000)
+        self.bn2 = nn.BatchNorm1d(num_features=128)
+        self.dropout = nn.Dropout(0.4)
+
+        for i in range(num_layer):
+            setattr(self, "linear%d_1" % i, nn.Linear(1000,512))
+            setattr(self, "batch_norm%d_1" % i, nn.BatchNorm1d(num_features=512))
+            setattr(self, "linear%d_2" % i, nn.Linear(512,64))
+            # setattr(self, "linear%d_3" % i, nn.Linear(128,64))
+
+        for m in self.modules():
+
+            # if isinstance(m, nn.Conv2d):
+            #     nn.init.kaiming_normal_(
+            #         m.weight, mode='fan_out', nonlinearity='leaky_relu')
+            # elif isinstance(m, nn.BatchNorm2d):
+            #     nn.init.constant_(m.weight, 1)
+            #     nn.init.constant_(m.bias, 0)
+
+            if isinstance(m, nn.Linear):
+
+                # m.weight = nn.parameter.Parameter(torch.randn(m.out_features,m.in_features) * torch.sqrt(torch.tensor(2/m.in_features,requires_grad = True)))
+                y = m.in_features
+                y = random.randint(m.in_features/2,m.in_features)
+                m.weight.data.normal_(0.0, 1/np.sqrt(y))
+                # m.bias.data should be 0
+                # m.bias.data.fill_(0)
+
+        self.act = nn.ReLU()
+
+    def forward(self, x):
+        x = self.classifier(x)
+        feat = []
+
+        # print(self.classifier.classifier.weight)
+
+        # feat = torch.zeros((self.num_layer,256))
+        for i in range(self.num_layer):
+
+            x1 = self.act(getattr(self,"batch_norm%d_1" % i)(getattr(self, "linear%d_1" % i)(x)))
+            x2 = getattr(self, "linear%d_2" % i)(x1)
+            feat.append(x2)
+            # print(getattr(self, "linear%d_1" % i).weight)
+            # weights.append(getattr(self, "linear%d_1" % i).weight)
+
+
+        feat = torch.stack(feat ,dim = 0)
+        # weights = torch.stack(weights,dim = 0)
+        # print(feat.size())
+        return feat
+
+
+
+
+def conv_block(in_channels, out_channels):
+    '''
+    returns a block conv-bn-relu-pool
+    '''
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(),
+        nn.MaxPool2d(2)
+    )
+
+
+class ProtoNet(nn.Module):
+    '''
+    Model as described in the reference paper,
+    source: https://github.com/jakesnell/prototypical-networks/blob/f0c48808e496989d01db59f86d4449d7aee9ab0c/protonets/models/few_shot.py#L62-L84
+    '''
+    def __init__(self, x_dim=3, hid_dim=64, z_dim=64):
+        super(ProtoNet, self).__init__()
+        self.encoder = nn.Sequential(
+            conv_block(x_dim, hid_dim),
+            conv_block(hid_dim, hid_dim),
+            conv_block(hid_dim, hid_dim),
+            conv_block(hid_dim, z_dim),
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        return [x.view(x.size(0), -1)]
+
+
+def resnet12(keep_prob=1.0, avg_pool=False,num_layer=5, **kwargs):
+    """Constructs a ResNet-12 model.
+    """
+    # model = ResNet(BasicBlock, keep_prob=keep_prob,
+    #                avg_pool=avg_pool, **kwargs)
+    model = custom_model(num_layer=num_layer)
+    # model = ProtoNet()
+
+    return model
+

models/classification_heads.py

@@ -0,0 +1,367 @@
+import os
+import sys
+
+import torch
+from torch.autograd import Variable
+import torch.nn as nn
+#from qpth.qp import QPFunction
+import torch.nn.functional as F
+
+
+def sqrt_newton_schulz(A, numIters):
+    dim = A.shape[0]
+    normA = A.mul(A).sum(dim=0).sum(dim=0).sqrt()
+    Y = A.div(normA.expand_as(A))
+    I = torch.eye(dim, dim).float().cuda()
+    Z = torch.eye(dim, dim).float().cuda()
+    for i in range(numIters):
+        T = 0.5 * (3.0 * I - Z.mm(Y))
+        Y = Y.mm(T)
+        Z = T.mm(Z)
+
+    # sA = Y * torch.sqrt(normA).view(batchSize, 1, 1).expand_as(A)
+
+    # sA = Y * torch.sqrt(normA).expand_as(A)
+
+    sZ = Z * 1. / torch.sqrt(normA).expand_as(A)
+    return sZ
+
+
+def polar_decompose(input):
+    # square_mat = input.mm(input.transpose(0, 1))
+    # square_mat = square_mat/torch.norm(torch.diag(square_mat), p=1)
+    # ortho_mat = self.sqrt_newton_schulz(square_mat, numIters=1)
+
+    square_mat = input.transpose(0, 1).mm(input)
+    sA_minushalf = sqrt_newton_schulz(square_mat, 1)
+    ortho_mat = input.mm(sA_minushalf)
+
+    # return ortho_mat
+
+    return ortho_mat.mm(ortho_mat.transpose(0, 1))
+
+
+def computeGramMatrix(A, B):
+    """
+    Constructs a linear kernel matrix between A and B.
+    We assume that each row in A and B represents a d-dimensional feature vector.
+    
+    Parameters:
+      A:  a (n_batch, n, d) Tensor.
+      B:  a (n_batch, m, d) Tensor.
+    Returns: a (n_batch, n, m) Tensor.
+    """
+    
+    assert(A.dim() == 3)
+    assert(B.dim() == 3)
+    assert(A.size(0) == B.size(0) and A.size(2) == B.size(2))
+
+    return torch.bmm(A, B.transpose(1,2))
+
+
+def binv(b_mat):
+    """
+    Computes an inverse of each matrix in the batch.
+    Pytorch 0.4.1 does not support batched matrix inverse.
+    Hence, we are solving AX=I.
+    
+    Parameters:
+      b_mat:  a (n_batch, n, n) Tensor.
+    Returns: a (n_batch, n, n) Tensor.
+    """
+
+    id_matrix = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat).cuda()
+    b_inv, _ = torch.gesv(id_matrix, b_mat)
+    
+    return b_inv
+
+
+def one_hot(indices, depth):
+    """
+    Returns a one-hot tensor.
+    This is a PyTorch equivalent of Tensorflow's tf.one_hot.
+        
+    Parameters:
+      indices:  a (n_batch, m) Tensor or (m) Tensor.
+      depth: a scalar. Represents the depth of the one hot dimension.
+    Returns: a (n_batch, m, depth) Tensor or (m, depth) Tensor.
+    """
+
+    encoded_indicies = torch.zeros(indices.size() + torch.Size([depth])).cuda()
+    index = indices.view(indices.size()+torch.Size([1]))
+    encoded_indicies = encoded_indicies.scatter_(1,index,1)
+    
+    return encoded_indicies
+
+def batched_kronecker(matrix1, matrix2):
+    matrix1_flatten = matrix1.reshape(matrix1.size()[0], -1)
+    matrix2_flatten = matrix2.reshape(matrix2.size()[0], -1)
+    return torch.bmm(matrix1_flatten.unsqueeze(2), matrix2_flatten.unsqueeze(1)).reshape([matrix1.size()[0]] + list(matrix1.size()[1:]) + list(matrix2.size()[1:])).permute([0, 1, 3, 2, 4]).reshape(matrix1.size(0), matrix1.size(1) * matrix2.size(1), matrix1.size(2) * matrix2.size(2))
+
+
+
+
+#################  uncomment this if you have installed QPFunction and run Ridge
+# def MetaOptNetHead_Ridge(query, support, support_labels, n_way, n_shot, lambda_reg=50.0, double_precision=True):
+#     """
+#     Fits the support set with ridge regression and
+#     returns the classification score on the query set.
+#
+#     Parameters:
+#       query:  a (tasks_per_batch, n_query, d) Tensor.
+#       support:  a (tasks_per_batch, n_support, d) Tensor.
+#       support_labels: a (tasks_per_batch, n_support) Tensor.
+#       n_way: a scalar. Represents the number of classes in a few-shot classification task.
+#       n_shot: a scalar. Represents the number of support examples given per class.
+#       lambda_reg: a scalar. Represents the strength of L2 regularization.
+#     Returns: a (tasks_per_batch, n_query, n_way) Tensor.
+#     """
+#
+#     tasks_per_batch = query.size(0)
+#     n_support = support.size(1)
+#     n_query = query.size(1)
+#
+#     assert(query.dim() == 3)
+#     assert(support.dim() == 3)
+#     assert(query.size(0) == support.size(0) and query.size(2) == support.size(2))
+#     assert(n_support == n_way * n_shot)      # n_support must equal to n_way * n_shot
+#
+#     #Here we solve the dual problem:
+#     #Note that the classes are indexed by m & samples are indexed by i.
+#     #min_{\alpha}  0.5 \sum_m ||w_m(\alpha)||^2 + \sum_i \sum_m e^m_i alpha^m_i
+#
+#     #where w_m(\alpha) = \sum_i \alpha^m_i x_i,
+#
+#     #\alpha is an (n_support, n_way) matrix
+#     kernel_matrix = computeGramMatrix(support, support)
+#     kernel_matrix += lambda_reg * torch.eye(n_support).expand(tasks_per_batch, n_support, n_support).cuda()
+#
+#     block_kernel_matrix = kernel_matrix.repeat(n_way, 1, 1) #(n_way * tasks_per_batch, n_support, n_support)
+#
+#     support_labels_one_hot = one_hot(support_labels.view(tasks_per_batch * n_support), n_way) # (tasks_per_batch * n_support, n_way)
+#     support_labels_one_hot = support_labels_one_hot.transpose(0, 1) # (n_way, tasks_per_batch * n_support)
+#     support_labels_one_hot = support_labels_one_hot.reshape(n_way * tasks_per_batch, n_support)     # (n_way*tasks_per_batch, n_support)
+#
+#     G = block_kernel_matrix
+#     e = -2.0 * support_labels_one_hot
+#
+#     #This is a fake inequlity constraint as qpth does not support QP without an inequality constraint.
+#     id_matrix_1 = torch.zeros(tasks_per_batch*n_way, n_support, n_support)
+#     C = Variable(id_matrix_1)
+#     h = Variable(torch.zeros((tasks_per_batch*n_way, n_support)))
+#     dummy = Variable(torch.Tensor()).cuda()      # We want to ignore the equality constraint.
+#
+#     #if double_precision:
+#     G, e, C, h = [x.double().cuda() for x in [G, e, C, h]]
+#
+#
+#     qp_sol = QPFunction(verbose=False)(G, e.detach(), C.detach(), h.detach(), dummy.detach(), dummy.detach())
+#     qp_sol = qp_sol.reshape(n_way, tasks_per_batch, n_support)
+#     qp_sol = qp_sol.permute(1, 2, 0)
+#
+#
+#     # Compute the classification score.
+#     compatibility = computeGramMatrix(support, query)
+#     compatibility = compatibility.float()
+#     compatibility = compatibility.unsqueeze(3).expand(tasks_per_batch, n_support, n_query, n_way)
+#     qp_sol = qp_sol.reshape(tasks_per_batch, n_support, n_way)
+#     logits = qp_sol.float().unsqueeze(2).expand(tasks_per_batch, n_support, n_query, n_way)
+#     logits = logits * compatibility
+#     logits = torch.sum(logits, 1)
+#
+#     return logits
+
+def R2D2Head(query, support, support_labels, n_way, n_shot, l2_regularizer_lambda=50.0):
+    """
+    Fits the support set with ridge regression and 
+    returns the classification score on the query set.
+    
+    This model is the classification head described in:
+    Meta-learning with differentiable closed-form solvers
+    (Bertinetto et al., in submission to NIPS 2018).
+    
+    Parameters:
+      query:  a (tasks_per_batch, n_query, d) Tensor.
+      support:  a (tasks_per_batch, n_support, d) Tensor.
+      support_labels: a (tasks_per_batch, n_support) Tensor.
+      n_way: a scalar. Represents the number of classes in a few-shot classification task.
+      n_shot: a scalar. Represents the number of support examples given per class.
+      l2_regularizer_lambda: a scalar. Represents the strength of L2 regularization.
+    Returns: a (tasks_per_batch, n_query, n_way) Tensor.
+    """
+    
+    tasks_per_batch = query.size(0)
+    n_support = support.size(1)
+
+    assert(query.dim() == 3)
+    assert(support.dim() == 3)
+    assert(query.size(0) == support.size(0) and query.size(2) == support.size(2))
+    assert(n_support == n_way * n_shot)      # n_support must equal to n_way * n_shot
+    
+    support_labels_one_hot = one_hot(support_labels.view(tasks_per_batch * n_support), n_way)
+    support_labels_one_hot = support_labels_one_hot.view(tasks_per_batch, n_support, n_way)
+
+    id_matrix = torch.eye(n_support).expand(tasks_per_batch, n_support, n_support).cuda()
+    
+    # Compute the dual form solution of the ridge regression.
+    # W = X^T(X X^T - lambda * I)^(-1) Y
+    ridge_sol = computeGramMatrix(support, support) + l2_regularizer_lambda * id_matrix
+    ridge_sol = binv(ridge_sol)
+    ridge_sol = torch.bmm(support.transpose(1,2), ridge_sol)
+    ridge_sol = torch.bmm(ridge_sol, support_labels_one_hot)
+    
+    # Compute the classification score.
+    # score = W X
+    logits = torch.bmm(query, ridge_sol)
+
+    return logits
+
+
+
+def ProtoNetHead(query, support, support_labels, n_way, n_shot, normalize=True):
+    """
+    Constructs the prototype representation of each class(=mean of support vectors of each class) and 
+    returns the classification score (=L2 distance to each class prototype) on the query set.
+    
+    This model is the classification head described in:
+    Prototypical Networks for Few-shot Learning
+    (Snell et al., NIPS 2017).
+    
+    Parameters:
+      query:  a (tasks_per_batch, n_query, d) Tensor.
+      support:  a (tasks_per_batch, n_support, d) Tensor.
+      support_labels: a (tasks_per_batch, n_support) Tensor.
+      n_way: a scalar. Represents the number of classes in a few-shot classification task.
+      n_shot: a scalar. Represents the number of support examples given per class.
+      normalize: a boolean. Represents whether if we want to normalize the distances by the embedding dimension.
+    Returns: a (tasks_per_batch, n_query, n_way) Tensor.
+    """
+    
+    tasks_per_batch = query.size(0)
+    n_support = support.size(1)
+    n_query = query.size(1)
+    d = query.size(2)
+    
+    assert(query.dim() == 3)
+    assert(support.dim() == 3)
+    assert(query.size(0) == support.size(0) and query.size(2) == support.size(2))
+    assert(n_support == n_way * n_shot)      # n_support must equal to n_way * n_shot
+    
+    support_labels_one_hot = one_hot(support_labels.view(tasks_per_batch * n_support), n_way)
+    support_labels_one_hot = support_labels_one_hot.view(tasks_per_batch, n_support, n_way)
+    
+    # From:
+    # https://github.com/gidariss/FewShotWithoutForgetting/blob/master/architectures/PrototypicalNetworksHead.py
+    #************************* Compute Prototypes **************************
+    labels_train_transposed = support_labels_one_hot.transpose(1,2)
+
+    prototypes = torch.bmm(labels_train_transposed, support)
+    # Divide with the number of examples per novel category.
+    prototypes = prototypes.div(
+        labels_train_transposed.sum(dim=2, keepdim=True).expand_as(prototypes)
+    )
+
+    # Distance Matrix Vectorization Trick
+    AB = computeGramMatrix(query, prototypes)
+    AA = (query * query).sum(dim=2, keepdim=True)
+    BB = (prototypes * prototypes).sum(dim=2, keepdim=True).reshape(tasks_per_batch, 1, n_way)
+    logits = AA.expand_as(AB) - 2 * AB + BB.expand_as(AB)
+    logits = -logits
+    
+    if normalize:
+        logits = logits / d
+
+    return logits
+
+
+
+def SubspaceNetHead(query, support, support_labels, n_way, n_shot, normalize=True):
+    """
+       Constructs the subspace representation of each class(=mean of support vectors of each class) and
+       returns the classification score (=L2 distance to each class prototype) on the query set.
+
+        Our algorithm using subspaces here
+
+       Parameters:
+         query:  a (tasks_per_batch, n_query, d) Tensor.
+         support:  a (tasks_per_batch, n_support, d) Tensor.
+         support_labels: a (tasks_per_batch, n_support) Tensor.
+         n_way: a scalar. Represents the number of classes in a few-shot classification task.
+         n_shot: a scalar. Represents the number of support examples given per class.
+         normalize: a boolean. Represents whether if we want to normalize the distances by the embedding dimension.
+       Returns: a (tasks_per_batch, n_query, n_way) Tensor.
+       """
+
+    tasks_per_batch = query.size(0)
+    n_support = support.size(1)
+    n_query = query.size(1)
+    d = query.size(2)
+
+    assert(query.dim() == 3)
+    assert(support.dim() == 3)
+    assert(query.size(0) == support.size(0) and query.size(2) == support.size(2))
+    assert(n_support == n_way * n_shot)      # n_support must equal to n_way * n_shot
+
+    support_labels_one_hot = one_hot(support_labels.view(tasks_per_batch * n_support), n_way)
+    #support_labels_one_hot = support_labels_one_hot.view(tasks_per_batch, n_support, n_way)
+
+
+    support_reshape = support.view(tasks_per_batch * n_support, -1)
+
+    support_labels_reshaped = support_labels.contiguous().view(-1)
+    class_representatives = []
+    for nn in range(n_way):
+        idxss = torch.nonzero((support_labels_reshaped == nn),as_tuple=False)
+        all_support_perclass = support_reshape[idxss, :]
+        class_representatives.append(all_support_perclass.view(tasks_per_batch, n_shot, -1))
+
+    class_representatives = torch.stack(class_representatives)
+    class_representatives = class_representatives.transpose(0, 1) #tasks_per_batch, n_way, n_support, -1
+    class_representatives = class_representatives.transpose(2, 3).contiguous().view(tasks_per_batch*n_way, -1, n_shot)
+
+    dist = []
+    for cc in range(tasks_per_batch*n_way):
+        batch_idx = cc//n_way
+        qq = query[batch_idx]
+        uu, _, _ = torch.svd(class_representatives[cc].double())
+        uu = uu.float()
+        subspace = uu[:, :n_shot-1].transpose(0, 1)
+        projection = subspace.transpose(0, 1).mm(subspace.mm(qq.transpose(0, 1))).transpose(0, 1)
+        dist_perclass = torch.sum((qq - projection)**2, dim=-1)
+        dist.append(dist_perclass)
+
+    dist = torch.stack(dist).view(tasks_per_batch, n_way, -1).transpose(1, 2)
+    logits = -dist
+
+    if normalize:
+        logits = logits / d
+
+    return logits
+
+
+
+
+class ClassificationHead(nn.Module):
+    def __init__(self, base_learner='MetaOptNet', enable_scale=True):
+        super(ClassificationHead, self).__init__()
+        if ('Subspace' in base_learner):
+            self.head = SubspaceNetHead
+        elif ('Ridge' in base_learner):
+            self.head = MetaOptNetHead_Ridge
+        elif ('R2D2' in base_learner):
+            self.head = R2D2Head
+        elif ('Proto' in base_learner):
+            self.head = ProtoNetHead
+        else:
+            print ("Cannot recognize the base learner type")
+            assert(False)
+        
+        # Add a learnable scale
+        self.enable_scale = enable_scale
+        self.scale = nn.Parameter(torch.FloatTensor([1.0]))
+        
+    def forward(self, query, support, support_labels, n_way, n_shot, **kwargs):
+        if self.enable_scale:
+            return self.scale * self.head(query, support, support_labels, n_way, n_shot, **kwargs)
+        else:
+            return self.head(query, support, support_labels, n_way, n_shot, **kwargs)

models/closerlook_classifier.py

@@ -0,0 +1,28 @@
+import torch.nn as nn
+import torch
+from torch.nn.utils.weight_norm import WeightNorm
+
+
+class distLinear(nn.Module):
+    def __init__(self, indim, outdim):
+        super(distLinear, self).__init__()
+        self.L = nn.Linear( indim, outdim, bias = False)
+        self.class_wise_learnable_norm = True  #See the issue#4&8 in the github
+        if self.class_wise_learnable_norm:
+            WeightNorm.apply(self.L, 'weight', dim=0) #split the weight update component to direction and norm
+
+        if outdim <=200:
+            self.scale_factor = 2; #a fixed scale factor to scale the output of cos value into a reasonably large input for softmax, for to reproduce the result of CUB with ResNet10, use 4. see the issue#31 in the github
+        else:
+            self.scale_factor = 10; #in omniglot, a larger scale factor is required to handle >1000 output classes.
+
+    def forward(self, x):
+        x_norm = torch.norm(x, p=2, dim =1).unsqueeze(1).expand_as(x)
+        x_normalized = x.div(x_norm+ 0.00001)
+        if not self.class_wise_learnable_norm:
+            L_norm = torch.norm(self.L.weight.data, p=2, dim =1).unsqueeze(1).expand_as(self.L.weight.data)
+            self.L.weight.data = self.L.weight.data.div(L_norm + 0.00001)
+        cos_dist = self.L(x_normalized) #matrix product by forward function, but when using WeightNorm, this also multiply the cosine distance by a class-wise learnable norm, see the issue#4&8 in the github
+        scores = self.scale_factor* (cos_dist)
+
+        return scores

models/dropblock.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.distributions import Bernoulli
+
+
+class DropBlock(nn.Module):
+    def __init__(self, block_size):
+        super(DropBlock, self).__init__()
+
+        self.block_size = block_size
+        #self.gamma = gamma
+        #self.bernouli = Bernoulli(gamma)
+
+    def forward(self, x, gamma):
+        # shape: (bsize, channels, height, width)
+
+        if self.training:
+            batch_size, channels, height, width = x.shape
+            
+            bernoulli = Bernoulli(gamma)
+            mask = bernoulli.sample((batch_size, channels, height - (self.block_size - 1), width - (self.block_size - 1))).cuda()
+            #print((x.sample[-2], x.sample[-1]))
+            block_mask = self._compute_block_mask(mask)
+            #print (block_mask.size())
+            #print (x.size())
+            countM = block_mask.size()[0] * block_mask.size()[1] * block_mask.size()[2] * block_mask.size()[3]
+            count_ones = block_mask.sum()
+
+            return block_mask * x * (countM / count_ones)
+        else:
+            return x
+
+    def _compute_block_mask(self, mask):
+        left_padding = int((self.block_size-1) / 2)
+        right_padding = int(self.block_size / 2)
+        
+        batch_size, channels, height, width = mask.shape
+        #print ("mask", mask[0][0])
+        non_zero_idxs = torch.nonzero(mask,as_tuple=False)
+        # print(type(non_zero_idxs))
+        # print(type(non_zero_idxs))
+        nr_blocks = non_zero_idxs.size(0)
+
+        offsets = torch.stack(
+            [
+                torch.arange(self.block_size).view(-1, 1).expand(self.block_size, self.block_size).reshape(-1), # - left_padding,
+                torch.arange(self.block_size).repeat(self.block_size), #- left_padding
+            ]
+        ).t().cuda()
+        offsets = torch.cat((torch.zeros(self.block_size**2, 2).cuda().long(), offsets.long()), 1)
+        
+        if nr_blocks > 0:
+            non_zero_idxs = non_zero_idxs.repeat(self.block_size ** 2, 1)
+            offsets = offsets.repeat(nr_blocks, 1).view(-1, 4)
+            offsets = offsets.long()
+
+            block_idxs = non_zero_idxs + offsets
+            #block_idxs += left_padding
+            padded_mask = F.pad(mask, (left_padding, right_padding, left_padding, right_padding))
+            padded_mask[block_idxs[:, 0], block_idxs[:, 1], block_idxs[:, 2], block_idxs[:, 3]] = 1.
+        else:
+            padded_mask = F.pad(mask, (left_padding, right_padding, left_padding, right_padding))
+            
+        block_mask = 1 - padded_mask#[:height, :width]
+        return block_mask

models/protonet_embedding.py

@@ -0,0 +1,43 @@
+import torch.nn as nn
+import math
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, retain_activation=True):
+        super(ConvBlock, self).__init__()
+        
+        self.block = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels)
+        )
+        
+        if retain_activation:
+            self.block.add_module("ReLU", nn.ReLU(inplace=True))
+        self.block.add_module("MaxPool2d", nn.MaxPool2d(kernel_size=2, stride=2, padding=0))
+        
+    def forward(self, x):
+        out = self.block(x)
+        return out
+
+# Embedding network used in Matching Networks (Vinyals et al., NIPS 2016), Meta-LSTM (Ravi & Larochelle, ICLR 2017),
+# MAML (w/ h_dim=z_dim=32) (Finn et al., ICML 2017), Prototypical Networks (Snell et al. NIPS 2017).
+
+class ProtoNetEmbedding(nn.Module):
+    def __init__(self, x_dim=3, h_dim=64, z_dim=64, retain_last_activation=True):
+        super(ProtoNetEmbedding, self).__init__()
+        self.encoder = nn.Sequential(
+          ConvBlock(x_dim, h_dim),
+          ConvBlock(h_dim, h_dim),
+          ConvBlock(h_dim, h_dim),
+          ConvBlock(h_dim, z_dim, retain_activation=retain_last_activation),
+        )
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.encoder(x)
+        return x.view(x.size(0), -1)

models/vit.py

@@ -0,0 +1,127 @@
+# https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit_pytorch.py
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+MIN_NUM_PATCHES = 16
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + x
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dropout = 0.):
+        super().__init__()
+        self.heads = heads
+        self.scale = dim ** -0.5
+
+        self.to_qkv = nn.Linear(dim, dim * 3, bias = False)
+        self.to_out = nn.Sequential(
+            nn.Linear(dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, mask = None):
+        b, n, _, h = *x.shape, self.heads
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+        dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale
+
+        if mask is not None:
+            mask = F.pad(mask.flatten(1), (1, 0), value = True)
+            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
+            mask = mask[:, None, :] * mask[:, :, None]
+            dots.masked_fill_(~mask, float('-inf'))
+            del mask
+
+        attn = dots.softmax(dim=-1)
+
+        out = torch.einsum('bhij,bhjd->bhid', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        out =  self.to_out(out)
+        return out
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, mlp_dim, dropout):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Residual(PreNorm(dim, Attention(dim, heads = heads, dropout = dropout))),
+                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
+            ]))
+    def forward(self, x, mask = None):
+        for attn, ff in self.layers:
+            x = attn(x, mask = mask)
+            x = ff(x)
+        return x
+
+class ViT(nn.Module):
+    def __init__(self, *, image_size, patch_size, dim, depth, heads, mlp_dim, channels = 3, dropout = 0., emb_dropout = 0.):
+        super().__init__()
+        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = channels * patch_size ** 2
+        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective. try decreasing your patch size'
+
+        self.patch_size = patch_size
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.patch_to_embedding = nn.Linear(patch_dim, dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, mlp_dim, dropout)
+
+        self.to_cls_token = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, mlp_dim),
+            nn.GELU(),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, img, mask = None):
+        p = self.patch_size
+
+        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
+        x = self.patch_to_embedding(x)
+        b, n, _ = x.shape
+
+        cls_tokens = self.cls_token.expand(b, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x, mask)
+
+        x = self.to_cls_token(x[:, 0])
+        return self.mlp_head(x)