init submit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+PuzzleTuning/Counterpart[[:space:]]PreTrain[[:space:]]Methods/gcmae/nohup.out filter=lfs diff=lfs merge=lfs -text

PuzzleTuning/Archive/PuzzleTuning_demoset.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:363cbdc5f8e944f99542a5727d4b2457dd79bbfca7eb70d76857de4b8be92858
+size 4987321

PuzzleTuning/Archive/build_CPIA_dataset.sh

@@ -0,0 +1,35 @@
+#!/bin/sh
+# go to the dataset location
+
+# altering the zip files
+
+zip -F L.zip --out L_Scale.zip
+zip -FF L_Scale.zip --out L.zip -fz
+zip -F M.zip --out M_Scale.zip
+zip -FF M_Scale.zip --out M.zip -fz
+
+rm -f L_Scale.zip
+rm -f L.z01
+rm -f M_Scale.zip
+rm -f M.z01
+rm -f M.z02
+
+# build a directory of datasets
+mkdir datasets
+mv L.zip datasets
+mv M.zip datasets
+mv S.zip datasets
+
+cd datasets
+unzip L.zip
+unzip M.zip
+unzip S.zip
+
+rm -f L.zip
+rm -f M.zip
+rm -f S.zip
+
+mkdir All
+cp -r L/* All/ &
+cp -r M/* All/ &
+cp -r S/* All/

PuzzleTuning/Archive/warwick_CLS.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9813cdc86e2420476e06638965df6040cec0197b148be786f0bf88b020e445f6
+size 15098031

PuzzleTuning/Backbone/GetPromptModel.py

@@ -0,0 +1,87 @@
+"""
+build_promptmodel   Script  ver： Oct 17th 14:20
+
+"""
+
+try:
+    from Backbone.VPT_structure import *
+except:
+    from Backbone.VPT_structure import *
+
+
+def build_promptmodel(num_classes=1000, edge_size=224, model_idx='ViT', patch_size=16,
+                      Prompt_Token_num=20, VPT_type="Deep", prompt_state_dict=None, base_state_dict='timm'):
+    """
+    following the https://github.com/sagizty/VPT
+    this build the VPT (prompt version of ViT), with additional prompt tokens,
+    each layer the information become [B, N_patch + N_prompt, Dim]
+
+    During training only the prompt tokens and the head layer are
+    set to be learnable while the rest of Transformer layers are frozen
+
+    # VPT_type = "Shallow" / "Deep"
+        - Shallow: concatenate N_prompt of prompt tokens before the first Transformer Encoder block,
+                each layer the information become [B, N_patch + N_prompt, Dim]
+        - Deep: concatenate N_prompt of prompt tokens to each Transformer Encoder block,
+                this will replace the output prompt tokens learnt form previous encoder.
+    """
+
+    if model_idx[0:3] == 'ViT':
+
+        if base_state_dict is None:
+            basic_state_dict = None
+
+        elif type(base_state_dict) == str:
+            if base_state_dict == 'timm':
+                # ViT_Prompt
+                import timm
+                # from pprint import pprint
+                # model_names = timm.list_models('*vit*')
+                # pprint(model_names)
+
+                basic_model = timm.create_model('vit_base_patch' + str(patch_size) + '_' + str(edge_size),
+                                                pretrained=True)
+                basic_state_dict = basic_model.state_dict()
+                print('in prompt model building, timm ViT loaded for base_state_dict')
+
+            else:
+                basic_state_dict = None
+                print('in prompt model building, no vaild str for base_state_dict')
+
+        else:  # state dict: collections.OrderedDict
+            basic_state_dict = base_state_dict
+            print('in prompt model building, a .pth base_state_dict loaded')
+
+        model = VPT_ViT(img_size=edge_size, patch_size=patch_size, Prompt_Token_num=Prompt_Token_num,
+                        VPT_type=VPT_type, basic_state_dict=basic_state_dict)
+
+        model.New_CLS_head(num_classes)
+
+        if prompt_state_dict is not None:
+            try:
+                model.load_prompt(prompt_state_dict)
+            except:
+                print('erro in .pth prompt_state_dict')
+            else:
+                print('in prompt model building, a .pth prompt_state_dict loaded')
+
+        model.Freeze()
+    else:
+        print("The model is not difined in the Prompt script！！")
+        return -1
+
+    try:
+        img = torch.randn(1, 3, edge_size, edge_size)
+        preds = model(img)  # (1, class_number)
+        print('Build VPT model with in/out shape: ', img.shape, ' -> ', preds.shape)
+
+    except:
+        print("Problem exist in the model defining process！！")
+        return -1
+    else:
+        print('model is ready now!')
+        return model
+
+
+if __name__ == '__main__':
+    model = build_promptmodel(prompt_state_dict=None, base_state_dict='timm', num_classes=0)

PuzzleTuning/Backbone/ResHybrid.py

@@ -0,0 +1,481 @@
+"""
+MSHT
+
+Models  ver： OCT 27th 20：00 official release
+
+by the authors, check our github page:
+https://github.com/sagizty/Multi-Stage-Hybrid-Transformer
+
+
+ResNet stages' feature map
+
+# input = 3, 384, 384
+torch.Size([1, 256, 96, 96])
+torch.Size([1, 512, 48, 48])
+torch.Size([1, 1024, 24, 24])
+torch.Size([1, 2048, 12, 12])
+torch.Size([1, 1000])
+
+# input = 3, 224, 224
+torch.Size([1, 256, 56, 56])
+torch.Size([1, 512, 28, 28])
+torch.Size([1, 1024, 14, 14])
+torch.Size([1, 2048, 7, 7])
+torch.Size([1, 1000])
+
+ref
+https://note.youdao.com/ynoteshare1/index.html?id=5a7dbe1a71713c317062ddeedd97d98e&type=note
+"""
+import torch
+from torch import nn
+from functools import partial
+from torchsummary import summary
+import os
+from Backbone import Transformer_blocks
+
+
+# ResNet Bottleneck_block_constructor
+class Bottleneck_block_constructor(nn.Module):
+
+    extention = 4
+
+    # 定义初始化的网络和参数
+    def __init__(self, inplane, midplane, stride, downsample=None):
+        super(Bottleneck_block_constructor, self).__init__()
+
+        outplane = midplane * self.extention
+
+        self.conv1 = nn.Conv2d(inplane, midplane, kernel_size=1, stride=stride, bias=False)
+        self.bn1 = nn.BatchNorm2d(midplane)
+
+        self.conv2 = nn.Conv2d(midplane, midplane, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(midplane)
+
+        self.conv3 = nn.Conv2d(midplane, outplane, kernel_size=1, stride=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(midplane * self.extention)
+
+        self.relu = nn.ReLU(inplace=False)
+
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.relu(self.bn3(self.conv3(out)))
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        else:
+            residual = x
+
+        out += residual
+
+        out = self.relu(out)
+
+        return out
+
+
+# Hybrid_backbone of ResNets
+class Hybrid_backbone_4(nn.Module):
+
+    def __init__(self, block_constructor, bottleneck_channels_setting=None, identity_layers_setting=None,
+                 stage_stride_setting=None, fc_num_classes=None, feature_idx=None):
+
+        if bottleneck_channels_setting is None:
+            bottleneck_channels_setting = [64, 128, 256, 512]
+        if identity_layers_setting is None:
+            identity_layers_setting = [3, 4, 6, 3]
+        if stage_stride_setting is None:
+            stage_stride_setting = [1, 2, 2, 2]
+
+        self.inplane = 64
+        self.fc_num_classes = fc_num_classes
+        self.feature_idx = feature_idx
+
+        super(Hybrid_backbone_4, self).__init__()
+
+        self.block_constructor = block_constructor  # Bottleneck_block_constructor
+        self.bcs = bottleneck_channels_setting  # [64, 128, 256, 512]
+        self.ils = identity_layers_setting  # [3, 4, 6, 3]
+        self.sss = stage_stride_setting  # [1, 2, 2, 2]
+
+        # stem
+        # alter the RGB pic chanel to match inplane
+        self.conv1 = nn.Conv2d(3, self.inplane, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplane)
+        self.relu = nn.ReLU()
+        self.maxpool = nn.MaxPool2d(kernel_size=3, padding=1, stride=2)
+
+        # ResNet stages
+        self.layer1 = self.make_stage_layer(self.block_constructor, self.bcs[0], self.ils[0], self.sss[0])
+        self.layer2 = self.make_stage_layer(self.block_constructor, self.bcs[1], self.ils[1], self.sss[1])
+        self.layer3 = self.make_stage_layer(self.block_constructor, self.bcs[2], self.ils[2], self.sss[2])
+        self.layer4 = self.make_stage_layer(self.block_constructor, self.bcs[3], self.ils[3], self.sss[3])
+
+        # cls head
+        if self.fc_num_classes is not None:
+            self.avgpool = nn.AvgPool2d(7)
+            self.fc = nn.Linear(512 * self.block_constructor.extention, fc_num_classes)
+
+    def forward(self, x):
+
+        # stem
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        stem_out = self.maxpool(out)
+
+        # Resnet block of 4 stages
+        stage1_out = self.layer1(stem_out)
+        stage2_out = self.layer2(stage1_out)
+        stage3_out = self.layer3(stage2_out)
+        stage4_out = self.layer4(stage3_out)
+
+        if self.fc_num_classes is not None:
+            # connect to cls head mlp if asked
+            fc_out = self.avgpool(stage4_out)
+            fc_out = torch.flatten(fc_out, 1)
+            fc_out = self.fc(fc_out)
+
+        # get what we need for different usage
+        if self.feature_idx == 'stages':
+            if self.fc_num_classes is not None:
+                return stage1_out, stage2_out, stage3_out, stage4_out, fc_out
+            else:
+                return stage1_out, stage2_out, stage3_out, stage4_out
+        elif self.feature_idx == 'features':
+            if self.fc_num_classes is not None:
+                return stem_out, stage1_out, stage2_out, stage3_out, stage4_out, fc_out
+            else:
+                return stem_out, stage1_out, stage2_out, stage3_out, stage4_out
+        else:  # self.feature_idx is None
+            if self.fc_num_classes is not None:
+                return fc_out
+            else:
+                return stage4_out
+
+    def make_stage_layer(self, block_constractor, midplane, block_num, stride=1):
+        """
+        block:
+        midplane：usually = output chanel/4
+        block_num：
+        stride：stride of the ResNet Conv Block
+        """
+
+        block_list = []
+
+        outplane = midplane * block_constractor.extention  # extention
+
+        if stride != 1 or self.inplane != outplane:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplane, outplane, stride=stride, kernel_size=1, bias=False),
+                nn.BatchNorm2d(midplane * block_constractor.extention)
+            )
+        else:
+            downsample = None
+
+        # Conv Block
+        conv_block = block_constractor(self.inplane, midplane, stride=stride, downsample=downsample)
+        block_list.append(conv_block)
+
+        self.inplane = outplane  # update inplane for the next stage
+
+        # Identity Block
+        for i in range(1, block_num):
+            block_list.append(block_constractor(self.inplane, midplane, stride=1, downsample=None))
+
+        return nn.Sequential(*block_list)  # stack blocks
+
+
+class Hybrid_backbone_3(nn.Module):  # 3 stages version
+
+    def __init__(self, block_constructor, bottleneck_channels_setting=None, identity_layers_setting=None,
+                 stage_stride_setting=None, fc_num_classes=None, feature_idx=None):
+
+        if bottleneck_channels_setting is None:
+            bottleneck_channels_setting = [64, 128, 256]
+        if identity_layers_setting is None:
+            identity_layers_setting = [3, 4, 6]
+        if stage_stride_setting is None:
+            stage_stride_setting = [1, 2, 2]
+
+        self.inplane = 64
+        self.fc_num_classes = fc_num_classes
+        self.feature_idx = feature_idx
+
+        super(Hybrid_backbone_3, self).__init__()
+
+        self.block_constructor = block_constructor  # Bottleneck_block_constructor
+        self.bcs = bottleneck_channels_setting  # [64, 128, 256]
+        self.ils = identity_layers_setting  # [3, 4, 6]
+        self.sss = stage_stride_setting  # [1, 2, 2]
+
+        # stem
+        self.conv1 = nn.Conv2d(3, self.inplane, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplane)
+        self.relu = nn.ReLU()
+        self.maxpool = nn.MaxPool2d(kernel_size=3, padding=1, stride=2)
+
+        # ResNet 3 stages
+        self.layer1 = self.make_stage_layer(self.block_constructor, self.bcs[0], self.ils[0], self.sss[0])
+        self.layer2 = self.make_stage_layer(self.block_constructor, self.bcs[1], self.ils[1], self.sss[1])
+        self.layer3 = self.make_stage_layer(self.block_constructor, self.bcs[2], self.ils[2], self.sss[2])
+
+        if self.fc_num_classes is not None:
+            self.avgpool = nn.AvgPool2d(24)  # 224-14  384-24
+            self.fc = nn.Linear(self.bcs[-1] * self.block_constructor.extention, fc_num_classes)
+
+    def forward(self, x):
+        # stem:conv+bn+relu+maxpool
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        stem_out = self.maxpool(out)
+
+        # Resnet 3 stages
+        stage1_out = self.layer1(stem_out)
+        stage2_out = self.layer2(stage1_out)
+        stage3_out = self.layer3(stage2_out)
+
+        if self.fc_num_classes is not None:
+            fc_out = self.avgpool(stage3_out)
+            fc_out = torch.flatten(fc_out, 1)
+            fc_out = self.fc(fc_out)
+
+        if self.feature_idx == 'stages':
+            if self.fc_num_classes is not None:
+                return stage1_out, stage2_out, stage3_out, fc_out
+            else:
+                return stage1_out, stage2_out, stage3_out
+        elif self.feature_idx == 'features':
+            if self.fc_num_classes is not None:
+                return stem_out, stage1_out, stage2_out, stage3_out, fc_out
+            else:
+                return stem_out, stage1_out, stage2_out, stage3_out
+        else:  # self.feature_idx is None
+            if self.fc_num_classes is not None:
+                return fc_out
+            else:
+                return stage3_out
+
+    def make_stage_layer(self, block_constractor, midplane, block_num, stride=1):
+        """
+        block:
+        midplane:
+        block_num:
+        stride:
+        """
+
+        block_list = []
+
+        outplane = midplane * block_constractor.extention  # extention
+
+        if stride != 1 or self.inplane != outplane:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplane, outplane, stride=stride, kernel_size=1, bias=False),
+                nn.BatchNorm2d(midplane * block_constractor.extention)
+            )
+        else:
+            downsample = None
+
+        # Conv Block
+        conv_block = block_constractor(self.inplane, midplane, stride=stride, downsample=downsample)
+        block_list.append(conv_block)
+
+        self.inplane = outplane
+
+        # Identity Block
+        for i in range(1, block_num):
+            block_list.append(block_constractor(self.inplane, midplane, stride=1, downsample=None))
+
+        return nn.Sequential(*block_list)
+
+
+def Hybrid_a(backbone, img_size=224, patch_size=1, in_chans=3, num_classes=1000, embed_dim=768, depth=8,
+             num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, drop_rate=0., attn_drop_rate=0.,
+             drop_path_rate=0., norm_layer=None, act_layer=None):
+    # directly stack CNNs and Transformer blocks
+    embed_layer = partial(Transformer_blocks.Hybrid_feature_map_Embed, backbone=backbone)
+
+    Hybrid_model = Transformer_blocks.VisionTransformer(img_size, patch_size, in_chans, num_classes, embed_dim, depth,
+                                                        num_heads, mlp_ratio, qkv_bias, representation_size,
+                                                        drop_rate, attn_drop_rate, drop_path_rate, embed_layer,
+                                                        norm_layer, act_layer)
+
+    return Hybrid_model
+
+
+def create_model(model_idx, edge_size, pretrained=True, num_classes=2, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., use_cls_token=True, use_pos_embedding=True, use_att_module='SimAM'):
+    """
+    get one of MSHT models
+
+    :param model_idx: the model we are going to use. by the format of Model_size_other_info
+    :param edge_size: the input edge size of the dataloder
+    :param pretrained: The backbone CNN is initiate randomly or by its official Pretrained models
+    :param num_classes: classification required number of your dataset
+
+    :param drop_rate: The dropout layer's probility of proposed models
+    :param attn_drop_rate: The dropout layer(right after the MHSA block or MHGA block)'s probility of proposed models
+    :param drop_path_rate: The probility of stochastic depth
+
+    :param use_cls_token: To use the class token
+    :param use_pos_embedding: To use the positional enbedding
+    :param use_att_module: To use which attention module in the FGD Focus block
+    # use_att_module in ['SimAM', 'CBAM', 'SE']  different attention module we applied in the ablation study
+
+    :return: prepared model
+    """
+
+    if pretrained:
+        from torchvision import models
+        backbone_weights = models.resnet50(pretrained=True).state_dict()
+        # True for pretrained Resnet50 model, False will randomly initiate
+    else:
+        backbone_weights = None
+
+    if model_idx[0:11] == 'Hybrid1_224' and edge_size == 224:  # ablation study： no focus depth=8 edge_size == 224
+        backbone = Hybrid_backbone_4(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256, 512],
+                                     identity_layers_setting=[3, 4, 6, 3],
+                                     stage_stride_setting=[1, 2, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx=None)
+
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Hybrid_a(backbone, img_size=edge_size, patch_size=1, in_chans=3, num_classes=num_classes, embed_dim=768,
+                         depth=8, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None,
+                         drop_rate=drop_rate, attn_drop_rate=attn_drop_rate, drop_path_rate=drop_path_rate,
+                         norm_layer=None, act_layer=None)
+
+    elif model_idx[0:11] == 'Hybrid1_384' and edge_size == 384:  # ablation study： no focus depth=8 edge_size == 384
+        backbone = Hybrid_backbone_4(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256, 512],
+                                     identity_layers_setting=[3, 4, 6, 3],
+                                     stage_stride_setting=[1, 2, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx=None)
+
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Hybrid_a(backbone, img_size=edge_size, patch_size=1, in_chans=3, num_classes=num_classes, embed_dim=768,
+                         depth=8, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None,
+                         drop_rate=drop_rate, attn_drop_rate=attn_drop_rate, drop_path_rate=drop_path_rate,
+                         norm_layer=None, act_layer=None)
+
+    elif model_idx[0:11] == 'Hybrid2_224' and edge_size == 224:  # Proposed model ablation study： edge_size==224
+        backbone = Hybrid_backbone_4(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256, 512],
+                                     identity_layers_setting=[3, 4, 6, 3],
+                                     stage_stride_setting=[1, 2, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx='stages')
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Transformer_blocks.Stage_wise_hybrid_Transformer(backbone, num_classes=num_classes,
+                                                                 drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                                                 drop_path_rate=drop_path_rate,
+                                                                 use_cls_token=use_cls_token,
+                                                                 use_pos_embedding=use_pos_embedding,
+                                                                 use_att_module=use_att_module,
+                                                                 stage_size=(56, 28, 14, 7),
+                                                                 stage_dim=[256, 512, 1024, 2048])
+
+    elif model_idx[0:11] == 'Hybrid2_384' and edge_size == 384:  # Proposed model 384 !!!
+        backbone = Hybrid_backbone_4(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256, 512],
+                                     identity_layers_setting=[3, 4, 6, 3],
+                                     stage_stride_setting=[1, 2, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx='stages')
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Transformer_blocks.Stage_wise_hybrid_Transformer(backbone, num_classes=num_classes,
+                                                                 drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                                                 drop_path_rate=drop_path_rate,
+                                                                 use_cls_token=use_cls_token,
+                                                                 use_pos_embedding=use_pos_embedding,
+                                                                 use_att_module=use_att_module,
+                                                                 stage_size=(96, 48, 24, 12),
+                                                                 stage_dim=[256, 512, 1024, 2048])
+
+    elif model_idx[0:11] == 'Hybrid3_224' and edge_size == 224:  # Proposed model ablation study： edge_size==224
+        backbone = Hybrid_backbone_3(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256],
+                                     identity_layers_setting=[3, 4, 6],
+                                     stage_stride_setting=[1, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx='stages')
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Transformer_blocks.Stage_wise_hybrid_Transformer(backbone, num_classes=num_classes,
+                                                                 drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                                                 drop_path_rate=drop_path_rate,
+                                                                 use_cls_token=use_cls_token,
+                                                                 use_pos_embedding=use_pos_embedding,
+                                                                 use_att_module=use_att_module,
+                                                                 stage_size=(56, 28, 14),
+                                                                 stage_dim=[256, 512, 1024])
+
+    elif model_idx[0:11] == 'Hybrid3_384' and edge_size == 384:  # Proposed model 384 !!!
+        backbone = Hybrid_backbone_3(block_constructor=Bottleneck_block_constructor,
+                                     bottleneck_channels_setting=[64, 128, 256],
+                                     identity_layers_setting=[3, 4, 6],
+                                     stage_stride_setting=[1, 2, 2],
+                                     fc_num_classes=None,
+                                     feature_idx='stages')
+        if pretrained:
+            try:
+                backbone.load_state_dict(backbone_weights, False)
+            except:
+                print("backbone not loaded")
+            else:
+                print("backbone loaded")
+
+        model = Transformer_blocks.Stage_wise_hybrid_Transformer(backbone, num_classes=num_classes,
+                                                                 drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                                                 drop_path_rate=drop_path_rate,
+                                                                 use_cls_token=use_cls_token,
+                                                                 use_pos_embedding=use_pos_embedding,
+                                                                 use_att_module=use_att_module,
+                                                                 stage_size=(96, 48, 24),
+                                                                 stage_dim=[256, 512, 1024])
+
+    else:
+        print('not a valid hybrid model')
+        return -1
+
+    return model

PuzzleTuning/Backbone/Transformer_blocks.py

@@ -0,0 +1,1631 @@
+"""
+Transformer blocks script  ver： OCT 28th 15：00
+
+bug fix: 'Cross-attn' name is used in MHGA for compareability
+
+by the authors, check our github page:
+https://github.com/sagizty/Multi-Stage-Hybrid-Transformer
+
+based on：timm
+https://www.freeaihub.com/post/94067.html
+
+"""
+
+import math
+import logging
+from functools import partial
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+from timm.models.layers import StdConv2dSame, DropPath, to_2tuple, trunc_normal_
+
+from .attention_modules import simam_module, cbam_module, se_module
+
+
+class FFN(nn.Module):  # Mlp from timm
+    """
+    FFN (from timm)
+
+    :param in_features:
+    :param hidden_features:
+    :param out_features:
+    :param act_layer:
+    :param drop:
+    """
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+
+        return x
+
+
+class Attention(nn.Module):  # qkv Transform + MSA(MHSA) (Attention from timm)
+    """
+    qkv Transform + MSA(MHSA) (from timm)
+
+    # input  x.shape = batch, patch_number, patch_dim
+    # output  x.shape = batch, patch_number, patch_dim
+
+    :param dim: dim=CNN feature dim, because the patch size is 1x1
+    :param num_heads:
+    :param qkv_bias:
+    :param qk_scale: by default head_dim ** -0.5  (squre root)
+    :param attn_drop: dropout rate after MHSA
+    :param proj_drop:
+
+    """
+
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        # input x.shape = batch, patch_number, patch_dim
+        batch, patch_number, patch_dim = x.shape
+
+        # mlp transform + head split [N, P, D] -> [N, P, 3D] -> [N, P, 3, H, D/H] -> [3, N, H, P, D/H]
+        qkv = self.qkv(x).reshape(batch, patch_number, 3, self.num_heads, patch_dim //
+                                  self.num_heads).permute(2, 0, 3, 1, 4)
+        # 3 [N, H, P, D/H]
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        # [N, H, P, D/H] -> [N, H, P, D/H]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+
+        attn = self.attn_drop(attn)  # Dropout
+
+        # head fusion [N, H, P, D/H] -> [N, P, H, D/H] -> [N, P, D]
+        x = (attn @ v).transpose(1, 2).reshape(batch, patch_number, patch_dim)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)  # mlp
+
+        # output x.shape = batch, patch_number, patch_dim
+        return x
+
+
+class Encoder_Block(nn.Module):  # teansformer Block from timm
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        """
+        # input x.shape = batch, patch_number, patch_dim
+        # output x.shape = batch, patch_number, patch_dim
+
+        :param dim: dim
+        :param num_heads:
+        :param mlp_ratio: FFN
+        :param qkv_bias:
+        :param qk_scale: by default head_dim ** -0.5  (squre root)
+        :param drop:
+        :param attn_drop: dropout rate after Attention
+        :param drop_path: dropout rate after sd
+        :param act_layer: FFN act
+        :param norm_layer: Pre Norm
+        """
+        super().__init__()
+        # Pre Norm
+        self.norm1 = norm_layer(dim)  # Transformer used the nn.LayerNorm
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+                              proj_drop=drop)
+        # NOTE from timm: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()  # stochastic depth
+
+        # Add & Norm
+        self.norm2 = norm_layer(dim)
+
+        # FFN
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = FFN(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class Guided_Attention(nn.Module):  # q1 k1 v0 Transform + MSA(MHSA) (based on timm Attention)
+    """
+    notice the q abd k is guided information from Focus module
+    qkv Transform + MSA(MHSA) (from timm)
+
+    # 3 input of x.shape = batch, patch_number, patch_dim
+    # 1 output of x.shape = batch, patch_number, patch_dim
+
+    :param dim: dim = CNN feature dim, because the patch size is 1x1
+    :param num_heads:
+    :param qkv_bias:
+    :param qk_scale: by default head_dim ** -0.5  (squre root)
+    :param attn_drop:
+    :param proj_drop:
+
+    """
+
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qT = nn.Linear(dim, dim, bias=qkv_bias)
+        self.kT = nn.Linear(dim, dim, bias=qkv_bias)
+        self.vT = nn.Linear(dim, dim, bias=qkv_bias)
+
+        self.attn_drop = nn.Dropout(attn_drop)
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, q_encoder, k_encoder, v_input):
+        # 3 input of x.shape = batch, patch_number, patch_dim
+        batch, patch_number, patch_dim = v_input.shape
+
+        q = self.qT(q_encoder).reshape(batch, patch_number, 1, self.num_heads,
+                                       patch_dim // self.num_heads).permute(2, 0, 3, 1, 4)
+        k = self.kT(k_encoder).reshape(batch, patch_number, 1, self.num_heads,
+                                       patch_dim // self.num_heads).permute(2, 0, 3, 1, 4)
+        v = self.vT(v_input).reshape(batch, patch_number, 1, self.num_heads,
+                                     patch_dim // self.num_heads).permute(2, 0, 3, 1, 4)
+        q = q[0]
+        k = k[0]
+        v = v[0]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+
+        attn = self.attn_drop(attn)  # Dropout
+
+        x = (attn @ v).transpose(1, 2).reshape(batch, patch_number, patch_dim)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)  # mlp Dropout
+
+        # output of x.shape = batch, patch_number, patch_dim
+        return x
+
+
+class Decoder_Block(nn.Module):
+    # FGD Decoder (Transformer encoder + Guided Attention block block)
+    def __init__(self, dim, num_heads=8, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        """
+        # input x.shape = batch, patch_number, patch_dim
+        # output x.shape = batch, patch_number, patch_dim
+
+        :param dim: dim=CNN feature dim, because the patch size is 1x1
+        :param num_heads: multi-head
+        :param mlp_ratio: FFN expand ratio
+        :param qkv_bias: qkv MLP bias
+        :param qk_scale: by default head_dim ** -0.5  (squre root)
+        :param drop: the MLP after MHSA equipt a dropout rate
+        :param attn_drop: dropout rate after attention block
+        :param drop_path: dropout rate for stochastic depth
+        :param act_layer: FFN act
+        :param norm_layer: Pre Norm strategy with norm layer
+        """
+        super().__init__()
+        # Pre Norm
+        self.norm0 = norm_layer(dim)  # nn.LayerNorm
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop,
+                              proj_drop=drop)
+        # stochastic depth
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        # Pre Norm
+        self.norm1 = norm_layer(dim)
+
+        # FFN1
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.FFN1 = FFN(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        # Guided_Attention
+        self.Cross_attn = Guided_Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                           attn_drop=attn_drop, proj_drop=drop)
+
+        # Add & Norm
+        self.norm2 = norm_layer(dim)
+        # FFN2
+        self.FFN2 = FFN(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        # Add & Norm
+        self.norm3 = norm_layer(dim)
+
+    def forward(self, q_encoder, k_encoder, v_input):
+        v_self = v_input + self.drop_path(self.attn(self.norm0(v_input)))
+
+        v_self = v_self + self.drop_path(self.FFN1(self.norm1(v_self)))
+
+        # norm layer for v only, the normalization of q and k is inside FGD Focus block
+        v_self = v_self + self.drop_path(self.Cross_attn(q_encoder, k_encoder, self.norm2(v_self)))
+
+        v_self = v_self + self.drop_path(self.FFN2(self.norm3(v_self)))
+
+        return v_self
+
+
+'''
+# testing example
+
+model=Decoder_Block(dim=768)
+k = torch.randn(7, 49, 768)
+q = torch.randn(7, 49, 768)
+v = torch.randn(7, 49, 768)
+x = model(k,q,v)
+print(x.shape)
+'''
+
+
+# MViT modules
+# from https://github.com/facebookresearch/SlowFast/slowfast/models/attention.py
+def attention_pool(tensor, pool, thw_shape, has_cls_embed=True, norm=None):
+    """
+    attention pooling constructor
+
+    input:
+    tensor of (B, Head, N, C) or (B, N, C)
+    thw_shape: T, H, W  对应CNN的特征图形状（2D形状）T is video frams
+
+    numpy.prob(T, H, W) == N(Num_patches) - 1 (cls token if it is there)
+
+    output:
+    tensor of (B, Head, N_O, C) or (B, N_O, C)
+    thw_shape: T_O, H_O, W_O
+
+    :param tensor: input feature patches
+    :param pool: pooling/conv layer
+    :param thw_shape: reconstruction feature map shape
+    :param has_cls_embed: if cls token is used
+    :param norm:  norm layer
+
+    """
+    if pool is None:  # no pool
+        return tensor, thw_shape
+
+    tensor_dim = tensor.ndim
+
+    # fix dim: [B, Head, N, C]
+    # N is Num_patches in Transformer modeling
+
+    if tensor_dim == 4:
+        pass
+    elif tensor_dim == 3:  # [B, N, C] -> [B, Head(1), N, C]
+        tensor = tensor.unsqueeze(1)
+    else:
+        raise NotImplementedError(f"Unsupported input dimension {tensor.shape}")
+
+    if has_cls_embed:
+        cls_tok, tensor = tensor[:, :, :1, :], tensor[:, :, 1:, :]
+
+    B, Head, N, C = tensor.shape
+    T, H, W = thw_shape  # numpy.prob(T, H, W) == N(Num_patches) - 1 (cls token if it is there)
+
+    # [B, Head, N, C] -> [B * Head, T, H, W, C] -> [B * Head, C, T, H, W]
+    tensor = (tensor.reshape(B * Head, T, H, W, C).permute(0, 4, 1, 2, 3).contiguous())
+    # use tensor.contiguous() to matain its memory location
+
+    # [B * Head, C, T, H, W] -> [B * Head, C, T_O, H_O, W_O]
+    tensor = pool(tensor)  # 3D Pooling/ 3D Conv
+
+    # output T, H, W
+    thw_shape = [tensor.shape[2], tensor.shape[3], tensor.shape[4]]
+    # output Num_patches: numpy.prob(T, H, W)
+    N_pooled = tensor.shape[2] * tensor.shape[3] * tensor.shape[4]
+
+    # [B * Head, C, T_O, H_O, W_O] -> [B, Head, C, N_O(T_O*H_O*W_O)] -> [B, Head, N_O, C]
+    tensor = tensor.reshape(B, Head, C, N_pooled).transpose(2, 3)
+
+    if has_cls_embed:
+        # [B, Head, N_O, C] -> [B, Head, N_O+1(cls token), C]
+        tensor = torch.cat((cls_tok, tensor), dim=2)
+
+    # norm
+    if norm is not None:
+        tensor = norm(tensor)
+
+    # Assert tensor_dim in [3, 4]
+    if tensor_dim == 4:  # [B, Head, N_O, C] multi-head
+        pass
+    else:  # tensor_dim == 3: this is a single Head
+        tensor = tensor.squeeze(1)  # [B, N_O, C]
+
+    return tensor, thw_shape
+
+
+'''
+# case 1 single-head no pooling scale
+x = torch.randn(1, 197, 768)
+thw_shape = [1, 14, 14]
+pool = nn.MaxPool3d((1, 1, 1), (1, 1, 1), (0, 0, 0), ceil_mode=False)
+y, thw = attention_pool(x, pool, thw_shape)
+
+print(y.shape)  # torch.Size([1, 197, 768])
+print(thw)  # [1, 14, 14]
+
+
+# case 2  multi-head no pooling scale
+x = torch.randn(1, 8, 197, 96)  # [B, Head, N_O, C] multi-head
+thw_shape = [1, 14, 14]
+pool = nn.MaxPool3d((1, 1, 1), (1, 1, 1), (0, 0, 0), ceil_mode=False)
+y, thw = attention_pool(x, pool, thw_shape)
+
+print(y.shape)  # torch.Size([1, 8, 197, 96])
+print(thw)  # [1, 14, 14]
+
+
+# case 3 pooling scale
+x = torch.randn(1, 197, 768)
+thw_shape = [1, 14, 14]
+pool = nn.MaxPool3d((1, 2, 2), (1, 2, 2), (0, 0, 0), ceil_mode=False)
+y, thw = attention_pool(x, pool, thw_shape)
+
+print(y.shape)  # torch.Size([1, 50, 768])
+print(thw)  # [1, 7, 7]
+
+
+# case 4 multi-head pooling scale
+x = torch.randn(1, 8, 197, 96)  # [B, Head, N_O, C] multi-head
+thw_shape = [1, 14, 14]
+pool = nn.MaxPool3d((1, 2, 2), (1, 2, 2), (0, 0, 0), ceil_mode=False)
+y, thw = attention_pool(x, pool, thw_shape)
+
+print(y.shape)  # torch.Size([1, 8, 50, 96])
+print(thw)  # [1, 7, 7]
+'''
+
+
+class MultiScaleAttention(nn.Module):  # Attention module
+    """
+    Attention module constructor
+
+        input:
+        tensor of (B, N, C)
+        thw_shape: T, H, W  对应CNN的特征图形状（2D形状）T is video frams
+
+        numpy.prob(T, H, W) == N(Num_patches) - 1 (cls token if it is there)
+
+        output:
+        tensor of (B, N_O, C)
+        thw_shape: T_O, H_O, W_O
+
+        :param dim: Transformer feature dim
+        :param num_heads: Transformer heads
+        :param qkv_bias: projecting bias
+        :param drop_rate: dropout rate after attention calculation and mlp
+
+        :param kernel_q: pooling kernal size for q
+        :param kernel_kv: pooling kernal size for k and v
+        :param stride_q: pooling kernal stride for q
+        :param stride_kv: pooling kernal stride for k and v
+
+        :param norm_layer:  norm layer
+        :param has_cls_embed: if cls token is used
+        :param mode: mode for attention pooling(downsampling) Options include `conv`, `avg`, and `max`.
+        :param pool_first: process pooling(downsampling) before liner projecting
+
+    """
+
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=False,
+            drop_rate=0.0,
+            kernel_q=(1, 1, 1),
+            kernel_kv=(1, 1, 1),
+            stride_q=(1, 1, 1),
+            stride_kv=(1, 1, 1),
+            norm_layer=nn.LayerNorm,
+            has_cls_embed=True,
+            # Options include `conv`, `avg`, and `max`.
+            mode="conv",
+            # If True, perform pool before projection.
+            pool_first=False,
+    ):
+        super().__init__()
+
+        self.pool_first = pool_first
+        self.drop_rate = drop_rate
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5  # squre root
+        self.has_cls_embed = has_cls_embed
+
+        padding_q = [int(q // 2) for q in kernel_q]  # 以半个kernal size进行padding，向下取整
+        padding_kv = [int(kv // 2) for kv in kernel_kv]
+
+        # projecting mlp
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+
+        if drop_rate > 0.0:
+            self.proj_drop = nn.Dropout(drop_rate)
+
+        # Skip pooling with kernel and stride size of (1, 1, 1).
+        if np.prod(kernel_q) == 1 and np.prod(stride_q) == 1:
+            kernel_q = ()  # clear
+        if np.prod(kernel_kv) == 1 and np.prod(stride_kv) == 1:
+            kernel_kv = ()
+
+        if mode in ("avg", "max"):  # use nn.MaxPool3d or nn.AvgPool3d
+            pool_op = nn.MaxPool3d if mode == "max" else nn.AvgPool3d
+            self.pool_q = (
+                pool_op(kernel_q, stride_q, padding_q, ceil_mode=False)
+                if len(kernel_q) > 0
+                else None  # Skip pooling if kernel is cleared
+            )
+            self.pool_k = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.pool_v = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+
+        elif mode == "conv":  # use nn.Conv3d with depth wise conv and fixed channel setting
+            self.pool_q = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_q,
+                    stride=stride_q,
+                    padding=padding_q,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_q) > 0
+                else None
+            )
+            self.norm_q = norm_layer(head_dim) if len(kernel_q) > 0 else None
+
+            self.pool_k = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_kv,
+                    stride=stride_kv,
+                    padding=padding_kv,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.norm_k = norm_layer(head_dim) if len(kernel_kv) > 0 else None
+
+            self.pool_v = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_kv,
+                    stride=stride_kv,
+                    padding=padding_kv,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.norm_v = norm_layer(head_dim) if len(kernel_kv) > 0 else None
+        else:
+            raise NotImplementedError(f"Unsupported model {mode}")
+
+    def forward(self, x, thw_shape):
+        """
+        x: Transformer feature patches
+        thw_shape: reconstruction feature map shape
+        """
+
+        B, N, C = x.shape
+
+        # step 1: duplicate projecting + head split: [B, N, C] -> [B, H, N, C/H]
+
+        if self.pool_first:  # step a.1 embedding
+            # head split [B, N, C] -> [B, N, H, C/H] -> [B, H, N, C/H]
+            x = x.reshape(B, N, self.num_heads, C // self.num_heads).permute(
+                0, 2, 1, 3
+            )
+            q = k = v = x
+
+        else:  # step b.1 projecting first
+            # mlp transform + head split: [B, N, C] -> [B, N, H, C/H] -> [B, H, N, C/H]
+            # todo 这里我觉得可能共享mlp映射更好，能有更好的交互，但是分离mlp更节约计算量
+            q = k = v = x
+            q = (
+                self.q(q)
+                    .reshape(B, N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+            k = (
+                self.k(k)
+                    .reshape(B, N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+            v = (
+                self.v(v)
+                    .reshape(B, N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+
+        # step 2: calculate attention_pool feature sequence and its shape
+        # [B, H, N0, C/H] -> [B, H, N1, C/H]
+        q, q_shape = attention_pool(
+            q,
+            self.pool_q,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_q if hasattr(self, "norm_q") else None,
+        )
+        k, k_shape = attention_pool(
+            k,
+            self.pool_k,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_k if hasattr(self, "norm_k") else None,
+        )
+        v, v_shape = attention_pool(
+            v,
+            self.pool_v,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_v if hasattr(self, "norm_v") else None,
+        )
+
+        if self.pool_first:  # step a.3 MLP projecting
+            # calculate patch number, q_N, k_N, v_N
+            q_N = (
+                np.prod(q_shape) + 1
+                if self.has_cls_embed
+                else np.prod(q_shape)
+            )
+            k_N = (
+                np.prod(k_shape) + 1
+                if self.has_cls_embed
+                else np.prod(k_shape)
+            )
+            v_N = (
+                np.prod(v_shape) + 1
+                if self.has_cls_embed
+                else np.prod(v_shape)
+            )
+
+            # [B, H, N1, C/H] -> [B, N1, H, C/H] -> [B, N1, C] -> MLP
+            # -> [B, N1, C] -> [B, N1, H, C/H] -> [B, H, N1, C/H]
+            q = q.permute(0, 2, 1, 3).reshape(B, q_N, C)
+            q = (
+                self.q(q)
+                    .reshape(B, q_N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+
+            v = v.permute(0, 2, 1, 3).reshape(B, v_N, C)
+            v = (
+                self.v(v)
+                    .reshape(B, v_N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+
+            k = k.permute(0, 2, 1, 3).reshape(B, k_N, C)
+            k = (
+                self.k(k)
+                    .reshape(B, k_N, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+            )
+
+        # step 3: attention calculation
+        # multi-head self attention [B, H, N1, C/H] -> [B, H, N1, C/H]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+
+        # head squeeze [B, H, N1, C/H] -> [B, N1, H, C/H] -> [B, N1, C]
+        N = q.shape[2]
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+
+        # step 4: mlp stablization and dropout [B, N1, C] -> [B, N1, C]
+        x = self.proj(x)
+        if self.drop_rate > 0.0:
+            x = self.proj_drop(x)
+
+        return x, q_shape
+
+
+'''
+# case 1
+model = MultiScaleAttention(768)
+x = torch.randn(1, 197, 768)
+y, thw = model(x, [1, 14, 14])
+print(y.shape)
+
+
+# case 2
+kernel_q = (1, 2, 2)
+kernel_kv = (1, 2, 2)
+stride_q = (1, 2, 2)
+stride_kv = (1, 2, 2)
+# MultiScaleAttention 中设计以半个kernal size进行padding，向下取整
+
+model = MultiScaleAttention(768, kernel_q=kernel_q, kernel_kv=kernel_kv, stride_q=stride_q, stride_kv=stride_kv)
+x = torch.randn(1, 197, 768)
+y, thw = model(x, [1, 14, 14])
+
+print(y.shape)  # 输出torch.Size([1, 65, 768])：不padding是7*7 由于padding变成8*8， 之后加上cls token
+'''
+
+
+class MultiScaleBlock(nn.Module):  # MViT Encoder
+    """
+    Attention module constructor
+
+        input:
+        tensor of (B, N, C)
+        thw_shape: T, H, W  对应CNN的特征图形状（2D形状）T is video frams
+
+        numpy.prob(T, H, W) == N(Num_patches) - 1 (cls token if it is there)
+
+        output:
+        tensor of (B, N_O, C)
+        thw_shape: T_O, H_O, W_O
+
+        :param dim: Transformer feature dim
+        :param dim_out:
+
+        :param num_heads: Transformer heads
+        :param mlp_ratio: FFN hidden expansion
+        :param qkv_bias: projecting bias
+        :param drop_rate: dropout rate after attention calculation and mlp
+        :param drop_path: dropout rate for SD
+        :param act_layer: FFN act
+        :param norm_layer: Pre Norm
+
+        :param up_rate:
+        :param kernel_q: pooling kernal size for q
+        :param kernel_kv: pooling kernal size for k and v
+        :param stride_q: pooling kernal stride for q
+        :param stride_kv: pooling kernal stride for k and v
+
+        :param has_cls_embed: if cls token is used
+        :param mode: mode for attention pooling(downsampling) Options include `conv`, `avg`, and `max`.
+        :param pool_first: process pooling(downsampling) before liner projecting
+
+    """
+
+    def __init__(
+            self,
+            dim,
+            dim_out,
+            num_heads=8,
+            mlp_ratio=4.0,
+            qkv_bias=False,
+            drop_rate=0.0,
+            drop_path=0.0,
+            act_layer=nn.GELU,
+            norm_layer=nn.LayerNorm,
+            up_rate=None,
+            kernel_q=(1, 1, 1),
+            kernel_kv=(1, 1, 1),
+            stride_q=(1, 1, 1),
+            stride_kv=(1, 1, 1),
+            has_cls_embed=True,
+            mode="conv",
+            pool_first=False,
+    ):
+        super().__init__()
+
+        self.has_cls_embed = has_cls_embed
+
+        # step 1: Attention projecting
+        self.dim = dim
+        self.dim_out = dim_out
+        self.norm1 = norm_layer(dim)  # pre-norm
+
+        self.attn = MultiScaleAttention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            drop_rate=drop_rate,
+            kernel_q=kernel_q,
+            kernel_kv=kernel_kv,
+            stride_q=stride_q,
+            stride_kv=stride_kv,
+            norm_layer=nn.LayerNorm,
+            has_cls_embed=self.has_cls_embed,
+            mode=mode,
+            pool_first=pool_first,
+            )
+
+        self.drop_path = (DropPath(drop_path) if drop_path > 0.0 else nn.Identity())
+
+        # residual connection for Attention projecting
+        kernel_skip = kernel_q  # fixme ori: [s + 1 if s > 1 else s for s in stride_q]
+        stride_skip = stride_q
+        padding_skip = [int(skip // 2) for skip in kernel_skip]  # 以半个kernal size进行padding，向下取整
+
+        self.pool_skip = (
+            nn.MaxPool3d(kernel_skip, stride_skip, padding_skip, ceil_mode=False)
+            if len(kernel_skip) > 0
+            else None)
+
+        self.norm2 = norm_layer(dim)  # pre-norm
+
+        # step 2: FFN projecting
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        # here use FFN to encode feature into abstractive information in the dimension
+        # TODO: check the use case for up_rate, and merge the following lines
+        if up_rate is not None and up_rate > 1:
+            mlp_dim_out = dim * up_rate
+        else:
+            mlp_dim_out = dim_out
+
+        self.mlp = FFN(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            out_features=mlp_dim_out,
+            act_layer=act_layer,
+            drop=drop_rate,
+        )
+
+        # residual connection for FFN projecting
+        if dim != dim_out:
+            self.proj = nn.Linear(dim, dim_out)
+
+    def forward(self, x, thw_shape):
+        # step 1: Attention projecting
+        x_block, thw_shape_new = self.attn(self.norm1(x), thw_shape)
+        # residual connection for Attention projecting
+        x_res, _ = attention_pool(x, self.pool_skip, thw_shape, has_cls_embed=self.has_cls_embed)
+        x = x_res + self.drop_path(x_block)
+
+        # step 2: FFN projecting
+        x_norm = self.norm2(x)
+        x_mlp = self.mlp(x_norm)
+        # residual connection for FFN projecting
+        if self.dim != self.dim_out:
+            x = self.proj(x_norm)
+        x = x + self.drop_path(x_mlp)
+
+        return x, thw_shape_new
+
+
+'''
+# case 1
+model = MultiScaleBlock(768,1024)
+x = torch.randn(1, 197, 768)
+y, thw = model(x, [1, 14, 14])
+print(y.shape)  # torch.Size([1, 197, 1024])
+
+
+# case 2
+kernel_q = (1, 2, 2)
+kernel_kv = (1, 2, 2)
+stride_q = (1, 2, 2)
+stride_kv = (1, 2, 2)
+# MultiScaleAttention 中设计以半个kernal size进行padding，向下取整
+
+model = MultiScaleBlock(768, 1024, kernel_q=kernel_q, kernel_kv=kernel_kv, stride_q=stride_q, stride_kv=stride_kv)
+x = torch.randn(1, 197, 768)
+y, thw = model(x, [1, 14, 14])
+
+print(y.shape)  # 输出torch.Size([1, 65, 1024])：不padding是7*7 由于padding变成8*8， 之后加上cls token
+'''
+
+
+class PatchEmbed(nn.Module):  # PatchEmbed from timm
+    """
+    Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+
+        # x: (B, 14*14, 768)
+        return x
+
+
+class Hybrid_feature_map_Embed(nn.Module):  # HybridEmbed from timm
+    """
+    CNN Feature Map Embedding, required backbone which is just for referance here
+    Extract feature map from CNN, flatten, project to embedding dim.
+
+    # input x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+    # output x.shape = batch, patch_number, patch_dim
+    """
+
+    def __init__(self, backbone, img_size=224, patch_size=1, feature_size=None, feature_dim=None,
+                 in_chans=3, embed_dim=768):
+        super().__init__()
+
+        assert isinstance(backbone, nn.Module)
+
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.backbone = backbone
+
+        if feature_size is None or feature_dim is None:  # backbone output feature_size
+            with torch.no_grad():
+                # NOTE Most reliable way of determining output dims is to run forward pass
+                training = backbone.training
+                if training:
+                    backbone.eval()
+                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))
+                if isinstance(o, (list, tuple)):
+                    o = o[-1]  # last feature if backbone outputs list/tuple of features
+                feature_size = o.shape[-2:]
+                feature_dim = o.shape[1]
+                backbone.train(training)
+        else:
+            feature_size = to_2tuple(feature_size)
+            '''
+            if hasattr(self.backbone, 'feature_info'):
+                feature_dim = self.backbone.feature_info.channels()[-1]
+            else:
+                feature_dim = self.backbone.num_features
+            '''
+
+        assert feature_size[0] % patch_size[0] == 0 and feature_size[1] % patch_size[1] == 0
+
+        self.grid_size = (feature_size[0] // patch_size[0], feature_size[1] // patch_size[1])  # patchlize
+
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        self.proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                              kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        x = self.proj(x).flatten(2).transpose(1, 2)  # shape = ( )
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output: x.shape = batch, patch_number, patch_dim
+        return x
+
+
+class Last_feature_map_Embed(nn.Module):
+    """
+    use this block to connect last CNN stage to the first Transformer block
+    Extract feature map from CNN, flatten, project to embedding dim.
+
+    # input x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+    # output x.shape = batch, patch_number, patch_dim
+    """
+
+    def __init__(self, patch_size=1, feature_size=(7, 7), feature_dim=2048, embed_dim=768,
+                 Attention_module=None):
+        super().__init__()
+
+        # Attention module
+        if Attention_module is not None:
+            if Attention_module == 'SimAM':
+                self.Attention_module = simam_module(e_lambda=1e-4)
+            elif Attention_module == 'CBAM':
+                self.Attention_module = cbam_module(gate_channels=feature_dim)
+            elif Attention_module == 'SE':
+                self.Attention_module = se_module(channel=feature_dim)
+        else:
+            self.Attention_module = None
+
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+
+        feature_size = to_2tuple(feature_size)
+
+        # feature map should be matching the size
+        assert feature_size[0] % self.patch_size[0] == 0 and feature_size[1] % self.patch_size[1] == 0
+
+        self.grid_size = (feature_size[0] // self.patch_size[0], feature_size[1] // self.patch_size[1])  # patch
+
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        # use the conv to split the patch by the following design:
+        self.proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                              kernel_size=self.patch_size, stride=self.patch_size)
+
+    def forward(self, x):
+        if self.Attention_module is not None:
+            x = self.Attention_module(x)
+
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output 格式 x.shape = batch, patch_number, patch_dim
+        return x
+
+
+class Focus_Embed(nn.Module):  # Attention guided module for hybridzing the early stages CNN feature
+    """
+    FGD Focus module
+    Extract feature map from CNN, flatten, project to embedding dim. and use them as attention guidance
+
+    input: x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+
+    Firstly, an attention block will be used to stable the feature projecting process
+
+    Secondly, for each feature map，the focus will be 2 path: gaze and glance
+    in gaze path Max pool will be applied to get prominent information
+    in glance path Avg pool will be applied to get general information
+
+    after the dual pooling path 2 seperate CNNs will be used to project the dimension
+    Finally, flattern and transpose will be applied
+
+    output 2 attention guidance: gaze, glance
+    x.shape = batch, patch_number, patch_dim
+
+
+    ref:
+    ResNet50's feature map from different stages (edge size of 224)
+    stage 1 output feature map: torch.Size([b, 256, 56, 56])
+    stage 2 output feature map: torch.Size([b, 512, 28, 28])
+    stage 3 output feature map: torch.Size([b, 1024, 14, 14])
+    stage 4 output feature map: torch.Size([b, 2048, 7, 7])
+    """
+
+    def __init__(self, patch_size=1, target_feature_size=(7, 7), feature_size=(56, 56), feature_dim=256, embed_dim=768,
+                 Attention_module=None, norm_layer=nn.LayerNorm):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        feature_size = to_2tuple(feature_size)  # patch size of the current feature map
+
+        target_feature_size = to_2tuple(target_feature_size)  # patch size of the last feature map
+
+        # cheak feature map can be patchlize to target_feature_size
+        assert feature_size[0] % target_feature_size[0] == 0 and feature_size[1] % target_feature_size[1] == 0
+
+        # cheak target_feature map can be patchlize to patch
+        assert target_feature_size[0] % patch_size[0] == 0 and target_feature_size[1] % patch_size[1] == 0
+
+        # Attention block
+        if Attention_module is not None:
+            if Attention_module == 'SimAM':
+                self.Attention_module = simam_module(e_lambda=1e-4)
+            elif Attention_module == 'CBAM':
+                self.Attention_module = cbam_module(gate_channels=feature_dim)
+            elif Attention_module == 'SE':
+                self.Attention_module = se_module(channel=feature_dim)
+        else:
+            self.Attention_module = None
+
+        # split focus ROI
+        self.focus_size = (feature_size[0] // target_feature_size[0], feature_size[1] // target_feature_size[1])
+        self.num_focus = self.focus_size[0] * self.focus_size[1]
+        # by kernel_size=focus_size, stride=focus_size design
+        # output_size=target_feature_size=7x7 so as to match the minist feature map
+
+        self.gaze = nn.MaxPool2d(self.focus_size, stride=self.focus_size)
+        self.glance = nn.AvgPool2d(self.focus_size, stride=self.focus_size)
+        # x.shape:  batch, feature_dim, target_feature_size[0], target_feature_size[1]
+
+        # split patch
+        self.grid_size = (target_feature_size[0] // patch_size[0], target_feature_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        # use CNN to project dim to patch_dim
+        self.gaze_proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                                   kernel_size=patch_size, stride=patch_size)
+        self.glance_proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                                     kernel_size=patch_size, stride=patch_size)
+
+        self.norm_q = norm_layer(embed_dim)  # Transformer nn.LayerNorm
+        self.norm_k = norm_layer(embed_dim)  # Transformer nn.LayerNorm
+
+    def forward(self, x):
+        if self.Attention_module is not None:
+            x = self.Attention_module(x)
+
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        q = self.norm_q(self.gaze_proj(self.gaze(x)).flatten(2).transpose(1, 2))
+        k = self.norm_k(self.glance_proj(self.glance(x)).flatten(2).transpose(1, 2))
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        gaze/glance(x).shape:  batch, feature_dim, target_feature_size[0], target_feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output x.shape = batch, patch_number, patch_dim
+        return q, k
+
+
+'''
+# test sample
+model = Focus_Embed()
+x = torch.randn(4, 256, 56, 56)
+y1,y2 = model(x)
+print(y1.shape)
+print(y2.shape)
+'''
+
+
+class Focus_SEmbed(nn.Module):  # Attention guided module for hybridzing the early stages CNN feature
+    """
+
+    self focus (q=k)  based on FGD Focus block
+
+    Extract feature map from CNN, flatten, project to embedding dim. and use them as attention guidance
+
+    input: x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+
+    Firstly, an attention block will be used to stable the feature projecting process
+
+    Secondly, for each feature map，the focus will be 1 path: glance
+    in glance path Avg pool will be applied to get general information
+
+    after the pooling process 1 CNN will be used to project the dimension
+    Finally, flattern and transpose will be applied
+
+    output 2 attention guidance: glance, glance
+    x.shape = batch, patch_number, patch_dim
+    """
+
+    def __init__(self, patch_size=1, target_feature_size=(7, 7), feature_size=(56, 56), feature_dim=256, embed_dim=768,
+                 Attention_module=None, norm_layer=nn.LayerNorm):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        feature_size = to_2tuple(feature_size)
+
+        target_feature_size = to_2tuple(target_feature_size)
+
+        assert feature_size[0] % target_feature_size[0] == 0 and feature_size[1] % target_feature_size[1] == 0
+
+        assert target_feature_size[0] % patch_size[0] == 0 and target_feature_size[1] % patch_size[1] == 0
+
+        if Attention_module is not None:
+            if Attention_module == 'SimAM':
+                self.Attention_module = simam_module(e_lambda=1e-4)
+            elif Attention_module == 'CBAM':
+                self.Attention_module = cbam_module(gate_channels=feature_dim)
+            elif Attention_module == 'SE':
+                self.Attention_module = se_module(channel=feature_dim)
+        else:
+            self.Attention_module = None
+
+        self.focus_size = (feature_size[0] // target_feature_size[0], feature_size[1] // target_feature_size[1])
+        self.num_focus = self.focus_size[0] * self.focus_size[1]
+
+        self.gaze = nn.MaxPool2d(self.focus_size, stride=self.focus_size)
+
+        self.grid_size = (target_feature_size[0] // patch_size[0], target_feature_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+        self.proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim, kernel_size=patch_size,
+                              stride=patch_size)
+
+        self.norm_f = norm_layer(embed_dim)
+
+    def forward(self, x):
+        if self.Attention_module is not None:
+            x = self.Attention_module(x)
+
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        q = self.norm_f(self.proj(self.gaze(x)).flatten(2).transpose(1, 2))
+        k = q
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        gaze/glance(x).shape:  batch, feature_dim, target_feature_size[0], target_feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output x.shape = batch, patch_number, patch_dim
+        return q, k
+
+
+class Focus_Aggressive(nn.Module):  # Attention guided module for hybridzing the early stages CNN feature
+    """
+    Aggressive CNN Focus based on FGD Focus block
+
+    Extract feature map from CNN, flatten, project to embedding dim. and use them as attention guidance
+
+    input: x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+
+    Firstly, an attention block will be used to stable the feature projecting process
+
+    Secondly, 2 CNNs will be used to project the dimension
+
+    Finally, flattern and transpose will be applied
+
+    output 2 attention guidance: gaze, glance
+    x.shape = batch, patch_number, patch_dim
+
+    """
+
+    def __init__(self, patch_size=1, target_feature_size=(7, 7), feature_size=(56, 56), feature_dim=256, embed_dim=768,
+                 Attention_module=None, norm_layer=nn.LayerNorm):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)  # patch size of the last feature map
+        feature_size = to_2tuple(feature_size)
+
+        target_feature_size = to_2tuple(target_feature_size)
+
+        assert feature_size[0] % target_feature_size[0] == 0 and feature_size[1] % target_feature_size[1] == 0
+
+        assert target_feature_size[0] % patch_size[0] == 0 and target_feature_size[1] % patch_size[1] == 0
+
+        if Attention_module is not None:
+            if Attention_module == 'SimAM':
+                self.Attention_module = simam_module(e_lambda=1e-4)
+            elif Attention_module == 'CBAM':
+                self.Attention_module = cbam_module(gate_channels=feature_dim)
+            elif Attention_module == 'SE':
+                self.Attention_module = se_module(channel=feature_dim)
+        else:
+            self.Attention_module = None
+
+        self.focus_size = (feature_size[0] // target_feature_size[0], feature_size[1] // target_feature_size[1])
+
+        self.grid_size = (self.focus_size[0] * patch_size[0], self.focus_size[1] * patch_size[1])
+        self.num_patches = (feature_size[0] // self.grid_size[0]) * (feature_size[1] // self.grid_size[1])
+
+        self.gaze_proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                                   kernel_size=self.grid_size, stride=self.grid_size)
+        self.glance_proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                                     kernel_size=self.grid_size, stride=self.grid_size)
+
+        self.norm_q = norm_layer(embed_dim)
+        self.norm_k = norm_layer(embed_dim)
+
+    def forward(self, x):
+        if self.Attention_module is not None:
+            x = self.Attention_module(x)
+
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        q = self.norm_q(self.gaze_proj(x).flatten(2).transpose(1, 2))
+        k = self.norm_k(self.glance_proj(x).flatten(2).transpose(1, 2))
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output x.shape = batch, patch_number, patch_dim
+        return q, k
+
+
+class Focus_SAggressive(nn.Module):  # Attention guided module for hybridzing the early stages CNN feature
+    """
+    Aggressive CNN self Focus
+    Extract feature map from CNN, flatten, project to embedding dim. and use them as attention guidance
+
+    input: x.shape = batch, feature_dim, feature_size[0], feature_size[1]
+
+    Firstly, an attention block will be used to stable the feature projecting process
+
+    Secondly, 1 CNN will be used to project the dimension
+
+    Finally, flattern and transpose will be applied
+
+    output 2 attention guidance: glance, glance
+    x.shape = batch, patch_number, patch_dim
+    """
+
+    def __init__(self, patch_size=1, target_feature_size=(7, 7), feature_size=(56, 56), feature_dim=256, embed_dim=768,
+                 Attention_module=None, norm_layer=nn.LayerNorm):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        feature_size = to_2tuple(feature_size)
+
+        target_feature_size = to_2tuple(target_feature_size)
+
+        assert feature_size[0] % target_feature_size[0] == 0 and feature_size[1] % target_feature_size[1] == 0
+
+        assert target_feature_size[0] % patch_size[0] == 0 and target_feature_size[1] % patch_size[1] == 0
+
+        if Attention_module is not None:
+            if Attention_module == 'SimAM':
+                self.Attention_module = simam_module(e_lambda=1e-4)
+            elif Attention_module == 'CBAM':
+                self.Attention_module = cbam_module(gate_channels=feature_dim)
+            elif Attention_module == 'SE':
+                self.Attention_module = se_module(channel=feature_dim)
+        else:
+            self.Attention_module = None
+
+        self.focus_size = (feature_size[0] // target_feature_size[0], feature_size[1] // target_feature_size[1])
+
+        self.grid_size = (self.focus_size[0] * patch_size[0], self.focus_size[1] * patch_size[1])
+        self.num_patches = (feature_size[0] // self.grid_size[0]) * (feature_size[1] // self.grid_size[1])
+
+        self.proj = nn.Conv2d(in_channels=feature_dim, out_channels=embed_dim,
+                              kernel_size=self.grid_size, stride=self.grid_size)
+
+        self.norm_f = norm_layer(embed_dim)
+
+    def forward(self, x):
+        if self.Attention_module is not None:
+            x = self.Attention_module(x)
+
+        if isinstance(x, (list, tuple)):
+            x = x[-1]  # last feature if backbone outputs list/tuple of features
+
+        q = self.norm_f(self.proj(x).flatten(2).transpose(1, 2))
+        k = q
+        """
+        x.shape:  batch, feature_dim, feature_size[0], feature_size[1]
+        proj(x).shape:  batch, embed_dim, patch_height_num, patch_width_num
+        flatten(2).shape:  batch, embed_dim, patch_num
+        .transpose(1, 2).shape:  batch feature_patch_number feature_patch_dim
+        """
+        # output x.shape = batch, patch_number, patch_dim
+        return q, k
+
+
+class VisionTransformer(nn.Module):  # From timm to review the ViT and ViT_resn5
+    """
+    Vision Transformer
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`
+        - https://arxiv.org/abs/2010.11929
+    Includes distillation token & head support for `DeiT: Data-efficient Image Transformers`
+        - https://arxiv.org/abs/2012.12877
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., embed_layer=PatchEmbed, norm_layer=None, act_layer=None):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            embed_layer (nn.Module): patch embedding layer
+            norm_layer: (nn.Module): normalization layer
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_tokens = 1
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        act_layer = act_layer or nn.GELU
+
+        self.patch_embed = embed_layer(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        self.blocks = nn.Sequential(*[
+            Encoder_Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
+                          attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer)
+            for i in range(depth)])
+
+        self.norm = norm_layer(embed_dim)
+
+        # Representation layer
+        if representation_size:
+            self.num_features = representation_size
+            self.pre_logits = nn.Sequential(OrderedDict([
+                ('fc', nn.Linear(embed_dim, representation_size)),
+                ('act', nn.Tanh())
+            ]))
+        else:
+            self.pre_logits = nn.Identity()
+
+        # Classifier head(s)
+        self.head = nn.Linear(self.num_features, self.num_classes) if self.num_classes > 0 else nn.Identity()
+        self.head_dist = None
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        # print(x.shape,self.pos_embed.shape)
+        cls_token = self.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_token, x), dim=1)
+        x = self.pos_drop(x + self.pos_embed)
+
+        x = self.blocks(x)
+        x = self.norm(x)
+        return self.pre_logits(x[:, 0])  # use cls token for cls head
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+
+class Stage_wise_hybrid_Transformer(nn.Module):
+    """
+    MSHT: Multi Stage Backbone Transformer
+    Stem + 4 ResNet stages（Backbone）is used as backbone
+    then, last feature map patch embedding is used to connect the CNN output to the decoder1 input
+
+    horizonally, 4 ResNet Stage has its feature map connecting to the Focus module
+    which we be use as attention guidance into the FGD decoder
+    """
+
+    def __init__(self, backbone, num_classes=1000, patch_size=1, embed_dim=768, depth=4, num_heads=8, mlp_ratio=4.,
+                 qkv_bias=True, representation_size=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
+                 use_cls_token=True, use_pos_embedding=True, use_att_module='SimAM', stage_size=(56, 28, 14, 7),
+                 stage_dim=(256, 512, 1024, 2048), norm_layer=None, act_layer=None):
+        """
+        Args:
+            backbone (nn.Module): input backbone = stem + 4 ResNet stages
+            num_classes (int): number of classes for classification head
+            patch_size (int, tuple): patch size
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+
+            use_cls_token(bool): classification token
+            use_pos_embedding(bool): use positional embedding
+            use_att_module(str or None): use which attention module in embedding
+
+            stage_size (int, tuple): the stage feature map size of ResNet stages
+            stage_dim (int, tuple): the stage feature map dimension of ResNet stages
+            norm_layer: (nn.Module): normalization layer
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        if len(stage_dim) != len(stage_size):
+            raise TypeError('stage_dim and stage_size mismatch!')
+        else:
+            self.stage_num = len(stage_dim)
+
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.cls_token_num = 1 if use_cls_token else 0
+        self.use_pos_embedding = use_pos_embedding
+
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        act_layer = act_layer or nn.GELU
+
+        # backbone CNN
+        self.backbone = backbone
+
+        # Attention module
+        if use_att_module is not None:
+            if use_att_module in ['SimAM', 'CBAM', 'SE']:
+                Attention_module = use_att_module
+            else:
+                Attention_module = None
+        else:
+            Attention_module = None
+
+        self.patch_embed = Last_feature_map_Embed(patch_size=patch_size, feature_size=stage_size[-1],
+                                                  feature_dim=stage_dim[-1], embed_dim=self.embed_dim,
+                                                  Attention_module=Attention_module)
+        num_patches = self.patch_embed.num_patches
+
+        # global sharing cls token and positional embedding
+        self.cls_token_0 = nn.Parameter(torch.zeros(1, 1, embed_dim))  # like message token
+        if self.use_pos_embedding:
+            self.pos_embed_0 = nn.Parameter(torch.zeros(1, num_patches + self.cls_token_num, embed_dim))
+
+        '''
+        self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed_1 = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+
+        self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed_2 = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+
+        self.cls_token_3 = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed_3 = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+
+        self.cls_token_4 = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed_4 = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+        '''
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        self.dec1 = Decoder_Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                                  drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[0], norm_layer=norm_layer,
+                                  act_layer=act_layer)
+        self.Fo1 = Focus_Embed(patch_size=patch_size, target_feature_size=stage_size[-1], feature_size=stage_size[0],
+                               feature_dim=stage_dim[0], embed_dim=embed_dim, Attention_module=Attention_module,
+                               norm_layer=norm_layer)
+
+        self.dec2 = Decoder_Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                                  drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[1], norm_layer=norm_layer,
+                                  act_layer=act_layer)
+        self.Fo2 = Focus_Embed(patch_size=patch_size, target_feature_size=stage_size[-1], feature_size=stage_size[1],
+                               feature_dim=stage_dim[1], embed_dim=embed_dim, Attention_module=Attention_module,
+                               norm_layer=norm_layer)
+
+        self.dec3 = Decoder_Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                                  drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[2], norm_layer=norm_layer,
+                                  act_layer=act_layer)
+        self.Fo3 = Focus_Embed(patch_size=patch_size, target_feature_size=stage_size[-1], feature_size=stage_size[2],
+                               feature_dim=stage_dim[2], embed_dim=embed_dim, Attention_module=Attention_module,
+                               norm_layer=norm_layer)
+
+        if self.stage_num == 4:
+            self.dec4 = Decoder_Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                                      drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[3], norm_layer=norm_layer,
+                                      act_layer=act_layer)
+            self.Fo4 = Focus_Embed(patch_size=patch_size, target_feature_size=stage_size[-1],
+                                   feature_size=stage_size[-1],
+                                   feature_dim=stage_dim[-1], embed_dim=embed_dim, Attention_module=Attention_module,
+                                   norm_layer=norm_layer)
+
+        self.norm = norm_layer(embed_dim)
+
+        # Representation layer
+        if representation_size:
+            self.num_features = representation_size
+            self.pre_logits = nn.Sequential(OrderedDict([
+                ('fc', nn.Linear(embed_dim, representation_size)),
+                ('act', nn.Tanh())
+            ]))
+        else:
+            self.pre_logits = nn.Identity()
+
+        # Classifier head(s)
+        self.head = nn.Linear(self.num_features, self.num_classes) if self.num_classes > 0 else nn.Identity()
+        self.head_dist = None
+
+    def forward_features(self, x):
+        if self.stage_num == 3:
+            stage1_out, stage2_out, stage3_out = self.backbone(x)
+            # embedding the last feature map
+            x = self.patch_embed(stage3_out)
+
+        elif self.stage_num == 4:
+            stage1_out, stage2_out, stage3_out, stage4_out = self.backbone(x)
+            # embedding the last feature map
+            x = self.patch_embed(stage4_out)
+        else:
+            raise TypeError('stage_dim is not legal !')
+
+        # get guidance info
+        s1_q, s1_k = self.Fo1(stage1_out)
+        s2_q, s2_k = self.Fo2(stage2_out)
+        s3_q, s3_k = self.Fo3(stage3_out)
+        if self.stage_num == 4:
+            s4_q, s4_k = self.Fo4(stage4_out)
+
+        if self.cls_token_num != 0:  # concat cls token
+            # process the（cls token / message token）
+            cls_token_0 = self.cls_token_0.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+            x = torch.cat((cls_token_0, x), dim=1)  # 增加classification head patch
+
+            s1_q = torch.cat((cls_token_0, s1_q), dim=1)
+            s1_k = torch.cat((cls_token_0, s1_k), dim=1)
+            s2_q = torch.cat((cls_token_0, s2_q), dim=1)
+            s2_k = torch.cat((cls_token_0, s2_k), dim=1)
+            s3_q = torch.cat((cls_token_0, s3_q), dim=1)
+            s3_k = torch.cat((cls_token_0, s3_k), dim=1)
+            if self.stage_num == 4:
+                s4_q = torch.cat((cls_token_0, s4_q), dim=1)
+                s4_k = torch.cat((cls_token_0, s4_k), dim=1)
+
+        if self.use_pos_embedding:
+
+            s1_q = self.pos_drop(s1_q + self.pos_embed_0)
+            s1_k = self.pos_drop(s1_k + self.pos_embed_0)
+            s2_q = self.pos_drop(s2_q + self.pos_embed_0)
+            s2_k = self.pos_drop(s2_k + self.pos_embed_0)
+            s3_q = self.pos_drop(s3_q + self.pos_embed_0)
+            s3_k = self.pos_drop(s3_k + self.pos_embed_0)
+            if self.stage_num == 4:
+                s4_q = self.pos_drop(s4_q + self.pos_embed_0)
+                s4_k = self.pos_drop(s4_k + self.pos_embed_0)
+
+            # plus to encoding positional infor
+            x = self.pos_drop(x + self.pos_embed_0)
+
+        else:
+
+            s1_q = self.pos_drop(s1_q)
+            s1_k = self.pos_drop(s1_k)
+            s2_q = self.pos_drop(s2_q)
+            s2_k = self.pos_drop(s2_k)
+            s3_q = self.pos_drop(s3_q)
+            s3_k = self.pos_drop(s3_k)
+            if self.stage_num == 4:
+                s4_q = self.pos_drop(s4_q)
+                s4_k = self.pos_drop(s4_k)
+
+            # stem's feature map
+            x = self.pos_drop(x)
+
+        # Decoder module use the guidance to help global modeling process
+
+        x = self.dec1(s1_q, s1_k, x)
+
+        x = self.dec2(s2_q, s2_k, x)
+
+        x = self.dec3(s3_q, s3_k, x)
+
+        if self.stage_num == 4:
+            x = self.dec4(s4_q, s4_k, x)
+
+        x = self.norm(x)
+        return self.pre_logits(x[:, 0])  # take the first cls token
+
+    def forward(self, x):
+        x = self.forward_features(x)  # connect the cls token to the cls head
+        x = self.head(x)
+        return x

PuzzleTuning/Backbone/VPT_structure.py

@@ -0,0 +1,133 @@
+"""
+VPT     Script  ver： Oct 17th 14:30
+
+based on
+timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
+"""
+
+import torch
+import torch.nn as nn
+
+from timm.models.vision_transformer import VisionTransformer, PatchEmbed
+
+
+class VPT_ViT(VisionTransformer):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=True, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
+                 embed_layer=PatchEmbed, norm_layer=None, act_layer=None, Prompt_Token_num=1,
+                 VPT_type="Shallow", basic_state_dict=None):
+
+        # Recreate ViT
+        super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans, num_classes=num_classes,
+                         embed_dim=embed_dim, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio,
+                         qkv_bias=qkv_bias, drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                         drop_path_rate=drop_path_rate, embed_layer=embed_layer,
+                         norm_layer=norm_layer, act_layer=act_layer)
+
+        # load basic state_dict
+        if basic_state_dict is not None:
+            self.load_state_dict(basic_state_dict, False)
+
+        self.VPT_type = VPT_type
+        if VPT_type == "Deep":
+            self.Prompt_Tokens = nn.Parameter(torch.zeros(depth, Prompt_Token_num, embed_dim))
+        else:  # "Shallow"
+            self.Prompt_Tokens = nn.Parameter(torch.zeros(1, Prompt_Token_num, embed_dim))
+
+    def New_CLS_head(self, new_classes=15):
+        if new_classes != 0:
+            self.head = nn.Linear(self.embed_dim, new_classes)
+        else:
+            self.head = nn.Identity()
+
+    def Freeze(self):
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.Prompt_Tokens.requires_grad = True
+        try:
+            for param in self.head.parameters():
+                param.requires_grad = True
+        except:
+            pass
+
+    def UnFreeze(self):
+        for param in self.parameters():
+            param.requires_grad = True
+
+    def obtain_prompt(self):
+        prompt_state_dict = {'head': self.head.state_dict(),
+                             'Prompt_Tokens': self.Prompt_Tokens}
+        # print(prompt_state_dict)
+        return prompt_state_dict
+
+    def load_prompt(self, prompt_state_dict):
+        try:
+            self.head.load_state_dict(prompt_state_dict['head'], False)
+        except:
+            print('head not match, so skip head')
+        else:
+            print('prompt head match')
+
+        if self.Prompt_Tokens.shape == prompt_state_dict['Prompt_Tokens'].shape:
+
+            # device check
+            Prompt_Tokens = nn.Parameter(prompt_state_dict['Prompt_Tokens'].cpu())
+            Prompt_Tokens.to(torch.device(self.Prompt_Tokens.device))
+
+            self.Prompt_Tokens = Prompt_Tokens
+
+        else:
+            print('\n !!! cannot load prompt')
+            print('shape of model req prompt', self.Prompt_Tokens.shape)
+            print('shape of model given prompt', prompt_state_dict['Prompt_Tokens'].shape)
+            print('')
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        # print(x.shape,self.pos_embed.shape)
+        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
+
+        # concatenate CLS token
+        x = torch.cat((cls_token, x), dim=1)
+        x = self.pos_drop(x + self.pos_embed)
+
+        if self.VPT_type == "Deep":
+
+            Prompt_Token_num = self.Prompt_Tokens.shape[1]
+
+            for i in range(len(self.blocks)):
+                # concatenate Prompt_Tokens
+                Prompt_Tokens = self.Prompt_Tokens[i].unsqueeze(0)
+                # firstly concatenate
+                x = torch.cat((x, Prompt_Tokens.expand(x.shape[0], -1, -1)), dim=1)
+                num_tokens = x.shape[1]
+                # lastly remove, a genius trick
+                x = self.blocks[i](x)[:, :num_tokens - Prompt_Token_num]
+
+        else:  # self.VPT_type == "Shallow"
+            Prompt_Token_num = self.Prompt_Tokens.shape[1]
+
+            # concatenate Prompt_Tokens
+            Prompt_Tokens = self.Prompt_Tokens.expand(x.shape[0], -1, -1)
+            x = torch.cat((x, Prompt_Tokens), dim=1)
+            num_tokens = x.shape[1]
+            # Sequntially procees
+            x = self.blocks(x)[:, :num_tokens - Prompt_Token_num]
+
+        x = self.norm(x)
+        return x
+
+    def forward(self, x):
+
+        x = self.forward_features(x)
+
+        # use cls token for cls head
+        try:
+            x = self.pre_logits(x[:, 0, :])
+        except:
+            x = self.fc_norm(x[:, 0, :])
+        else:
+            pass
+        x = self.head(x)
+        return x

PuzzleTuning/Backbone/attention_modules.py

@@ -0,0 +1,303 @@
+"""
+attention modules in ['SimAM', 'CBAM', 'SE', 'GAM']  were applied in the ablation study
+
+ver： Dec 24th 15：00
+
+
+ref:
+https://github.com/xmu-xiaoma666/External-Attention-pytorch
+"""
+
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+from torch.nn import init
+
+
+# help func
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True,
+                 bn=True, bias=False):
+        super(BasicConv, self).__init__()
+        self.out_channels = out_planes
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding,
+                              dilation=dilation, groups=groups, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None
+        self.relu = nn.ReLU() if relu else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+
+class ChannelGate(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
+        super(ChannelGate, self).__init__()
+        self.gate_channels = gate_channels
+        self.mlp = nn.Sequential(
+            Flatten(),
+            nn.Linear(gate_channels, int(gate_channels // reduction_ratio)),
+            nn.ReLU(),
+            nn.Linear(int(gate_channels // reduction_ratio), gate_channels)
+        )
+        self.pool_types = pool_types
+
+    def forward(self, x):
+        channel_att_sum = None
+        for pool_type in self.pool_types:
+            if pool_type == 'avg':
+                avg_pool = F.avg_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(avg_pool)
+            elif pool_type == 'max':
+                max_pool = F.max_pool2d(x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(max_pool)
+            elif pool_type == 'lp':
+                lp_pool = F.lp_pool2d(x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp(lp_pool)
+            elif pool_type == 'lse':
+                # LSE pool only
+                lse_pool = logsumexp_2d(x)
+                channel_att_raw = self.mlp(lse_pool)
+
+            if channel_att_sum is None:
+                channel_att_sum = channel_att_raw
+            else:
+                channel_att_sum = channel_att_sum + channel_att_raw
+
+        scale = F.sigmoid(channel_att_sum).unsqueeze(2).unsqueeze(3).expand_as(x)
+        return x * scale
+
+
+def logsumexp_2d(tensor):
+    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
+    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
+    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
+    return outputs
+
+
+class ChannelPool(nn.Module):
+    def forward(self, x):
+        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)
+
+
+class SpatialGate(nn.Module):
+    def __init__(self):
+        super(SpatialGate, self).__init__()
+        kernel_size = 7
+        self.compress = ChannelPool()
+        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=int((kernel_size - 1) // 2), relu=False)
+
+    def forward(self, x):
+        x_compress = self.compress(x)
+        x_out = self.spatial(x_compress)
+        scale = F.sigmoid(x_out)  # broadcasting
+        return x * scale
+
+
+# attention modules：
+class cbam_module(nn.Module):
+    """
+    module：CBAM
+
+    input、output= b, c, h, w
+
+    paper：
+    https://arxiv.org/abs/1807.06521
+    code：
+    https://github.com/ZjjConan/SimAM/blob/master/networks/attentions
+    """
+
+    def __init__(self, gate_channels, reduction=16, pool_types=['avg', 'max'], no_spatial=False):
+        super(cbam_module, self).__init__()
+        self.ChannelGate = ChannelGate(gate_channels, reduction, pool_types)
+        self.no_spatial = no_spatial
+        if not no_spatial:
+            self.SpatialGate = SpatialGate()
+
+    @staticmethod
+    def get_module_name():
+        return "cbam"
+
+    def forward(self, x):
+        x_out = self.ChannelGate(x)
+        if not self.no_spatial:
+            x_out = self.SpatialGate(x_out)
+        return x_out
+
+
+class se_module(nn.Module):
+    """
+    module: SE
+
+    input、output= b, c, h, w
+
+    from paper Squeeze-and-Excitation Networks
+    SE-Net  https://arxiv.org/abs/1709.01507
+    code：
+    https://github.com/ZjjConan/SimAM/blob/master/networks/attentions
+    """
+
+    def __init__(self, channel, reduction=16):
+        super(se_module, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, int(channel // reduction), bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(int(channel // reduction), channel, bias=False),
+            nn.Sigmoid()
+        )
+
+    @staticmethod
+    def get_module_name():
+        return "se"
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+class simam_module(torch.nn.Module):
+    """
+    module:SimAM
+
+    input、output= b, c, h, w
+
+    paper:(ICML)
+    SimAM: A Simple, Parameter-Free Attention Module for Convolutional Neural Networks
+    code:
+    https://github.com/ZjjConan/SimAM/blob/master/networks/attentions/simam_module.py
+    """
+
+    def __init__(self, channels=None, e_lambda=1e-4):
+        super(simam_module, self).__init__()
+
+        self.activaton = nn.Sigmoid()
+        self.e_lambda = e_lambda
+
+    def __repr__(self):
+        s = self.__class__.__name__ + '('
+        s += ('lambda=%f)' % self.e_lambda)
+        return s
+
+    @staticmethod
+    def get_module_name():
+        return "simam"
+
+    def forward(self, x):
+        b, c, h, w = x.size()
+
+        n = w * h - 1
+
+        x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
+        y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5
+
+        return x * self.activaton(y)
+
+
+class ResidualAttention(nn.Module):
+    """
+    module: ResidualAttention
+
+    input、output= b, c, h, w
+
+    Paper：ICCV 2021 Residual Attention: A Simple but Effective Method for Multi-Label Recognition
+    code：https://github.com/xmu-xiaoma666/External-Attention-pytorch/blob/master/attention/ResidualAttention.py
+    """
+
+    def __init__(self, channel=512, num_class=1000, la=0.2):
+        super().__init__()
+        self.la = la
+        self.fc = nn.Conv2d(in_channels=channel, out_channels=num_class, kernel_size=1, stride=1, bias=False)
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        y_raw = self.fc(x).flatten(2)  # b,num_class,hxw
+        y_avg = torch.mean(y_raw, dim=2)  # b,num_class
+        y_max = torch.max(y_raw, dim=2)[0]  # b,num_class
+        score = y_avg + self.la * y_max
+        return score
+
+
+class eca_module(nn.Module):
+    """Constructs a ECA module.
+
+    Args:
+        channel: Number of channels of the input feature map
+        k_size: Adaptive selection of kernel size
+    """
+    def __init__(self, channel, k_size=3):
+        super(eca_module, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        # x: input features with shape [b, c, h, w]
+        b, c, h, w = x.size()
+
+        # feature descriptor on the global spatial information
+        y = self.avg_pool(x)
+
+        # Two different branches of ECA module
+        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
+
+        # Multi-scale information fusion
+        y = self.sigmoid(y)
+
+        return x * y.expand_as(x)
+
+
+class GAM_Attention(nn.Module):
+    """
+    module：GAM
+
+    input= b, in_channels, h, w
+    output= b, out_channels, h, w
+
+    paper：
+    Global Attention Mechanism: Retain Information to Enhance Channel-Spatial Interactions
+    https://arxiv.org/abs/2112.05561
+    code：
+    https://mp.weixin.qq.com/s/VL6rXjyUDmHToYTqM32hUg
+    """
+    def __init__(self, in_channels, out_channels, rate=4):
+        super(GAM_Attention, self).__init__()
+
+        self.channel_attention = nn.Sequential(
+            nn.Linear(in_channels, int(in_channels / rate)),
+            nn.ReLU(inplace=True),
+            nn.Linear(int(in_channels / rate), in_channels)
+        )
+
+        self.spatial_attention = nn.Sequential(
+            nn.Conv2d(in_channels, int(in_channels / rate), kernel_size=7, padding=3),
+            nn.BatchNorm2d(int(in_channels / rate)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(int(in_channels / rate), out_channels, kernel_size=7, padding=3),
+            nn.BatchNorm2d(out_channels)
+        )
+
+    def forward(self, x):
+        b, c, h, w = x.shape
+        x_permute = x.permute(0, 2, 3, 1).view(b, -1, c)
+        x_att_permute = self.channel_attention(x_permute).view(b, h, w, c)
+        x_channel_att = x_att_permute.permute(0, 3, 1, 2)
+
+        x = x * x_channel_att
+
+        x_spatial_att = self.spatial_attention(x).sigmoid()
+        out = x * x_spatial_att
+
+        return out

PuzzleTuning/Backbone/counterpart_models/README.md

@@ -0,0 +1,25 @@
+Recent SOTA works in fine-grained Tasks
+
+
+CrossFormer
+
+Paper: 
+https://arxiv.org/pdf/2108.00154.pdf
+
+
+Code from:
+https://github.com/cheerss/CrossFormer
+
+
+
+Conformer
+Paper: 
+https://arxiv.org/pdf/2105.03889.pdf
+
+
+Code from:
+https://github.com/pengzhiliang/Conformer/blob/main/conformer.py
+
+
+both work will be compared with official pretrained backbone
+and a new MLP head (classification head).

PuzzleTuning/Backbone/counterpart_models/conformer.py

@@ -0,0 +1,483 @@
+"""
+From Conformer with alter: conv and trans cls head was changed to volting together
+ver： DEC 1st 16:00 official release
+
+ref: https://github.com/pengzhiliang/Conformer/blob/main/conformer.py
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+from timm.models.layers import DropPath, trunc_normal_
+
+
+class Mlp(nn.Module):  # FFN
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):  # MHSA
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape  # N is patch number, C is patch dimension
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)  # re arrange
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):  # Encoder from ViT
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=partial(nn.LayerNorm, eps=1e-6)):
+        super().__init__()
+        # pre norm 1
+        self.norm1 = norm_layer(dim)
+        # MHSA
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        # pre norm 2
+        self.norm2 = norm_layer(dim)
+
+        # FFN(MLP)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class ConvBlock(nn.Module):  # ResNet bottleneck Convblock actually
+
+    def __init__(self, inplanes, outplanes, stride=1, res_conv=False, act_layer=nn.ReLU, groups=1,
+                 norm_layer=partial(nn.BatchNorm2d, eps=1e-6), drop_block=None, drop_path=None):
+        super(ConvBlock, self).__init__()
+
+        expansion = 4
+        med_planes = outplanes // expansion
+
+        self.conv1 = nn.Conv2d(inplanes, med_planes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn1 = norm_layer(med_planes)
+        self.act1 = act_layer(inplace=True)
+
+        self.conv2 = nn.Conv2d(med_planes, med_planes, kernel_size=3, stride=stride, groups=groups, padding=1,
+                               bias=False)
+        self.bn2 = norm_layer(med_planes)
+        self.act2 = act_layer(inplace=True)
+
+        self.conv3 = nn.Conv2d(med_planes, outplanes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn3 = norm_layer(outplanes)
+        self.act3 = act_layer(inplace=True)
+
+        if res_conv:
+            self.residual_conv = nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=stride, padding=0, bias=False)
+            self.residual_bn = norm_layer(outplanes)
+
+        self.res_conv = res_conv
+        self.drop_block = drop_block
+        self.drop_path = drop_path
+
+    def zero_init_last_bn(self):
+        nn.init.zeros_(self.bn3.weight)
+
+    def forward(self, x, x_t=None, return_x_2=True):
+        residual = x
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+        x = self.act1(x)
+
+        x = self.conv2(x) if x_t is None else self.conv2(x + x_t)
+        x = self.bn2(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+        x2 = self.act2(x)
+
+        x = self.conv3(x2)
+        x = self.bn3(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        if self.res_conv:
+            residual = self.residual_conv(residual)
+            residual = self.residual_bn(residual)
+
+        x += residual
+        x = self.act3(x)
+
+        if return_x_2:
+            return x, x2
+        else:
+            return x
+
+
+class FCUDown(nn.Module):
+    """ CNN feature maps -> Transformer patch embeddings
+    """
+
+    def __init__(self, inplanes, outplanes, dw_stride, act_layer=nn.GELU,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6)):
+        super(FCUDown, self).__init__()
+        self.dw_stride = dw_stride
+
+        self.conv_project = nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=1, padding=0)  # fix dimension
+        self.sample_pooling = nn.AvgPool2d(kernel_size=dw_stride, stride=dw_stride)  # fix feature map size
+
+        self.ln = norm_layer(outplanes)
+        self.act = act_layer()
+
+    def forward(self, x, x_t):
+        x = self.conv_project(x)  # [N, C, H, W]
+
+        x = self.sample_pooling(x).flatten(2).transpose(1, 2)
+        x = self.ln(x)
+        x = self.act(x)
+
+        x = torch.cat([x_t[:, 0][:, None, :], x], dim=1)  # concatenate class token from x_t
+
+        return x
+
+
+class FCUUp(nn.Module):
+    """ Transformer patch embeddings -> CNN feature maps
+    by interpolate operation
+    """
+
+    def __init__(self, inplanes, outplanes, up_stride, act_layer=nn.ReLU,
+                 norm_layer=partial(nn.BatchNorm2d, eps=1e-6), ):
+        super(FCUUp, self).__init__()
+
+        self.up_stride = up_stride
+        self.conv_project = nn.Conv2d(inplanes, outplanes, kernel_size=1, stride=1, padding=0)
+        self.bn = norm_layer(outplanes)
+        self.act = act_layer()
+
+    def forward(self, x, H, W):  # interpolate to
+        B, _, C = x.shape
+
+        # [N, 197, 384] -> [N, 196, 384] -> [N, 384, 196] -> [N, 384, 14, 14]
+        x_r = x[:, 1:].transpose(1, 2).reshape(B, C, H, W)  # drop cls token of x_t
+
+        x_r = self.act(self.bn(self.conv_project(x_r)))
+
+        return F.interpolate(x_r, size=(H * self.up_stride, W * self.up_stride))  # interpolate operation
+
+
+class Med_ConvBlock(nn.Module):  # ResNet bottleneck indentity actually
+    """ special case for Convblock without down sampling,
+    """
+
+    def __init__(self, inplanes, act_layer=nn.ReLU, groups=1, norm_layer=partial(nn.BatchNorm2d, eps=1e-6),
+                 drop_block=None, drop_path=None):
+
+        super(Med_ConvBlock, self).__init__()
+
+        expansion = 4
+        med_planes = inplanes // expansion
+
+        self.conv1 = nn.Conv2d(inplanes, med_planes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn1 = norm_layer(med_planes)
+        self.act1 = act_layer(inplace=True)
+
+        self.conv2 = nn.Conv2d(med_planes, med_planes, kernel_size=3, stride=1, groups=groups, padding=1, bias=False)
+        self.bn2 = norm_layer(med_planes)
+        self.act2 = act_layer(inplace=True)
+
+        self.conv3 = nn.Conv2d(med_planes, inplanes, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn3 = norm_layer(inplanes)
+        self.act3 = act_layer(inplace=True)
+
+        self.drop_block = drop_block
+        self.drop_path = drop_path
+
+    def zero_init_last_bn(self):
+        nn.init.zeros_(self.bn3.weight)
+
+    def forward(self, x):
+        residual = x
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+        x = self.act1(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+        x = self.act2(x)
+
+        x = self.conv3(x)
+        x = self.bn3(x)
+        if self.drop_block is not None:
+            x = self.drop_block(x)
+
+        if self.drop_path is not None:
+            x = self.drop_path(x)
+
+        x += residual
+        x = self.act3(x)
+
+        return x
+
+
+class ConvTransBlock(nn.Module):
+    """
+    Basic module for ConvTransformer, keep feature maps for CNN block and patch embeddings for transformer encoder block
+    """
+
+    def __init__(self, inplanes, outplanes, res_conv, stride, dw_stride, embed_dim, num_heads=12, mlp_ratio=4.,
+                 qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., drop_path_rate=0.,
+                 last_fusion=False, num_med_block=0, groups=1):
+
+        super(ConvTransBlock, self).__init__()
+        expansion = 4
+        # ConvBlock
+        self.cnn_block = ConvBlock(inplanes=inplanes, outplanes=outplanes, res_conv=res_conv, stride=stride,
+                                   groups=groups)
+
+        if last_fusion:
+            self.fusion_block = ConvBlock(inplanes=outplanes, outplanes=outplanes, stride=2, res_conv=True,
+                                          groups=groups)
+        else:
+            self.fusion_block = ConvBlock(inplanes=outplanes, outplanes=outplanes, groups=groups)
+
+        # identity block
+        if num_med_block > 0:
+            self.med_block = []
+            for i in range(num_med_block):
+                self.med_block.append(Med_ConvBlock(inplanes=outplanes, groups=groups))
+
+            self.med_block = nn.ModuleList(self.med_block)  # nn.ModuleList
+
+        # FCU
+        self.squeeze_block = FCUDown(inplanes=outplanes // expansion, outplanes=embed_dim, dw_stride=dw_stride)
+
+        self.expand_block = FCUUp(inplanes=embed_dim, outplanes=outplanes // expansion, up_stride=dw_stride)
+
+        # Transformer Encoder block
+        self.trans_block = Block(
+            dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=drop_path_rate)
+
+        self.dw_stride = dw_stride
+        self.embed_dim = embed_dim
+        self.num_med_block = num_med_block
+        self.last_fusion = last_fusion
+
+    def forward(self, x, x_t):
+        x, x2 = self.cnn_block(x)
+
+        _, _, H, W = x2.shape
+
+        x_st = self.squeeze_block(x2, x_t)
+
+        x_t = self.trans_block(x_st + x_t)
+
+        if self.num_med_block > 0:
+            for m in self.med_block:
+                x = m(x)
+
+        x_t_r = self.expand_block(x_t, H // self.dw_stride, W // self.dw_stride)
+        x = self.fusion_block(x, x_t_r, return_x_2=False)
+
+        return x, x_t
+
+
+class Conformer(nn.Module):
+
+    def __init__(self, patch_size=16, in_chans=3, num_classes=1000, base_channel=64, channel_ratio=4, num_med_block=0,
+                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.):
+
+        # Transformer
+        super().__init__()
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        assert depth % 3 == 0
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.trans_dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+        # Classifier head
+        self.trans_norm = nn.LayerNorm(embed_dim)
+        self.trans_cls_head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        self.pooling = nn.AdaptiveAvgPool2d(1)
+        self.conv_cls_head = nn.Linear(int(256 * channel_ratio), num_classes)
+        self.cls_head = nn.Linear(int(2 * num_classes), num_classes)
+
+        # Stem stage: get the feature maps by conv block (copied form resnet.py)
+        self.conv1 = nn.Conv2d(in_chans, 64, kernel_size=7, stride=2, padding=3, bias=False)  # 1 / 2 [112, 112]
+        self.bn1 = nn.BatchNorm2d(64)
+        self.act1 = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # 1 / 4 [56, 56]
+
+        # 1 stage
+        stage_1_channel = int(base_channel * channel_ratio)
+        trans_dw_stride = patch_size // 4
+        self.conv_1 = ConvBlock(inplanes=64, outplanes=stage_1_channel, res_conv=True, stride=1)
+        # embedding
+        self.trans_patch_conv = nn.Conv2d(64, embed_dim, kernel_size=trans_dw_stride, stride=trans_dw_stride, padding=0)
+        self.trans_1 = Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                             qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=self.trans_dpr[0],
+                             )
+
+        # 2~4 stage
+        init_stage = 2
+        fin_stage = depth // 3 + 1
+        for i in range(init_stage, fin_stage):
+            self.add_module('conv_trans_' + str(i),
+                            ConvTransBlock(
+                                stage_1_channel, stage_1_channel, False, 1, dw_stride=trans_dw_stride,
+                                embed_dim=embed_dim,
+                                num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                drop_path_rate=self.trans_dpr[i - 1],
+                                num_med_block=num_med_block
+                            )
+                            )
+
+        stage_2_channel = int(base_channel * channel_ratio * 2)
+        # 5~8 stage
+        init_stage = fin_stage  # 5
+        fin_stage = fin_stage + depth // 3  # 9
+        for i in range(init_stage, fin_stage):
+            s = 2 if i == init_stage else 1
+            in_channel = stage_1_channel if i == init_stage else stage_2_channel
+            res_conv = True if i == init_stage else False
+            self.add_module('conv_trans_' + str(i),
+                            ConvTransBlock(
+                                in_channel, stage_2_channel, res_conv, s, dw_stride=trans_dw_stride // 2,
+                                embed_dim=embed_dim,
+                                num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                drop_path_rate=self.trans_dpr[i - 1],
+                                num_med_block=num_med_block
+                            )
+                            )
+
+        stage_3_channel = int(base_channel * channel_ratio * 2 * 2)
+        # 9~12 stage
+        init_stage = fin_stage  # 9
+        fin_stage = fin_stage + depth // 3  # 13
+        for i in range(init_stage, fin_stage):
+            s = 2 if i == init_stage else 1
+            in_channel = stage_2_channel if i == init_stage else stage_3_channel
+            res_conv = True if i == init_stage else False
+            last_fusion = True if i == depth else False
+            self.add_module('conv_trans_' + str(i),
+                            ConvTransBlock(
+                                in_channel, stage_3_channel, res_conv, s, dw_stride=trans_dw_stride // 4,
+                                embed_dim=embed_dim,
+                                num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                drop_path_rate=self.trans_dpr[i - 1],
+                                num_med_block=num_med_block, last_fusion=last_fusion
+                            )
+                            )
+        self.fin_stage = fin_stage
+
+        trunc_normal_(self.cls_token, std=.02)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+        elif isinstance(m, nn.BatchNorm2d):
+            nn.init.constant_(m.weight, 1.)
+            nn.init.constant_(m.bias, 0.)
+        elif isinstance(m, nn.GroupNorm):
+            nn.init.constant_(m.weight, 1.)
+            nn.init.constant_(m.bias, 0.)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'cls_token'}
+
+    def forward(self, x):
+        B = x.shape[0]
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+
+        # pdb.set_trace()
+        # stem stage [N, 3, 224, 224] -> [N, 64, 56, 56]
+        x_base = self.maxpool(self.act1(self.bn1(self.conv1(x))))
+
+        # 1 stage
+        x = self.conv_1(x_base, return_x_2=False)
+        # embedding: [N, 64, 56, 56] -> [N, d, p, p] -> [N, d, p^2] -> [N, p^2, d] -> [N, p^2 + 1, d]
+        x_t = self.trans_patch_conv(x_base).flatten(2).transpose(1, 2)
+        x_t = torch.cat([cls_tokens, x_t], dim=1)
+        x_t = self.trans_1(x_t)
+
+        # 2 ~ final 
+        for i in range(2, self.fin_stage):
+            x, x_t = eval('self.conv_trans_' + str(i))(x, x_t)
+
+        # conv classification
+        x_p = self.pooling(x).flatten(1)
+        conv_cls = self.conv_cls_head(x_p)
+
+        # trans classification
+        x_t = self.trans_norm(x_t)
+        tran_cls = self.trans_cls_head(x_t[:, 0])
+
+        # 加一个类别投票
+        cls = torch.cat([conv_cls, tran_cls], dim=1)
+        cls = self.cls_head(cls)
+        return cls
+
+        # return [conv_cls, tran_cls]

PuzzleTuning/Backbone/counterpart_models/crossformer.py

@@ -0,0 +1,624 @@
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class DynamicPosBias(nn.Module):
+    def __init__(self, dim, num_heads, residual):
+        super().__init__()
+        self.residual = residual
+        self.num_heads = num_heads
+        self.pos_dim = dim // 4
+        self.pos_proj = nn.Linear(2, self.pos_dim)
+        self.pos1 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos2 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim)
+        )
+        self.pos3 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.num_heads)
+        )
+    def forward(self, biases):
+        if self.residual:
+            pos = self.pos_proj(biases) # 2Wh-1 * 2Ww-1, heads
+            pos = pos + self.pos1(pos)
+            pos = pos + self.pos2(pos)
+            pos = self.pos3(pos)
+        else:
+            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
+        return pos
+
+    def flops(self, N):
+        flops = N * 2 * self.pos_dim
+        flops += N * self.pos_dim * self.pos_dim
+        flops += N * self.pos_dim * self.pos_dim
+        flops += N * self.pos_dim * self.num_heads
+        return flops
+
+class Attention(nn.Module):
+    r""" Multi-head self attention module with dynamic position bias.
+
+    Args:
+        dim (int): Number of input channels.
+        group_size (tuple[int]): The height and width of the group.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, group_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.,
+                 position_bias=True):
+
+        super().__init__()
+        self.dim = dim
+        self.group_size = group_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.position_bias = position_bias
+
+        if position_bias:
+            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
+            
+            # generate mother-set
+            position_bias_h = torch.arange(1 - self.group_size[0], self.group_size[0])
+            position_bias_w = torch.arange(1 - self.group_size[1], self.group_size[1])
+            biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w]))  # 2, 2Wh-1, 2W2-1
+            biases = biases.flatten(1).transpose(0, 1).float()
+            self.register_buffer("biases", biases)
+
+            # get pair-wise relative position index for each token inside the group
+            coords_h = torch.arange(self.group_size[0])
+            coords_w = torch.arange(self.group_size[1])
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords[:, :, 0] += self.group_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += self.group_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * self.group_size[1] - 1
+            relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+            self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_groups*B, N, C)
+            mask: (0/-inf) mask with shape of (num_groups, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        if self.position_bias:
+            pos = self.pos(self.biases) # 2Wh-1 * 2Ww-1, heads
+            # select position bias
+            relative_position_bias = pos[self.relative_position_index.view(-1)].view(
+                self.group_size[0] * self.group_size[1], self.group_size[0] * self.group_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, group_size={self.group_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 group with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        if self.position_bias:
+            flops += self.pos.flops(N)
+        return flops
+
+
+class CrossFormerBlock(nn.Module):
+    r""" CrossFormer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        group_size (int): Group size.
+        lsda_flag (int): use SDA or LDA, 0 for SDA and 1 for LDA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, group_size=7, lsda_flag=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, num_patch_size=1):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.group_size = group_size
+        self.lsda_flag = lsda_flag
+        self.mlp_ratio = mlp_ratio
+        self.num_patch_size = num_patch_size
+        if min(self.input_resolution) <= self.group_size:
+            # if group size is larger than input resolution, we don't partition groups
+            self.lsda_flag = 0
+            self.group_size = min(self.input_resolution)
+
+        self.norm1 = norm_layer(dim)
+
+        self.attn = Attention(
+            dim, group_size=to_2tuple(self.group_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
+            position_bias=True)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        attn_mask = None
+        self.register_buffer("attn_mask", attn_mask)
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size %d, %d, %d" % (L, H, W)
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # group embeddings
+        G = self.group_size
+        if self.lsda_flag == 0: # 0 for SDA
+            x = x.reshape(B, H // G, G, W // G, G, C).permute(0, 1, 3, 2, 4, 5)
+        else: # 1 for LDA
+            x = x.reshape(B, G, H // G, G, W // G, C).permute(0, 2, 4, 1, 3, 5)
+        x = x.reshape(B * H * W // G**2, G**2, C)
+
+        # multi-head self-attention
+        x = self.attn(x, mask=self.attn_mask)  # nW*B, G*G, C
+
+        # ungroup embeddings
+        x = x.reshape(B, H // G, W // G, G, G, C)
+        if self.lsda_flag == 0:
+            x = x.permute(0, 1, 3, 2, 4, 5).reshape(B, H, W, C)
+        else:
+            x = x.permute(0, 3, 1, 4, 2, 5).reshape(B, H, W, C)
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"group_size={self.group_size}, lsda_flag={self.lsda_flag}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # LSDA
+        nW = H * W / self.group_size / self.group_size
+        flops += nW * self.attn.flops(self.group_size * self.group_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm, patch_size=[2], num_input_patch_size=1):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reductions = nn.ModuleList()
+        self.patch_size = patch_size
+        self.norm = norm_layer(dim)
+
+        for i, ps in enumerate(patch_size):
+            if i == len(patch_size) - 1:
+                out_dim = 2 * dim // 2 ** i
+            else:
+                out_dim = 2 * dim // 2 ** (i + 1)
+            stride = 2
+            padding = (ps - stride) // 2
+            self.reductions.append(nn.Conv2d(dim, out_dim, kernel_size=ps, 
+                                                stride=stride, padding=padding))
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = self.norm(x)
+        x = x.view(B, H, W, C).permute(0, 3, 1, 2)
+
+        xs = []
+        for i in range(len(self.reductions)):
+            tmp_x = self.reductions[i](x).flatten(2).transpose(1, 2)
+            xs.append(tmp_x)
+        x = torch.cat(xs, dim=2)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        for i, ps in enumerate(self.patch_size):
+            if i == len(self.patch_size) - 1:
+                out_dim = 2 * self.dim // 2 ** i
+            else:
+                out_dim = 2 * self.dim // 2 ** (i + 1)
+            flops += (H // 2) * (W // 2) * ps * ps * out_dim * self.dim
+        return flops
+
+
+class Stage(nn.Module):
+    """ CrossFormer blocks for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        group_size (int): variable G in the paper, one group has GxG embeddings
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, group_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 patch_size_end=[4], num_patch_size=None):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            lsda_flag = 0 if (i % 2 == 0) else 1
+            self.blocks.append(CrossFormerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, group_size=group_size,
+                                 lsda_flag=lsda_flag,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer,
+                                 num_patch_size=num_patch_size))
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer, 
+                                         patch_size=patch_size_end, num_input_patch_size=num_patch_size)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: [4].
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=[4], in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        # patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[0] // patch_size[0]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.projs = nn.ModuleList()
+        for i, ps in enumerate(patch_size):
+            if i == len(patch_size) - 1:
+                dim = embed_dim // 2 ** i
+            else:
+                dim = embed_dim // 2 ** (i + 1)
+            stride = patch_size[0]
+            padding = (ps - patch_size[0]) // 2
+            self.projs.append(nn.Conv2d(in_chans, dim, kernel_size=ps, stride=stride, padding=padding))
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        xs = []
+        for i in range(len(self.projs)):
+            tx = self.projs[i](x).flatten(2).transpose(1, 2)
+            xs.append(tx)  # B Ph*Pw C
+        x = torch.cat(xs, dim=2)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = 0
+        for i, ps in enumerate(self.patch_size):
+            if i == len(self.patch_size) - 1:
+                dim = self.embed_dim // 2 ** i
+            else:
+                dim = self.embed_dim // 2 ** (i + 1)
+            flops += Ho * Wo * dim * self.in_chans * (self.patch_size[i] * self.patch_size[i])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class CrossFormer(nn.Module):
+    r""" CrossFormer
+        A PyTorch impl of : `CrossFormer: A Versatile Vision Transformer Based on Cross-scale Attention`  -
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each stage.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        group_size (int): Group size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=[4], in_chans=3, num_classes=1000,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 group_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, merge_size=[[2], [2], [2]], **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+
+        num_patch_sizes = [len(patch_size)] + [len(m) for m in merge_size]
+        for i_layer in range(self.num_layers):
+            patch_size_end = merge_size[i_layer] if i_layer < self.num_layers - 1 else None
+            num_patch_size = num_patch_sizes[i_layer]
+            layer = Stage(dim=int(embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               group_size=group_size[i_layer],
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=qkv_bias, qk_scale=qk_scale,
+                               drop=drop_rate, attn_drop=attn_drop_rate,
+                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                               norm_layer=norm_layer,
+                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               patch_size_end=patch_size_end,
+                               num_patch_size=num_patch_size)
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x):
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        flops += self.num_features * self.num_classes
+        return flops
+
+
+class cross_former_cls_head_warp(nn.Module):
+    def __init__(self, backbone, num_classes):
+        super().__init__()
+        embed_dim = 96
+        depths = [2, 2, 18, 2]
+        num_layers = len(depths)
+        num_features = int(embed_dim * 2 ** (num_layers - 1))
+        self.backbone = backbone
+        self.head = nn.Linear(num_features, num_classes)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.head(x)
+        return x

PuzzleTuning/Backbone/counterpart_models/crossformer_backbone.py

@@ -0,0 +1,659 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from mmdet.utils import get_root_logger
+from mmcv.runner import load_checkpoint
+
+NEG_INF = -1000000
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class DynamicPosBias(nn.Module):
+    def __init__(self, dim, num_heads, residual):
+        super().__init__()
+        self.residual = residual
+        self.num_heads = num_heads
+        self.pos_dim = dim // 4
+        self.pos_proj = nn.Linear(2, self.pos_dim)
+        self.pos1 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim),
+        )
+        self.pos2 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.pos_dim)
+        )
+        self.pos3 = nn.Sequential(
+            nn.LayerNorm(self.pos_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(self.pos_dim, self.num_heads)
+        )
+    def forward(self, biases):
+        if self.residual:
+            pos = self.pos_proj(biases) # 2Gh-1 * 2Gw-1, heads
+            pos = pos + self.pos1(pos)
+            pos = pos + self.pos2(pos)
+            pos = self.pos3(pos)
+        else:
+            pos = self.pos3(self.pos2(self.pos1(self.pos_proj(biases))))
+        return pos
+
+    def flops(self, N):
+        flops = N * 2 * self.pos_dim
+        flops += N * self.pos_dim * self.pos_dim
+        flops += N * self.pos_dim * self.pos_dim
+        flops += N * self.pos_dim * self.num_heads
+        return flops
+
+class Attention(nn.Module):
+    r""" Multi-head self attention module with relative position bias.
+
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.,
+                 position_bias=True):
+
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.position_bias = position_bias
+        if self.position_bias:
+            self.pos = DynamicPosBias(self.dim // 4, self.num_heads, residual=False)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, H, W, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Gh*Gw, Gh*Gw) or None
+        """
+        group_size = (H, W)
+        B_, N, C = x.shape
+        assert H*W == N
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous()
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1)) # (num_windows*B, N, N), N = Gh*Gw
+
+        if self.position_bias:
+            # generate mother-set
+            position_bias_h = torch.arange(1 - group_size[0], group_size[0], device=attn.device)
+            position_bias_w = torch.arange(1 - group_size[1], group_size[1], device=attn.device)
+            biases = torch.stack(torch.meshgrid([position_bias_h, position_bias_w]))  # 2, 2Gh-1, 2W2-1
+            biases = biases.flatten(1).transpose(0, 1).contiguous().float()
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(group_size[0], device=attn.device)
+            coords_w = torch.arange(group_size[1], device=attn.device)
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Gh, Gw
+            coords_flatten = torch.flatten(coords, 1)  # 2, Gh*Gw
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Gh*Gw, Gh*Gw
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Gh*Gw, Gh*Gw, 2
+            relative_coords[:, :, 0] += group_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += group_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * group_size[1] - 1
+            relative_position_index = relative_coords.sum(-1)  # Gh*Gw, Gh*Gw
+
+            pos = self.pos(biases) # 2Gh-1 * 2Gw-1, heads
+            # select position bias
+            relative_position_bias = pos[relative_position_index.view(-1)].view( 
+                group_size[0] * group_size[1], group_size[0] * group_size[1], -1)  # Gh*Gw,Gh*Gw,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Gh*Gw, Gh*Gw
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nG = mask.shape[0]
+            attn = attn.view(B_ // nG, nG, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) # (B, nG, nHead, N, N)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        excluded_flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        excluded_flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        excluded_flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        if self.position_bias:
+            flops += self.pos.flops(N)
+        return flops, excluded_flops
+
+
+class CrossFormerBlock(nn.Module):
+    r""" CrossFormer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        group_size (int): Window size.
+        lsda_flag (int): use SDA or LDA, 0 for SDA and 1 for LDA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, group_size=7, interval=8, lsda_flag=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, num_patch_size=1):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.group_size = group_size
+        self.interval = interval
+        self.lsda_flag = lsda_flag
+        self.mlp_ratio = mlp_ratio
+        self.num_patch_size = num_patch_size
+
+        self.norm1 = norm_layer(dim)
+
+        self.attn = Attention(
+            dim, num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop,
+            position_bias=True)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, H, W):
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size %d, %d, %d" % (L, H, W)
+
+        if min(H, W) <= self.group_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.lsda_flag = 0
+            self.group_size = min(H, W)
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # padding
+        size_div = self.interval if self.lsda_flag == 1 else self.group_size
+        pad_l = pad_t = 0
+        pad_r = (size_div - W % size_div) % size_div
+        pad_b = (size_div - H % size_div) % size_div
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        mask = torch.zeros((1, Hp, Wp, 1), device=x.device)
+        if pad_b > 0:
+            mask[:, -pad_b:, :, :] = -1
+        if pad_r > 0:
+            mask[:, :, -pad_r:, :] = -1
+
+        # group embeddings and generate attn_mask
+        if self.lsda_flag == 0: # SDA
+            G = Gh = Gw = self.group_size
+            x = x.reshape(B, Hp // G, G, Wp // G, G, C).permute(0, 1, 3, 2, 4, 5).contiguous()
+            x = x.reshape(B * Hp * Wp // G**2, G**2, C)
+            nG = Hp * Wp // G**2
+            # attn_mask
+            if pad_r > 0 or pad_b > 0:
+                mask = mask.reshape(1, Hp // G, G, Wp // G, G, 1).permute(0, 1, 3, 2, 4, 5).contiguous()
+                mask = mask.reshape(nG, 1, G * G)
+                attn_mask = torch.zeros((nG, G * G, G * G), device=x.device)
+                attn_mask = attn_mask.masked_fill(mask < 0, NEG_INF)
+            else:
+                attn_mask = None
+        else: # LDA
+            I, Gh, Gw = self.interval, Hp // self.interval, Wp // self.interval
+            x = x.reshape(B, Gh, I, Gw, I, C).permute(0, 2, 4, 1, 3, 5).contiguous()
+            x = x.reshape(B * I * I, Gh * Gw, C)
+            nG = I ** 2
+            # attn_mask
+            if pad_r > 0 or pad_b > 0:
+                mask = mask.reshape(1, Gh, I, Gw, I, 1).permute(0, 2, 4, 1, 3, 5).contiguous()
+                mask = mask.reshape(nG, 1, Gh * Gw)
+                attn_mask = torch.zeros((nG, Gh * Gw, Gh * Gw), device=x.device)
+                attn_mask = attn_mask.masked_fill(mask < 0, NEG_INF)
+            else:
+                attn_mask = None
+
+        # multi-head self-attention
+        x = self.attn(x, Gh, Gw, mask=attn_mask)  # nG*B, G*G, C
+        
+        # ungroup embeddings
+        if self.lsda_flag == 0:
+            x = x.reshape(B, Hp // G, Wp // G, G, G, C).permute(0, 1, 3, 2, 4, 5).contiguous() # B, Hp//G, G, Wp//G, G, C
+        else:
+            x = x.reshape(B, I, I, Gh, Gw, C).permute(0, 3, 1, 4, 2, 5).contiguous() # B, Gh, I, Gw, I, C
+        x = x.reshape(B, Hp, Wp, C)
+
+        # remove padding
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"group_size={self.group_size}, lsda_flag={self.lsda_flag}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # Attention
+        size_div = self.interval if self.lsda_flag == 1 else self.group_size
+        Hp = math.ceil(H / size_div) * size_div
+        Wp = math.ceil(W / size_div) * size_div
+        Gh = Hp / size_div if self.lsda_flag == 1 else self.group_size
+        Gw = Wp / size_div if self.lsda_flag == 1 else self.group_size
+        nG = Hp * Wp / Gh / Gw
+        attn_flops, attn_excluded_flops = self.attn.flops(Gh * Gw)
+        flops += nG * attn_flops
+        excluded_flops = nG * attn_excluded_flops
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops, excluded_flops
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm, patch_size=[2], num_input_patch_size=1):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reductions = nn.ModuleList()
+        self.patch_size = patch_size
+        self.norm = norm_layer(dim)
+
+        for i, ps in enumerate(patch_size):
+            if i == len(patch_size) - 1:
+                out_dim = 2 * dim // 2 ** i
+            else:
+                out_dim = 2 * dim // 2 ** (i + 1)
+            stride = 2
+            padding = (ps - stride) // 2
+            self.reductions.append(nn.Conv2d(dim, out_dim, kernel_size=ps, 
+                                                stride=stride, padding=padding))
+
+    def forward(self, x, H, W):
+        """
+        x: B, H*W, C
+        """
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = self.norm(x)
+        x = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+
+        xs = []
+        for i in range(len(self.reductions)):
+            tmp_x = self.reductions[i](x).flatten(2).transpose(1, 2).contiguous()
+            xs.append(tmp_x)
+        x = torch.cat(xs, dim=2)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        for i, ps in enumerate(self.patch_size):
+            if i == len(self.patch_size) - 1:
+                out_dim = 2 * self.dim // 2 ** i
+            else:
+                out_dim = 2 * self.dim // 2 ** (i + 1)
+            flops += (H // 2) * (W // 2) * ps * ps * out_dim * self.dim
+        return flops
+
+
+class Stage(nn.Module):
+    """ CrossFormer blocks for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        group_size (int): Group size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Ghether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, group_size, interval,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 patch_size_end=[4], num_patch_size=None):
+
+        super().__init__()
+        self.dim = dim
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList()
+        for i in range(depth):
+            lsda_flag = 0 if (i % 2 == 0) else 1
+            self.blocks.append(CrossFormerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, group_size=group_size, interval=interval,
+                                 lsda_flag=lsda_flag,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer,
+                                 num_patch_size=num_patch_size))
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer, 
+                                         patch_size=patch_size_end, num_input_patch_size=num_patch_size)
+        else:
+            self.downsample = None
+
+    def forward(self, x, H, W):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x, H, W)
+
+        B, _, C = x.shape
+        feat = x.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+        if self.downsample is not None:
+            x = self.downsample(x, H, W)
+        return feat, x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        excluded_flops = 0
+        for blk in self.blocks:
+            blk_flops, blk_excluded_flops = blk.flops()
+            flops += blk_flops
+            excluded_flops += blk_excluded_flops
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops, excluded_flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=[4], in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        # patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // 4, img_size[1] // 4] # only for flops calculation
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.projs = nn.ModuleList()
+        for i, ps in enumerate(patch_size):
+            if i == len(patch_size) - 1:
+                dim = embed_dim // 2 ** i
+            else:
+                dim = embed_dim // 2 ** (i + 1)
+            stride = 4
+            padding = (ps - 4) // 2
+            self.projs.append(nn.Conv2d(in_chans, dim, kernel_size=ps, stride=stride, padding=padding))
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        xs = []
+        for i in range(len(self.projs)):
+            tx = self.projs[i](x).flatten(2).transpose(1, 2)
+            xs.append(tx)  # B Ph*Pw C
+        x = torch.cat(xs, dim=2)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x, H, W
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = 0
+        for i, ps in enumerate(self.patch_size):
+            if i == len(self.patch_size) - 1:
+                dim = self.embed_dim // 2 ** i
+            else:
+                dim = self.embed_dim // 2 ** (i + 1)
+            flops += Ho * Wo * dim * self.in_chans * (self.patch_size[i] * self.patch_size[i])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class CrossFormer(nn.Module):
+    r""" CrossFormer
+        A PyTorch impl of : `CrossFormer: A Versatile Vision Transformer Based on Cross-scale Attention`  -
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each stage.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        group_size (int): Group size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Ghether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=[4], in_chans=3, num_classes=1000,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 group_size=7, crs_interval=[8, 4, 2, 1], mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, patch_norm=True,
+                 use_checkpoint=False, merge_size=[[2], [2], [2]], **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution # [H//4, W//4] of original image size
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+
+        num_patch_sizes = [len(patch_size)] + [len(m) for m in merge_size]
+        for i_layer in range(self.num_layers):
+            patch_size_end = merge_size[i_layer] if i_layer < self.num_layers - 1 else None
+            num_patch_size = num_patch_sizes[i_layer]
+            layer = Stage(dim=int(embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               group_size=group_size[i_layer],
+                               interval=crs_interval[i_layer],
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=qkv_bias, qk_scale=qk_scale,
+                               drop=drop_rate, attn_drop=attn_drop_rate,
+                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                               norm_layer=norm_layer,
+                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               patch_size_end=patch_size_end,
+                               num_patch_size=num_patch_size)
+            self.layers.append(layer)
+
+        #  # classification
+        # self.norm = norm_layer(self.num_features)
+        # self.avgpool = nn.AdaptiveAvgPool1d(1)
+        # self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str):
+            logger = get_root_logger()
+            load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward(self, x):
+        x, H, W = self.patch_embed(x)
+        x = self.pos_drop(x)
+
+        outs = []
+        for i, layer in enumerate(self.layers):
+            feat, x = layer(x, H //4 //(2 ** i), W //4 //(2 ** i))
+            outs.append(feat)
+
+        # # classification
+        # x = self.norm(x)  # B L C
+        # x = self.avgpool(x.transpose(1, 2))  # B C 1
+        # x = torch.flatten(x, 1)
+        # x = self.head(x)
+        # return x
+
+        return outs
+
+    def flops(self):
+        flops = 0
+        excluded_flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            layer_flops, layer_excluded_flops = layer.flops()
+            flops += layer_flops
+            excluded_flops += layer_excluded_flops
+        # flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        # flops += self.num_features * self.num_classes
+        return flops, excluded_flops

PuzzleTuning/Backbone/getmodel.py

@@ -0,0 +1,392 @@
+"""
+get model func    Script  ver： Dec 5th 14:20
+"""
+import os
+import sys
+sys.path.append(os.path.realpath('.'))
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from Backbone import ResHybrid
+
+
+# get model
+def get_model(num_classes=1000, edge_size=224, model_idx=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0,
+              pretrained_backbone=True, use_cls_token=True, use_pos_embedding=True, use_att_module='SimAM'):
+    """
+    :param num_classes: classification required number of your dataset
+    :param edge_size: the input edge size of the dataloder
+    :param model_idx: the model we are going to use. by the format of Model_size_other_info
+
+    :param drop_rate: The dropout layer's probility of proposed models
+    :param attn_drop_rate: The dropout layer(right after the MHSA block or MHGA block)'s probility of proposed models
+    :param drop_path_rate: The probility of stochastic depth
+
+    :param pretrained_backbone: The backbone CNN is initiate randomly or by its official Pretrained models
+
+    :param use_cls_token: To use the class token
+    :param use_pos_embedding: To use the positional enbedding
+    :param use_att_module: To use which attention module in the FGD Focus block
+
+    :return: prepared model
+    """
+    if model_idx[0:5] == 'ViT_h':
+        # Transfer learning for ViT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit*')
+        pprint(model_names)
+        if edge_size == 224:
+            model = timm.create_model('vit_huge_patch14_224_in21k', pretrained=pretrained_backbone, num_classes=num_classes)
+        else:
+            print('not a avaliable image size with', model_idx)
+
+    elif model_idx[0:5] == 'ViT_l':
+        # Transfer learning for ViT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit*')
+        pprint(model_names)
+        if edge_size == 224:
+            model = timm.create_model('vit_large_patch16_224', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif edge_size == 384:
+            model = timm.create_model('vit_large_patch16_384', pretrained=pretrained_backbone, num_classes=num_classes)
+        else:
+            print('not a avaliable image size with', model_idx)
+
+    elif model_idx[0:5] == 'ViT_s':
+        # Transfer learning for ViT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit*')
+        pprint(model_names)
+        if edge_size == 224:
+            model = timm.create_model('vit_small_patch16_224', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif edge_size == 384:
+            model = timm.create_model('vit_small_patch16_384', pretrained=pretrained_backbone, num_classes=num_classes)
+        else:
+            print('not a avaliable image size with', model_idx)
+
+    elif model_idx[0:5] == 'ViT_t':
+        # Transfer learning for ViT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit*')
+        pprint(model_names)
+        if edge_size == 224:
+            model = timm.create_model('vit_tiny_patch16_224', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif edge_size == 384:
+            model = timm.create_model('vit_tiny_patch16_384', pretrained=pretrained_backbone, num_classes=num_classes)
+        else:
+            print('not a avaliable image size with', model_idx)
+
+    elif model_idx[0:5] == 'ViT_b' or model_idx[0:3] == 'ViT':  # vit_base
+        # Transfer learning for ViT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit*')
+        pprint(model_names)
+        if edge_size == 224:
+            model = timm.create_model('vit_base_patch16_224', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif edge_size == 384:
+            model = timm.create_model('vit_base_patch16_384', pretrained=pretrained_backbone, num_classes=num_classes)
+        else:
+            print('not a avaliable image size with', model_idx)
+
+    elif model_idx[0:3] == 'vgg':
+        # Transfer learning for vgg16_bn
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vgg*')
+        pprint(model_names)
+        if model_idx[0:8] == 'vgg16_bn':
+            model = timm.create_model('vgg16_bn', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif model_idx[0:5] == 'vgg16':
+            model = timm.create_model('vgg16', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif model_idx[0:8] == 'vgg19_bn':
+            model = timm.create_model('vgg19_bn', pretrained=pretrained_backbone, num_classes=num_classes)
+        elif model_idx[0:5] == 'vgg19':
+            model = timm.create_model('vgg19', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:4] == 'deit':  # Transfer learning for DeiT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*deit*')
+        pprint(model_names)
+        if edge_size == 384:
+            model = timm.create_model('deit_base_patch16_384', pretrained=pretrained_backbone, num_classes=2)
+        elif edge_size == 224:
+            model = timm.create_model('deit_base_patch16_224', pretrained=pretrained_backbone, num_classes=2)
+        else:
+            pass
+
+    elif model_idx[0:5] == 'twins':  # Transfer learning for twins
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*twins*')
+        pprint(model_names)
+        model = timm.create_model('twins_pcpvt_base', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:5] == 'pit_b' and edge_size == 224:  # Transfer learning for PiT
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*pit*')
+        pprint(model_names)
+        model = timm.create_model('pit_b_224', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:5] == 'gcvit' and edge_size == 224:  # Transfer learning for gcvit
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*gcvit*')
+        pprint(model_names)
+        model = timm.create_model('gcvit_base', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:6] == 'xcit_s':  # Transfer learning for XCiT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*xcit*')
+        pprint(model_names)
+        if edge_size == 384:
+            model = timm.create_model('xcit_small_12_p16_384_dist', pretrained=pretrained_backbone,
+                                      num_classes=num_classes)
+        elif edge_size == 224:
+            model = timm.create_model('xcit_small_12_p16_224_dist', pretrained=pretrained_backbone,
+                                      num_classes=num_classes)
+        else:
+            pass
+
+    elif model_idx[0:6] == 'xcit_m':  # Transfer learning for XCiT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*xcit*')
+        pprint(model_names)
+        if edge_size == 384:
+            model = timm.create_model('xcit_medium_24_p16_384_dist', pretrained=pretrained_backbone,
+                                      num_classes=num_classes)
+        elif edge_size == 224:
+            model = timm.create_model('xcit_medium_24_p16_224_dist', pretrained=pretrained_backbone,
+                                      num_classes=num_classes)
+        else:
+            pass
+
+    elif model_idx[0:6] == 'mvitv2':  # Transfer learning for MViT v2 small  fixme bug in model!
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*mvitv2*')
+        pprint(model_names)
+        model = timm.create_model('mvitv2_small_cls', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:6] == 'convit' and edge_size == 224:  # Transfer learning for ConViT fixme bug in model!
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*convit*')
+        pprint(model_names)
+        model = timm.create_model('convit_base', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:6] == 'ResNet':  # Transfer learning for the ResNets
+        if model_idx[0:8] == 'ResNet34':
+            model = models.resnet34(pretrained=pretrained_backbone)
+        elif model_idx[0:8] == 'ResNet50':
+            model = models.resnet50(pretrained=pretrained_backbone)
+        elif model_idx[0:9] == 'ResNet101':
+            model = models.resnet101(pretrained=pretrained_backbone)
+        else:
+            print('this model is not defined in get model')
+            return -1
+        num_ftrs = model.fc.in_features
+        model.fc = nn.Linear(num_ftrs, num_classes)
+
+    elif model_idx[0:6] == 'Backbone':  # ours: MSHT
+        # NOTICE: HERE 'pretrained' controls only The backbone CNN is initiate randomly
+        # or by its official Pretrained models
+        model = ResHybrid.create_model(model_idx, edge_size, pretrained=pretrained_backbone, num_classes=num_classes,
+                                       drop_rate=drop_rate, attn_drop_rate=attn_drop_rate,
+                                       drop_path_rate=drop_path_rate, use_cls_token=use_cls_token,
+                                       use_pos_embedding=use_pos_embedding, use_att_module=use_att_module)
+
+    elif model_idx[0:7] == 'bot_256' and edge_size == 256:  # Model: BoT
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*bot*')
+        pprint(model_names)
+        # NOTICE: we find no weight for BoT in timm
+        # ['botnet26t_256', 'botnet50ts_256', 'eca_botnext26ts_256']
+        model = timm.create_model('botnet26t_256', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:8] == 'densenet':  # Transfer learning for densenet
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*densenet*')
+        pprint(model_names)
+        model = timm.create_model('densenet121', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:8] == 'xception':  # Transfer learning for Xception
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*xception*')
+        pprint(model_names)
+        model = timm.create_model('xception', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:9] == 'pvt_v2_b0':  # Transfer learning for PVT v2 (todo not okey with torch summary)
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*pvt_v2*')
+        pprint(model_names)
+        model = timm.create_model('pvt_v2_b0', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:9] == 'visformer' and edge_size == 224:  # Transfer learning for Visformer
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*visformer*')
+        pprint(model_names)
+        model = timm.create_model('visformer_small', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:9] == 'conformer':  # Transfer learning for Conformer base
+        from Backbone.counterpart_models import conformer
+
+        embed_dim = 576
+        channel_ratio = 6
+
+        if pretrained_backbone:
+            model = conformer.Conformer(num_classes=1000, patch_size=16, channel_ratio=channel_ratio,
+                                        embed_dim=embed_dim, depth=12, num_heads=9, mlp_ratio=4, qkv_bias=True)
+            # this is the related path to <code>, not <Backbone>
+            save_model_path = '../saved_models/Conformer_base_patch16.pth'  # fixme model is downloaded at this path
+            # downloaded from official model state at https://github.com/pengzhiliang/Conformer
+            model.load_state_dict(torch.load(save_model_path), False)
+
+            model.trans_cls_head = nn.Linear(embed_dim, num_classes)
+            model.conv_cls_head = nn.Linear(int(256 * channel_ratio), num_classes)
+            model.cls_head = nn.Linear(int(2 * num_classes), num_classes)
+
+        else:
+            model = conformer.Conformer(num_classes=num_classes, patch_size=16, channel_ratio=channel_ratio,
+                                        embed_dim=embed_dim, depth=12, num_heads=9, mlp_ratio=4, qkv_bias=True)
+
+    elif model_idx[0:9] == 'coat_mini' and edge_size == 224:  # Transfer learning for coat_mini
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*coat*')
+        pprint(model_names)
+        model = timm.create_model('coat_mini', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:10] == 'swin_b_384' and edge_size == 384:  # Transfer learning for Swin Transformer (swin_b_384)
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*swin*')
+        pprint(model_names)  # swin_base_patch4_window12_384  swin_base_patch4_window12_384_in22k
+        model = timm.create_model('swin_base_patch4_window12_384', pretrained=pretrained_backbone,
+                                  num_classes=num_classes)
+
+    elif model_idx[0:10] == 'swin_b_224' and edge_size == 224:  # Transfer learning for Swin Transformer (swin_b_384)
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*swin*')
+        pprint(model_names)  # swin_base_patch4_window7_224  swin_base_patch4_window7_224_in22k
+        model = timm.create_model('swin_base_patch4_window7_224', pretrained=pretrained_backbone,
+                                  num_classes=num_classes)
+
+    elif model_idx[0:11] == 'mobilenetv3':  # Transfer learning for mobilenetv3
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*mobilenet*')
+        pprint(model_names)
+        model = timm.create_model('mobilenetv3_large_100', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:11] == 'mobilevit_s':  # Transfer learning for mobilevit_s
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*mobilevit*')
+        pprint(model_names)
+        model = timm.create_model('mobilevit_s', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:11] == 'inceptionv3':  # Transfer learning for Inception v3
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*inception*')
+        pprint(model_names)
+        model = timm.create_model('inception_v3', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:12] == 'cross_former' and edge_size == 224:  # Transfer learning for crossformer base
+        from Backbone.counterpart_models import crossformer
+        backbone = crossformer.CrossFormer(img_size=edge_size,
+                                           patch_size=[4, 8, 16, 32],
+                                           in_chans=3,
+                                           num_classes=0,  # get backbone only
+                                           embed_dim=96,
+                                           depths=[2, 2, 18, 2],
+                                           num_heads=[3, 6, 12, 24],
+                                           group_size=[7, 7, 7, 7],
+                                           mlp_ratio=4.,
+                                           qkv_bias=True,
+                                           qk_scale=None,
+                                           drop_rate=0.0,
+                                           drop_path_rate=0.3,
+                                           ape=False,
+                                           patch_norm=True,
+                                           use_checkpoint=False,
+                                           merge_size=[[2, 4], [2, 4], [2, 4]], )
+        if pretrained_backbone:
+            save_model_path = '../saved_models/crossformer-b.pth'  # fixme model is downloaded at this path
+            # downloaded from official model state at https://github.com/cheerss/CrossFormer
+            backbone.load_state_dict(torch.load(save_model_path)['model'], False)
+        model = crossformer.cross_former_cls_head_warp(backbone, num_classes)
+
+    elif model_idx[0:13] == 'crossvit_base':  # Transfer learning for crossvit_base  (todo not okey with torch summary)
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*crossvit_base*')
+        pprint(model_names)
+        model = timm.create_model('crossvit_base_240', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:14] == 'efficientnet_b':  # Transfer learning for efficientnet_b3,4
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*efficientnet*')
+        pprint(model_names)
+        model = timm.create_model(model_idx[0:15], pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:14] == 'ResN50_ViT_384':  # ResNet+ViT融合模型384
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*vit_base_resnet*')
+        pprint(model_names)
+        model = timm.create_model('vit_base_resnet50_384', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:15] == 'coat_lite_small' and edge_size == 224:  # Transfer learning for coat_lite_small
+        import timm
+        from pprint import pprint
+
+        model_names = timm.list_models('*coat*')
+        pprint(model_names)
+        model = timm.create_model('coat_lite_small', pretrained=pretrained_backbone, num_classes=num_classes)
+
+    elif model_idx[0:17] == 'efficientformer_l' and edge_size == 224:  # Transfer learning for efficientnet_b3,4
+        import timm
+        from pprint import pprint
+        model_names = timm.list_models('*efficientformer*')
+        pprint(model_names)
+        model = timm.create_model(model_idx[0:18], pretrained=pretrained_backbone, num_classes=num_classes)
+
+    else:
+        print('\nThe model', model_idx, 'with the edge size of', edge_size)
+        print("is not defined in the script！！", '\n')
+        return -1
+
+    try:
+        img = torch.randn(1, 3, edge_size, edge_size)
+        preds = model(img)  # (1, class_number)
+        print('test model output：', preds)
+    except:
+        print("Problem exist in the model defining process！！")
+        return -1
+    else:
+        print('model is ready now!')
+        return model

PuzzleTuning/Counterpart PreTrain Methods/ReadMe.md

@@ -0,0 +1,2 @@
+we have used MOCO-V3, Dino, MAE to pretrain the ViT-base-224 model.
+The official codes are implemented here.

PuzzleTuning/Counterpart PreTrain Methods/dino-main/LICENSE

@@ -0,0 +1,201 @@
+Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

PuzzleTuning/Counterpart PreTrain Methods/dino-main/README.md

@@ -0,0 +1,14 @@
+# Self-Supervised Vision Transformers with DINO
+
+The original repo of DINO could be found [here](https://github.com/facebookresearch/dino "DINO")
+
+Pip requirements:	timm == 0.4.9, PyTorch == 1.7.1, Torchvision == 0.8.2, Cuda == 11.0
+
+Typical BASH: 
+   ```console
+python -m torch.distributed.launch \
+--nproc_per_node=2 main_dino.py --arch vit_base --batch_size_per_gpu 128 \
+--lr 1.5e-4 --epochs 100 --data_path /root/autodl-tmp/All \
+--basic_state_dict /root/autodl-tmp/ViT_b16_224_Imagenet.pth \
+--num_workers 32 --output_dir the/path/of/CPIA
+   ```

PuzzleTuning/Counterpart PreTrain Methods/dino-main/eval_copy_detection.py

@@ -0,0 +1,301 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+import pickle
+import argparse
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+from torchvision import models as torchvision_models
+from torchvision import transforms as pth_transforms
+from PIL import Image, ImageFile
+import numpy as np
+
+import utils
+import vision_transformer as vits
+from eval_knn import extract_features
+
+
+class CopydaysDataset():
+    def __init__(self, basedir):
+        self.basedir = basedir
+        self.block_names = (
+            ['original', 'strong'] +
+            ['jpegqual/%d' % i for i in
+             [3, 5, 8, 10, 15, 20, 30, 50, 75]] +
+            ['crops/%d' % i for i in
+             [10, 15, 20, 30, 40, 50, 60, 70, 80]])
+        self.nblocks = len(self.block_names)
+
+        self.query_blocks = range(self.nblocks)
+        self.q_block_sizes = np.ones(self.nblocks, dtype=int) * 157
+        self.q_block_sizes[1] = 229
+        # search only among originals
+        self.database_blocks = [0]
+
+    def get_block(self, i):
+        dirname = self.basedir + '/' + self.block_names[i]
+        fnames = [dirname + '/' + fname
+                  for fname in sorted(os.listdir(dirname))
+                  if fname.endswith('.jpg')]
+        return fnames
+
+    def get_block_filenames(self, subdir_name):
+        dirname = self.basedir + '/' + subdir_name
+        return [fname
+                for fname in sorted(os.listdir(dirname))
+                if fname.endswith('.jpg')]
+
+    def eval_result(self, ids, distances):
+        j0 = 0
+        for i in range(self.nblocks):
+            j1 = j0 + self.q_block_sizes[i]
+            block_name = self.block_names[i]
+            I = ids[j0:j1]   # block size
+            sum_AP = 0
+            if block_name != 'strong':
+                # 1:1 mapping of files to names
+                positives_per_query = [[i] for i in range(j1 - j0)]
+            else:
+                originals = self.get_block_filenames('original')
+                strongs = self.get_block_filenames('strong')
+
+                # check if prefixes match
+                positives_per_query = [
+                    [j for j, bname in enumerate(originals)
+                     if bname[:4] == qname[:4]]
+                    for qname in strongs]
+
+            for qno, Iline in enumerate(I):
+                positives = positives_per_query[qno]
+                ranks = []
+                for rank, bno in enumerate(Iline):
+                    if bno in positives:
+                        ranks.append(rank)
+                sum_AP += score_ap_from_ranks_1(ranks, len(positives))
+
+            print("eval on %s mAP=%.3f" % (
+                block_name, sum_AP / (j1 - j0)))
+            j0 = j1
+
+
+# from the Holidays evaluation package
+def score_ap_from_ranks_1(ranks, nres):
+    """ Compute the average precision of one search.
+    ranks = ordered list of ranks of true positives
+    nres  = total number of positives in dataset
+    """
+
+    # accumulate trapezoids in PR-plot
+    ap = 0.0
+
+    # All have an x-size of:
+    recall_step = 1.0 / nres
+
+    for ntp, rank in enumerate(ranks):
+
+        # y-size on left side of trapezoid:
+        # ntp = nb of true positives so far
+        # rank = nb of retrieved items so far
+        if rank == 0:
+            precision_0 = 1.0
+        else:
+            precision_0 = ntp / float(rank)
+
+        # y-size on right side of trapezoid:
+        # ntp and rank are increased by one
+        precision_1 = (ntp + 1) / float(rank + 1)
+
+        ap += (precision_1 + precision_0) * recall_step / 2.0
+
+    return ap
+
+
+class ImgListDataset(torch.utils.data.Dataset):
+    def __init__(self, img_list, transform=None):
+        self.samples = img_list
+        self.transform = transform
+
+    def __getitem__(self, i):
+        with open(self.samples[i], 'rb') as f:
+            img = Image.open(f)
+            img = img.convert('RGB')
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, i
+
+    def __len__(self):
+        return len(self.samples)
+
+
+def is_image_file(s):
+    ext = s.split(".")[-1]
+    if ext in ['jpg', 'jpeg', 'png', 'ppm', 'bmp', 'pgm', 'tif', 'tiff', 'webp']:
+        return True
+    return False
+
+
+@torch.no_grad()
+def extract_features(image_list, model, args):
+    transform = pth_transforms.Compose([
+        pth_transforms.Resize((args.imsize, args.imsize), interpolation=3),
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    tempdataset = ImgListDataset(image_list, transform=transform)
+    data_loader = torch.utils.data.DataLoader(tempdataset, batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers, drop_last=False,
+        sampler=torch.utils.data.DistributedSampler(tempdataset, shuffle=False))
+    features = None
+    for samples, index in utils.MetricLogger(delimiter="  ").log_every(data_loader, 10):
+        samples, index = samples.cuda(non_blocking=True), index.cuda(non_blocking=True)
+        feats = model.get_intermediate_layers(samples, n=1)[0].clone()
+
+        cls_output_token = feats[:, 0, :]  #  [CLS] token
+        # GeM with exponent 4 for output patch tokens
+        b, h, w, d = len(samples), int(samples.shape[-2] / model.patch_embed.patch_size), int(samples.shape[-1] / model.patch_embed.patch_size), feats.shape[-1]
+        feats = feats[:, 1:, :].reshape(b, h, w, d)
+        feats = feats.clamp(min=1e-6).permute(0, 3, 1, 2)
+        feats = nn.functional.avg_pool2d(feats.pow(4), (h, w)).pow(1. / 4).reshape(b, -1)
+        # concatenate [CLS] token and GeM pooled patch tokens
+        feats = torch.cat((cls_output_token, feats), dim=1)
+
+        # init storage feature matrix
+        if dist.get_rank() == 0 and features is None:
+            features = torch.zeros(len(data_loader.dataset), feats.shape[-1])
+            if args.use_cuda:
+                features = features.cuda(non_blocking=True)
+
+        # get indexes from all processes
+        y_all = torch.empty(dist.get_world_size(), index.size(0), dtype=index.dtype, device=index.device)
+        y_l = list(y_all.unbind(0))
+        y_all_reduce = torch.distributed.all_gather(y_l, index, async_op=True)
+        y_all_reduce.wait()
+        index_all = torch.cat(y_l)
+
+        # share features between processes
+        feats_all = torch.empty(dist.get_world_size(), feats.size(0), feats.size(1),
+                                dtype=feats.dtype, device=feats.device)
+        output_l = list(feats_all.unbind(0))
+        output_all_reduce = torch.distributed.all_gather(output_l, feats, async_op=True)
+        output_all_reduce.wait()
+
+        # update storage feature matrix
+        if dist.get_rank() == 0:
+            if args.use_cuda:
+                features.index_copy_(0, index_all, torch.cat(output_l))
+            else:
+                features.index_copy_(0, index_all.cpu(), torch.cat(output_l).cpu())
+    return features  # features is still None for every rank which is not 0 (main)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Copy detection on Copydays')
+    parser.add_argument('--data_path', default='/path/to/copydays/', type=str,
+        help="See https://lear.inrialpes.fr/~jegou/data.php#copydays")
+    parser.add_argument('--whitening_path', default='/path/to/whitening_data/', type=str,
+        help="""Path to directory with images used for computing the whitening operator.
+        In our paper, we use 20k random images from YFCC100M.""")
+    parser.add_argument('--distractors_path', default='/path/to/distractors/', type=str,
+        help="Path to directory with distractors images. In our paper, we use 10k random images from YFCC100M.")
+    parser.add_argument('--imsize', default=320, type=int, help='Image size (square image)')
+    parser.add_argument('--batch_size_per_gpu', default=16, type=int, help='Per-GPU batch-size')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag)
+    parser.add_argument('--arch', default='vit_base', type=str, help='Architecture')
+    parser.add_argument('--patch_size', default=8, type=int, help='Patch resolution of the model.')
+    parser.add_argument("--checkpoint_key", default="teacher", type=str,
+        help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    args = parser.parse_args()
+
+    utils.init_distributed_mode(args)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    # ============ building network ... ============
+    if "vit" in args.arch:
+        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
+    else:
+        print(f"Architecture {args.arch} non supported")
+        sys.exit(1)
+    if args.use_cuda:
+        model.cuda()
+    model.eval()
+    utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
+
+    dataset = CopydaysDataset(args.data_path)
+
+    # ============ Extract features ... ============
+    # extract features for queries
+    queries = []
+    for q in dataset.query_blocks:
+        queries.append(extract_features(dataset.get_block(q), model, args))
+    if utils.get_rank() == 0:
+        queries = torch.cat(queries)
+        print(f"Extraction of queries features done. Shape: {queries.shape}")
+
+    # extract features for database
+    database = []
+    for b in dataset.database_blocks:
+        database.append(extract_features(dataset.get_block(b), model, args))
+
+    # extract features for distractors
+    if os.path.isdir(args.distractors_path):
+        print("Using distractors...")
+        list_distractors = [os.path.join(args.distractors_path, s) for s in os.listdir(args.distractors_path) if is_image_file(s)]
+        database.append(extract_features(list_distractors, model, args))
+    if utils.get_rank() == 0:
+        database = torch.cat(database)
+        print(f"Extraction of database and distractors features done. Shape: {database.shape}")
+
+    # ============ Whitening ... ============
+    if os.path.isdir(args.whitening_path):
+        print(f"Extracting features on images from {args.whitening_path} for learning the whitening operator.")
+        list_whit = [os.path.join(args.whitening_path, s) for s in os.listdir(args.whitening_path) if is_image_file(s)]
+        features_for_whitening = extract_features(list_whit, model, args)
+        if utils.get_rank() == 0:
+            # center
+            mean_feature = torch.mean(features_for_whitening, dim=0)
+            database -= mean_feature
+            queries -= mean_feature
+            pca = utils.PCA(dim=database.shape[-1], whit=0.5)
+            # compute covariance
+            cov = torch.mm(features_for_whitening.T, features_for_whitening) / features_for_whitening.shape[0]
+            pca.train_pca(cov.cpu().numpy())
+            database = pca.apply(database)
+            queries = pca.apply(queries)
+
+    # ============ Copy detection ... ============
+    if utils.get_rank() == 0:
+        # l2 normalize the features
+        database = nn.functional.normalize(database, dim=1, p=2)
+        queries = nn.functional.normalize(queries, dim=1, p=2)
+
+        # similarity
+        similarity = torch.mm(queries, database.T)
+        distances, indices = similarity.topk(20, largest=True, sorted=True)
+
+        # evaluate
+        retrieved = dataset.eval_result(indices, distances)
+    dist.barrier()
+

PuzzleTuning/Counterpart PreTrain Methods/dino-main/eval_image_retrieval.py

@@ -0,0 +1,201 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+import pickle
+import argparse
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+from torchvision import models as torchvision_models
+from torchvision import transforms as pth_transforms
+from PIL import Image, ImageFile
+import numpy as np
+
+import utils
+import vision_transformer as vits
+from eval_knn import extract_features
+
+
+class OxfordParisDataset(torch.utils.data.Dataset):
+    def __init__(self, dir_main, dataset, split, transform=None, imsize=None):
+        if dataset not in ['roxford5k', 'rparis6k']:
+            raise ValueError('Unknown dataset: {}!'.format(dataset))
+
+        # loading imlist, qimlist, and gnd, in cfg as a dict
+        gnd_fname = os.path.join(dir_main, dataset, 'gnd_{}.pkl'.format(dataset))
+        with open(gnd_fname, 'rb') as f:
+            cfg = pickle.load(f)
+        cfg['gnd_fname'] = gnd_fname
+        cfg['ext'] = '.jpg'
+        cfg['qext'] = '.jpg'
+        cfg['dir_data'] = os.path.join(dir_main, dataset)
+        cfg['dir_images'] = os.path.join(cfg['dir_data'], 'jpg')
+        cfg['n'] = len(cfg['imlist'])
+        cfg['nq'] = len(cfg['qimlist'])
+        cfg['im_fname'] = config_imname
+        cfg['qim_fname'] = config_qimname
+        cfg['dataset'] = dataset
+        self.cfg = cfg
+
+        self.samples = cfg["qimlist"] if split == "query" else cfg["imlist"]
+        self.transform = transform
+        self.imsize = imsize
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        path = os.path.join(self.cfg["dir_images"], self.samples[index] + ".jpg")
+        ImageFile.LOAD_TRUNCATED_IMAGES = True
+        with open(path, 'rb') as f:
+            img = Image.open(f)
+            img = img.convert('RGB')
+        if self.imsize is not None:
+            img.thumbnail((self.imsize, self.imsize), Image.ANTIALIAS)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, index
+
+
+def config_imname(cfg, i):
+    return os.path.join(cfg['dir_images'], cfg['imlist'][i] + cfg['ext'])
+
+
+def config_qimname(cfg, i):
+    return os.path.join(cfg['dir_images'], cfg['qimlist'][i] + cfg['qext'])
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Image Retrieval on revisited Paris and Oxford')
+    parser.add_argument('--data_path', default='/path/to/revisited_paris_oxford/', type=str)
+    parser.add_argument('--dataset', default='roxford5k', type=str, choices=['roxford5k', 'rparis6k'])
+    parser.add_argument('--multiscale', default=False, type=utils.bool_flag)
+    parser.add_argument('--imsize', default=224, type=int, help='Image size')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag)
+    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
+    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
+    parser.add_argument("--checkpoint_key", default="teacher", type=str,
+        help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    args = parser.parse_args()
+
+    utils.init_distributed_mode(args)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    # ============ preparing data ... ============
+    transform = pth_transforms.Compose([
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    dataset_train = OxfordParisDataset(args.data_path, args.dataset, split="train", transform=transform, imsize=args.imsize)
+    dataset_query = OxfordParisDataset(args.data_path, args.dataset, split="query", transform=transform, imsize=args.imsize)
+    sampler = torch.utils.data.DistributedSampler(dataset_train, shuffle=False)
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler,
+        batch_size=1,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    data_loader_query = torch.utils.data.DataLoader(
+        dataset_query,
+        batch_size=1,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    print(f"train: {len(dataset_train)} imgs / query: {len(dataset_query)} imgs")
+
+    # ============ building network ... ============
+    if "vit" in args.arch:
+        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
+    elif "xcit" in args.arch:
+        model = torch.hub.load('facebookresearch/xcit:main', args.arch, num_classes=0)
+    elif args.arch in torchvision_models.__dict__.keys():
+        model = torchvision_models.__dict__[args.arch](num_classes=0)
+    else:
+        print(f"Architecture {args.arch} non supported")
+        sys.exit(1)
+    if args.use_cuda:
+        model.cuda()
+    model.eval()
+
+    # load pretrained weights
+    if os.path.isfile(args.pretrained_weights):
+        state_dict = torch.load(args.pretrained_weights, map_location="cpu")
+        if args.checkpoint_key is not None and args.checkpoint_key in state_dict:
+            print(f"Take key {args.checkpoint_key} in provided checkpoint dict")
+            state_dict = state_dict[args.checkpoint_key]
+        # remove `module.` prefix
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        # remove `backbone.` prefix induced by multicrop wrapper
+        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+        msg = model.load_state_dict(state_dict, strict=False)
+        print('Pretrained weights found at {} and loaded with msg: {}'.format(args.pretrained_weights, msg))
+    elif args.arch == "vit_small" and args.patch_size == 16:
+        print("Since no pretrained weights have been provided, we load pretrained DINO weights on Google Landmark v2.")
+        model.load_state_dict(torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/dino_vitsmall16_googlelandmark_pretrain/dino_vitsmall16_googlelandmark_pretrain.pth"))
+    else:
+        print("Warning: We use random weights.")
+
+    ############################################################################
+    # Step 1: extract features
+    train_features = extract_features(model, data_loader_train, args.use_cuda, multiscale=args.multiscale)
+    query_features = extract_features(model, data_loader_query, args.use_cuda, multiscale=args.multiscale)
+
+    if utils.get_rank() == 0:  # only rank 0 will work from now on
+        # normalize features
+        train_features = nn.functional.normalize(train_features, dim=1, p=2)
+        query_features = nn.functional.normalize(query_features, dim=1, p=2)
+
+        ############################################################################
+        # Step 2: similarity
+        sim = torch.mm(train_features, query_features.T)
+        ranks = torch.argsort(-sim, dim=0).cpu().numpy()
+
+        ############################################################################
+        # Step 3: evaluate
+        gnd = dataset_train.cfg['gnd']
+        # evaluate ranks
+        ks = [1, 5, 10]
+        # search for easy & hard
+        gnd_t = []
+        for i in range(len(gnd)):
+            g = {}
+            g['ok'] = np.concatenate([gnd[i]['easy'], gnd[i]['hard']])
+            g['junk'] = np.concatenate([gnd[i]['junk']])
+            gnd_t.append(g)
+        mapM, apsM, mprM, prsM = utils.compute_map(ranks, gnd_t, ks)
+        # search for hard
+        gnd_t = []
+        for i in range(len(gnd)):
+            g = {}
+            g['ok'] = np.concatenate([gnd[i]['hard']])
+            g['junk'] = np.concatenate([gnd[i]['junk'], gnd[i]['easy']])
+            gnd_t.append(g)
+        mapH, apsH, mprH, prsH = utils.compute_map(ranks, gnd_t, ks)
+        print('>> {}: mAP M: {}, H: {}'.format(args.dataset, np.around(mapM*100, decimals=2), np.around(mapH*100, decimals=2)))
+        print('>> {}: mP@k{} M: {}, H: {}'.format(args.dataset, np.array(ks), np.around(mprM*100, decimals=2), np.around(mprH*100, decimals=2)))
+    dist.barrier()

PuzzleTuning/Counterpart PreTrain Methods/dino-main/eval_knn.py

@@ -0,0 +1,242 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+import argparse
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+from torchvision import datasets
+from torchvision import transforms as pth_transforms
+from torchvision import models as torchvision_models
+
+import utils
+import vision_transformer as vits
+
+
+def extract_feature_pipeline(args):
+    # ============ preparing data ... ============
+    transform = pth_transforms.Compose([
+        pth_transforms.Resize(256, interpolation=3),
+        pth_transforms.CenterCrop(224),
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    dataset_train = ReturnIndexDataset(os.path.join(args.data_path, "train"), transform=transform)
+    dataset_val = ReturnIndexDataset(os.path.join(args.data_path, "val"), transform=transform)
+    sampler = torch.utils.data.DistributedSampler(dataset_train, shuffle=False)
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    data_loader_val = torch.utils.data.DataLoader(
+        dataset_val,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=False,
+    )
+    print(f"Data loaded with {len(dataset_train)} train and {len(dataset_val)} val imgs.")
+
+    # ============ building network ... ============
+    if "vit" in args.arch:
+        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+        print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
+    elif "xcit" in args.arch:
+        model = torch.hub.load('facebookresearch/xcit:main', args.arch, num_classes=0)
+    elif args.arch in torchvision_models.__dict__.keys():
+        model = torchvision_models.__dict__[args.arch](num_classes=0)
+        model.fc = nn.Identity()
+    else:
+        print(f"Architecture {args.arch} non supported")
+        sys.exit(1)
+    model.cuda()
+    utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
+    model.eval()
+
+    # ============ extract features ... ============
+    print("Extracting features for train set...")
+    train_features = extract_features(model, data_loader_train, args.use_cuda)
+    print("Extracting features for val set...")
+    test_features = extract_features(model, data_loader_val, args.use_cuda)
+
+    if utils.get_rank() == 0:
+        train_features = nn.functional.normalize(train_features, dim=1, p=2)
+        test_features = nn.functional.normalize(test_features, dim=1, p=2)
+
+    train_labels = torch.tensor([s[-1] for s in dataset_train.samples]).long()
+    test_labels = torch.tensor([s[-1] for s in dataset_val.samples]).long()
+    # save features and labels
+    if args.dump_features and dist.get_rank() == 0:
+        torch.save(train_features.cpu(), os.path.join(args.dump_features, "trainfeat.pth"))
+        torch.save(test_features.cpu(), os.path.join(args.dump_features, "testfeat.pth"))
+        torch.save(train_labels.cpu(), os.path.join(args.dump_features, "trainlabels.pth"))
+        torch.save(test_labels.cpu(), os.path.join(args.dump_features, "testlabels.pth"))
+    return train_features, test_features, train_labels, test_labels
+
+
+@torch.no_grad()
+def extract_features(model, data_loader, use_cuda=True, multiscale=False):
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    features = None
+    for samples, index in metric_logger.log_every(data_loader, 10):
+        samples = samples.cuda(non_blocking=True)
+        index = index.cuda(non_blocking=True)
+        if multiscale:
+            feats = utils.multi_scale(samples, model)
+        else:
+            feats = model(samples).clone()
+
+        # init storage feature matrix
+        if dist.get_rank() == 0 and features is None:
+            features = torch.zeros(len(data_loader.dataset), feats.shape[-1])
+            if use_cuda:
+                features = features.cuda(non_blocking=True)
+            print(f"Storing features into tensor of shape {features.shape}")
+
+        # get indexes from all processes
+        y_all = torch.empty(dist.get_world_size(), index.size(0), dtype=index.dtype, device=index.device)
+        y_l = list(y_all.unbind(0))
+        y_all_reduce = torch.distributed.all_gather(y_l, index, async_op=True)
+        y_all_reduce.wait()
+        index_all = torch.cat(y_l)
+
+        # share features between processes
+        feats_all = torch.empty(
+            dist.get_world_size(),
+            feats.size(0),
+            feats.size(1),
+            dtype=feats.dtype,
+            device=feats.device,
+        )
+        output_l = list(feats_all.unbind(0))
+        output_all_reduce = torch.distributed.all_gather(output_l, feats, async_op=True)
+        output_all_reduce.wait()
+
+        # update storage feature matrix
+        if dist.get_rank() == 0:
+            if use_cuda:
+                features.index_copy_(0, index_all, torch.cat(output_l))
+            else:
+                features.index_copy_(0, index_all.cpu(), torch.cat(output_l).cpu())
+    return features
+
+
+@torch.no_grad()
+def knn_classifier(train_features, train_labels, test_features, test_labels, k, T, num_classes=1000):
+    top1, top5, total = 0.0, 0.0, 0
+    train_features = train_features.t()
+    num_test_images, num_chunks = test_labels.shape[0], 100
+    imgs_per_chunk = num_test_images // num_chunks
+    retrieval_one_hot = torch.zeros(k, num_classes).to(train_features.device)
+    for idx in range(0, num_test_images, imgs_per_chunk):
+        # get the features for test images
+        features = test_features[
+            idx : min((idx + imgs_per_chunk), num_test_images), :
+        ]
+        targets = test_labels[idx : min((idx + imgs_per_chunk), num_test_images)]
+        batch_size = targets.shape[0]
+
+        # calculate the dot product and compute top-k neighbors
+        similarity = torch.mm(features, train_features)
+        distances, indices = similarity.topk(k, largest=True, sorted=True)
+        candidates = train_labels.view(1, -1).expand(batch_size, -1)
+        retrieved_neighbors = torch.gather(candidates, 1, indices)
+
+        retrieval_one_hot.resize_(batch_size * k, num_classes).zero_()
+        retrieval_one_hot.scatter_(1, retrieved_neighbors.view(-1, 1), 1)
+        distances_transform = distances.clone().div_(T).exp_()
+        probs = torch.sum(
+            torch.mul(
+                retrieval_one_hot.view(batch_size, -1, num_classes),
+                distances_transform.view(batch_size, -1, 1),
+            ),
+            1,
+        )
+        _, predictions = probs.sort(1, True)
+
+        # find the predictions that match the target
+        correct = predictions.eq(targets.data.view(-1, 1))
+        top1 = top1 + correct.narrow(1, 0, 1).sum().item()
+        top5 = top5 + correct.narrow(1, 0, min(5, k)).sum().item()  # top5 does not make sense if k < 5
+        total += targets.size(0)
+    top1 = top1 * 100.0 / total
+    top5 = top5 * 100.0 / total
+    return top1, top5
+
+
+class ReturnIndexDataset(datasets.ImageFolder):
+    def __getitem__(self, idx):
+        img, lab = super(ReturnIndexDataset, self).__getitem__(idx)
+        return img, idx
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Evaluation with weighted k-NN on ImageNet')
+    parser.add_argument('--batch_size_per_gpu', default=128, type=int, help='Per-GPU batch-size')
+    parser.add_argument('--nb_knn', default=[10, 20, 100, 200], nargs='+', type=int,
+        help='Number of NN to use. 20 is usually working the best.')
+    parser.add_argument('--temperature', default=0.07, type=float,
+        help='Temperature used in the voting coefficient')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument('--use_cuda', default=True, type=utils.bool_flag,
+        help="Should we store the features on GPU? We recommend setting this to False if you encounter OOM")
+    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
+    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
+    parser.add_argument("--checkpoint_key", default="teacher", type=str,
+        help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--dump_features', default=None,
+        help='Path where to save computed features, empty for no saving')
+    parser.add_argument('--load_features', default=None, help="""If the features have
+        already been computed, where to find them.""")
+    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    parser.add_argument('--data_path', default='/path/to/imagenet/', type=str)
+    args = parser.parse_args()
+
+    utils.init_distributed_mode(args)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    if args.load_features:
+        train_features = torch.load(os.path.join(args.load_features, "trainfeat.pth"))
+        test_features = torch.load(os.path.join(args.load_features, "testfeat.pth"))
+        train_labels = torch.load(os.path.join(args.load_features, "trainlabels.pth"))
+        test_labels = torch.load(os.path.join(args.load_features, "testlabels.pth"))
+    else:
+        # need to extract features !
+        train_features, test_features, train_labels, test_labels = extract_feature_pipeline(args)
+
+    if utils.get_rank() == 0:
+        if args.use_cuda:
+            train_features = train_features.cuda()
+            test_features = test_features.cuda()
+            train_labels = train_labels.cuda()
+            test_labels = test_labels.cuda()
+
+        print("Features are ready!\nStart the k-NN classification.")
+        for k in args.nb_knn:
+            top1, top5 = knn_classifier(train_features, train_labels,
+                test_features, test_labels, k, args.temperature)
+            print(f"{k}-NN classifier result: Top1: {top1}, Top5: {top5}")
+    dist.barrier()

PuzzleTuning/Counterpart PreTrain Methods/dino-main/eval_linear.py

@@ -0,0 +1,281 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from torch import nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+from torchvision import datasets
+from torchvision import transforms as pth_transforms
+from torchvision import models as torchvision_models
+
+import utils
+import vision_transformer as vits
+
+
+def eval_linear(args):
+    utils.init_distributed_mode(args)
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    # ============ building network ... ============
+    # if the network is a Vision Transformer (i.e. vit_tiny, vit_small, vit_base)
+    if args.arch in vits.__dict__.keys():
+        model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+        embed_dim = model.embed_dim * (args.n_last_blocks + int(args.avgpool_patchtokens))
+    # if the network is a XCiT
+    elif "xcit" in args.arch:
+        model = torch.hub.load('facebookresearch/xcit:main', args.arch, num_classes=0)
+        embed_dim = model.embed_dim
+    # otherwise, we check if the architecture is in torchvision models
+    elif args.arch in torchvision_models.__dict__.keys():
+        model = torchvision_models.__dict__[args.arch]()
+        embed_dim = model.fc.weight.shape[1]
+        model.fc = nn.Identity()
+    else:
+        print(f"Unknow architecture: {args.arch}")
+        sys.exit(1)
+    model.cuda()
+    model.eval()
+    # load weights to evaluate
+    utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
+    print(f"Model {args.arch} built.")
+
+    linear_classifier = LinearClassifier(embed_dim, num_labels=args.num_labels)
+    linear_classifier = linear_classifier.cuda()
+    linear_classifier = nn.parallel.DistributedDataParallel(linear_classifier, device_ids=[args.gpu])
+
+    # ============ preparing data ... ============
+    val_transform = pth_transforms.Compose([
+        pth_transforms.Resize(256, interpolation=3),
+        pth_transforms.CenterCrop(224),
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    dataset_val = datasets.ImageFolder(os.path.join(args.data_path, "val"), transform=val_transform)
+    val_loader = torch.utils.data.DataLoader(
+        dataset_val,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    if args.evaluate:
+        utils.load_pretrained_linear_weights(linear_classifier, args.arch, args.patch_size)
+        test_stats = validate_network(val_loader, model, linear_classifier, args.n_last_blocks, args.avgpool_patchtokens)
+        print(f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
+        return
+
+    train_transform = pth_transforms.Compose([
+        pth_transforms.RandomResizedCrop(224),
+        pth_transforms.RandomHorizontalFlip(),
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    dataset_train = datasets.ImageFolder(os.path.join(args.data_path, "train"), transform=train_transform)
+    sampler = torch.utils.data.distributed.DistributedSampler(dataset_train)
+    train_loader = torch.utils.data.DataLoader(
+        dataset_train,
+        sampler=sampler,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+    print(f"Data loaded with {len(dataset_train)} train and {len(dataset_val)} val imgs.")
+
+    # set optimizer
+    optimizer = torch.optim.SGD(
+        linear_classifier.parameters(),
+        args.lr * (args.batch_size_per_gpu * utils.get_world_size()) / 256., # linear scaling rule
+        momentum=0.9,
+        weight_decay=0, # we do not apply weight decay
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min=0)
+
+    # Optionally resume from a checkpoint
+    to_restore = {"epoch": 0, "best_acc": 0.}
+    utils.restart_from_checkpoint(
+        os.path.join(args.output_dir, "checkpoint.pth.tar"),
+        run_variables=to_restore,
+        state_dict=linear_classifier,
+        optimizer=optimizer,
+        scheduler=scheduler,
+    )
+    start_epoch = to_restore["epoch"]
+    best_acc = to_restore["best_acc"]
+
+    for epoch in range(start_epoch, args.epochs):
+        train_loader.sampler.set_epoch(epoch)
+
+        train_stats = train(model, linear_classifier, optimizer, train_loader, epoch, args.n_last_blocks, args.avgpool_patchtokens)
+        scheduler.step()
+
+        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                     'epoch': epoch}
+        if epoch % args.val_freq == 0 or epoch == args.epochs - 1:
+            test_stats = validate_network(val_loader, model, linear_classifier, args.n_last_blocks, args.avgpool_patchtokens)
+            print(f"Accuracy at epoch {epoch} of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%")
+            best_acc = max(best_acc, test_stats["acc1"])
+            print(f'Max accuracy so far: {best_acc:.2f}%')
+            log_stats = {**{k: v for k, v in log_stats.items()},
+                         **{f'test_{k}': v for k, v in test_stats.items()}}
+        if utils.is_main_process():
+            with (Path(args.output_dir) / "log.txt").open("a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+            save_dict = {
+                "epoch": epoch + 1,
+                "state_dict": linear_classifier.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "scheduler": scheduler.state_dict(),
+                "best_acc": best_acc,
+            }
+            torch.save(save_dict, os.path.join(args.output_dir, "checkpoint.pth.tar"))
+    print("Training of the supervised linear classifier on frozen features completed.\n"
+                "Top-1 test accuracy: {acc:.1f}".format(acc=best_acc))
+
+
+def train(model, linear_classifier, optimizer, loader, epoch, n, avgpool):
+    linear_classifier.train()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    for (inp, target) in metric_logger.log_every(loader, 20, header):
+        # move to gpu
+        inp = inp.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+
+        # forward
+        with torch.no_grad():
+            if "vit" in args.arch:
+                intermediate_output = model.get_intermediate_layers(inp, n)
+                output = torch.cat([x[:, 0] for x in intermediate_output], dim=-1)
+                if avgpool:
+                    output = torch.cat((output.unsqueeze(-1), torch.mean(intermediate_output[-1][:, 1:], dim=1).unsqueeze(-1)), dim=-1)
+                    output = output.reshape(output.shape[0], -1)
+            else:
+                output = model(inp)
+        output = linear_classifier(output)
+
+        # compute cross entropy loss
+        loss = nn.CrossEntropyLoss()(output, target)
+
+        # compute the gradients
+        optimizer.zero_grad()
+        loss.backward()
+
+        # step
+        optimizer.step()
+
+        # log 
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss.item())
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def validate_network(val_loader, model, linear_classifier, n, avgpool):
+    linear_classifier.eval()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+    for inp, target in metric_logger.log_every(val_loader, 20, header):
+        # move to gpu
+        inp = inp.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
+
+        # forward
+        with torch.no_grad():
+            if "vit" in args.arch:
+                intermediate_output = model.get_intermediate_layers(inp, n)
+                output = torch.cat([x[:, 0] for x in intermediate_output], dim=-1)
+                if avgpool:
+                    output = torch.cat((output.unsqueeze(-1), torch.mean(intermediate_output[-1][:, 1:], dim=1).unsqueeze(-1)), dim=-1)
+                    output = output.reshape(output.shape[0], -1)
+            else:
+                output = model(inp)
+        output = linear_classifier(output)
+        loss = nn.CrossEntropyLoss()(output, target)
+
+        if linear_classifier.module.num_labels >= 5:
+            acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
+        else:
+            acc1, = utils.accuracy(output, target, topk=(1,))
+
+        batch_size = inp.shape[0]
+        metric_logger.update(loss=loss.item())
+        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        if linear_classifier.module.num_labels >= 5:
+            metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+    if linear_classifier.module.num_labels >= 5:
+        print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+    else:
+        print('* Acc@1 {top1.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(top1=metric_logger.acc1, losses=metric_logger.loss))
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+class LinearClassifier(nn.Module):
+    """Linear layer to train on top of frozen features"""
+    def __init__(self, dim, num_labels=1000):
+        super(LinearClassifier, self).__init__()
+        self.num_labels = num_labels
+        self.linear = nn.Linear(dim, num_labels)
+        self.linear.weight.data.normal_(mean=0.0, std=0.01)
+        self.linear.bias.data.zero_()
+
+    def forward(self, x):
+        # flatten
+        x = x.view(x.size(0), -1)
+
+        # linear layer
+        return self.linear(x)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Evaluation with linear classification on ImageNet')
+    parser.add_argument('--n_last_blocks', default=4, type=int, help="""Concatenate [CLS] tokens
+        for the `n` last blocks. We use `n=4` when evaluating ViT-Small and `n=1` with ViT-Base.""")
+    parser.add_argument('--avgpool_patchtokens', default=False, type=utils.bool_flag,
+        help="""Whether ot not to concatenate the global average pooled features to the [CLS] token.
+        We typically set this to False for ViT-Small and to True with ViT-Base.""")
+    parser.add_argument('--arch', default='vit_small', type=str, help='Architecture')
+    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument("--checkpoint_key", default="teacher", type=str, help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--epochs', default=100, type=int, help='Number of epochs of training.')
+    parser.add_argument("--lr", default=0.001, type=float, help="""Learning rate at the beginning of
+        training (highest LR used during training). The learning rate is linearly scaled
+        with the batch size, and specified here for a reference batch size of 256.
+        We recommend tweaking the LR depending on the checkpoint evaluated.""")
+    parser.add_argument('--batch_size_per_gpu', default=128, type=int, help='Per-GPU batch-size')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    parser.add_argument('--data_path', default='/path/to/imagenet/', type=str)
+    parser.add_argument('--num_workers', default=10, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument('--val_freq', default=1, type=int, help="Epoch frequency for validation.")
+    parser.add_argument('--output_dir', default=".", help='Path to save logs and checkpoints')
+    parser.add_argument('--num_labels', default=1000, type=int, help='Number of labels for linear classifier')
+    parser.add_argument('--evaluate', dest='evaluate', action='store_true', help='evaluate model on validation set')
+    args = parser.parse_args()
+    eval_linear(args)

PuzzleTuning/Counterpart PreTrain Methods/dino-main/eval_video_segmentation.py

@@ -0,0 +1,292 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Some parts are taken from https://github.com/Liusifei/UVC
+"""
+import os
+import copy
+import glob
+import queue
+from urllib.request import urlopen
+import argparse
+import numpy as np
+from tqdm import tqdm
+
+import cv2
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from PIL import Image
+from torchvision import transforms
+
+import utils
+import vision_transformer as vits
+
+
+@torch.no_grad()
+def eval_video_tracking_davis(args, model, frame_list, video_dir, first_seg, seg_ori, color_palette):
+    """
+    Evaluate tracking on a video given first frame & segmentation
+    """
+    video_folder = os.path.join(args.output_dir, video_dir.split('/')[-1])
+    os.makedirs(video_folder, exist_ok=True)
+
+    # The queue stores the n preceeding frames
+    que = queue.Queue(args.n_last_frames)
+
+    # first frame
+    frame1, ori_h, ori_w = read_frame(frame_list[0])
+    # extract first frame feature
+    frame1_feat = extract_feature(model, frame1).T #  dim x h*w
+
+    # saving first segmentation
+    out_path = os.path.join(video_folder, "00000.png")
+    imwrite_indexed(out_path, seg_ori, color_palette)
+    mask_neighborhood = None
+    for cnt in tqdm(range(1, len(frame_list))):
+        frame_tar = read_frame(frame_list[cnt])[0]
+
+        # we use the first segmentation and the n previous ones
+        used_frame_feats = [frame1_feat] + [pair[0] for pair in list(que.queue)]
+        used_segs = [first_seg] + [pair[1] for pair in list(que.queue)]
+
+        frame_tar_avg, feat_tar, mask_neighborhood = label_propagation(args, model, frame_tar, used_frame_feats, used_segs, mask_neighborhood)
+
+        # pop out oldest frame if neccessary
+        if que.qsize() == args.n_last_frames:
+            que.get()
+        # push current results into queue
+        seg = copy.deepcopy(frame_tar_avg)
+        que.put([feat_tar, seg])
+
+        # upsampling & argmax
+        frame_tar_avg = F.interpolate(frame_tar_avg, scale_factor=args.patch_size, mode='bilinear', align_corners=False, recompute_scale_factor=False)[0]
+        frame_tar_avg = norm_mask(frame_tar_avg)
+        _, frame_tar_seg = torch.max(frame_tar_avg, dim=0)
+
+        # saving to disk
+        frame_tar_seg = np.array(frame_tar_seg.squeeze().cpu(), dtype=np.uint8)
+        frame_tar_seg = np.array(Image.fromarray(frame_tar_seg).resize((ori_w, ori_h), 0))
+        frame_nm = frame_list[cnt].split('/')[-1].replace(".jpg", ".png")
+        imwrite_indexed(os.path.join(video_folder, frame_nm), frame_tar_seg, color_palette)
+
+
+def restrict_neighborhood(h, w):
+    # We restrict the set of source nodes considered to a spatial neighborhood of the query node (i.e. ``local attention'')
+    mask = torch.zeros(h, w, h, w)
+    for i in range(h):
+        for j in range(w):
+            for p in range(2 * args.size_mask_neighborhood + 1):
+                for q in range(2 * args.size_mask_neighborhood + 1):
+                    if i - args.size_mask_neighborhood + p < 0 or i - args.size_mask_neighborhood + p >= h:
+                        continue
+                    if j - args.size_mask_neighborhood + q < 0 or j - args.size_mask_neighborhood + q >= w:
+                        continue
+                    mask[i, j, i - args.size_mask_neighborhood + p, j - args.size_mask_neighborhood + q] = 1
+
+    mask = mask.reshape(h * w, h * w)
+    return mask.cuda(non_blocking=True)
+
+
+def norm_mask(mask):
+    c, h, w = mask.size()
+    for cnt in range(c):
+        mask_cnt = mask[cnt,:,:]
+        if(mask_cnt.max() > 0):
+            mask_cnt = (mask_cnt - mask_cnt.min())
+            mask_cnt = mask_cnt/mask_cnt.max()
+            mask[cnt,:,:] = mask_cnt
+    return mask
+
+
+def label_propagation(args, model, frame_tar, list_frame_feats, list_segs, mask_neighborhood=None):
+    """
+    propagate segs of frames in list_frames to frame_tar
+    """
+    ## we only need to extract feature of the target frame
+    feat_tar, h, w = extract_feature(model, frame_tar, return_h_w=True)
+
+    return_feat_tar = feat_tar.T # dim x h*w
+
+    ncontext = len(list_frame_feats)
+    feat_sources = torch.stack(list_frame_feats) # nmb_context x dim x h*w
+
+    feat_tar = F.normalize(feat_tar, dim=1, p=2)
+    feat_sources = F.normalize(feat_sources, dim=1, p=2)
+
+    feat_tar = feat_tar.unsqueeze(0).repeat(ncontext, 1, 1)
+    aff = torch.exp(torch.bmm(feat_tar, feat_sources) / 0.1) # nmb_context x h*w (tar: query) x h*w (source: keys)
+
+    if args.size_mask_neighborhood > 0:
+        if mask_neighborhood is None:
+            mask_neighborhood = restrict_neighborhood(h, w)
+            mask_neighborhood = mask_neighborhood.unsqueeze(0).repeat(ncontext, 1, 1)
+        aff *= mask_neighborhood
+
+    aff = aff.transpose(2, 1).reshape(-1, h * w) # nmb_context*h*w (source: keys) x h*w (tar: queries)
+    tk_val, _ = torch.topk(aff, dim=0, k=args.topk)
+    tk_val_min, _ = torch.min(tk_val, dim=0)
+    aff[aff < tk_val_min] = 0
+
+    aff = aff / torch.sum(aff, keepdim=True, axis=0)
+
+    list_segs = [s.cuda() for s in list_segs]
+    segs = torch.cat(list_segs)
+    nmb_context, C, h, w = segs.shape
+    segs = segs.reshape(nmb_context, C, -1).transpose(2, 1).reshape(-1, C).T # C x nmb_context*h*w
+    seg_tar = torch.mm(segs, aff)
+    seg_tar = seg_tar.reshape(1, C, h, w)
+    return seg_tar, return_feat_tar, mask_neighborhood
+ 
+
+def extract_feature(model, frame, return_h_w=False):
+    """Extract one frame feature everytime."""
+    out = model.get_intermediate_layers(frame.unsqueeze(0).cuda(), n=1)[0]
+    out = out[:, 1:, :]  # we discard the [CLS] token
+    h, w = int(frame.shape[1] / model.patch_embed.patch_size), int(frame.shape[2] / model.patch_embed.patch_size)
+    dim = out.shape[-1]
+    out = out[0].reshape(h, w, dim)
+    out = out.reshape(-1, dim)
+    if return_h_w:
+        return out, h, w
+    return out
+
+
+def imwrite_indexed(filename, array, color_palette):
+    """ Save indexed png for DAVIS."""
+    if np.atleast_3d(array).shape[2] != 1:
+      raise Exception("Saving indexed PNGs requires 2D array.")
+
+    im = Image.fromarray(array)
+    im.putpalette(color_palette.ravel())
+    im.save(filename, format='PNG')
+
+
+def to_one_hot(y_tensor, n_dims=None):
+    """
+    Take integer y (tensor or variable) with n dims &
+    convert it to 1-hot representation with n+1 dims.
+    """
+    if(n_dims is None):
+        n_dims = int(y_tensor.max()+ 1)
+    _,h,w = y_tensor.size()
+    y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)
+    n_dims = n_dims if n_dims is not None else int(torch.max(y_tensor)) + 1
+    y_one_hot = torch.zeros(y_tensor.size()[0], n_dims).scatter_(1, y_tensor, 1)
+    y_one_hot = y_one_hot.view(h,w,n_dims)
+    return y_one_hot.permute(2, 0, 1).unsqueeze(0)
+
+
+def read_frame_list(video_dir):
+    frame_list = [img for img in glob.glob(os.path.join(video_dir,"*.jpg"))]
+    frame_list = sorted(frame_list)
+    return frame_list
+
+
+def read_frame(frame_dir, scale_size=[480]):
+    """
+    read a single frame & preprocess
+    """
+    img = cv2.imread(frame_dir)
+    ori_h, ori_w, _ = img.shape
+    if len(scale_size) == 1:
+        if(ori_h > ori_w):
+            tw = scale_size[0]
+            th = (tw * ori_h) / ori_w
+            th = int((th // 64) * 64)
+        else:
+            th = scale_size[0]
+            tw = (th * ori_w) / ori_h
+            tw = int((tw // 64) * 64)
+    else:
+        th, tw = scale_size
+    img = cv2.resize(img, (tw, th))
+    img = img.astype(np.float32)
+    img = img / 255.0
+    img = img[:, :, ::-1]
+    img = np.transpose(img.copy(), (2, 0, 1))
+    img = torch.from_numpy(img).float()
+    img = color_normalize(img)
+    return img, ori_h, ori_w
+
+
+def read_seg(seg_dir, factor, scale_size=[480]):
+    seg = Image.open(seg_dir)
+    _w, _h = seg.size # note PIL.Image.Image's size is (w, h)
+    if len(scale_size) == 1:
+        if(_w > _h):
+            _th = scale_size[0]
+            _tw = (_th * _w) / _h
+            _tw = int((_tw // 64) * 64)
+        else:
+            _tw = scale_size[0]
+            _th = (_tw * _h) / _w
+            _th = int((_th // 64) * 64)
+    else:
+        _th = scale_size[1]
+        _tw = scale_size[0]
+    small_seg = np.array(seg.resize((_tw // factor, _th // factor), 0))
+    small_seg = torch.from_numpy(small_seg.copy()).contiguous().float().unsqueeze(0)
+    return to_one_hot(small_seg), np.asarray(seg)
+
+
+def color_normalize(x, mean=[0.485, 0.456, 0.406], std=[0.228, 0.224, 0.225]):
+    for t, m, s in zip(x, mean, std):
+        t.sub_(m)
+        t.div_(s)
+    return x
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Evaluation with video object segmentation on DAVIS 2017')
+    parser.add_argument('--pretrained_weights', default='', type=str, help="Path to pretrained weights to evaluate.")
+    parser.add_argument('--arch', default='vit_small', type=str,
+        choices=['vit_tiny', 'vit_small', 'vit_base'], help='Architecture (support only ViT atm).')
+    parser.add_argument('--patch_size', default=16, type=int, help='Patch resolution of the model.')
+    parser.add_argument("--checkpoint_key", default="teacher", type=str, help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument('--output_dir', default=".", help='Path where to save segmentations')
+    parser.add_argument('--data_path', default='/path/to/davis/', type=str)
+    parser.add_argument("--n_last_frames", type=int, default=7, help="number of preceeding frames")
+    parser.add_argument("--size_mask_neighborhood", default=12, type=int,
+        help="We restrict the set of source nodes considered to a spatial neighborhood of the query node")
+    parser.add_argument("--topk", type=int, default=5, help="accumulate label from top k neighbors")
+    parser.add_argument("--bs", type=int, default=6, help="Batch size, try to reduce if OOM")
+    args = parser.parse_args()
+
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+
+    # building network
+    model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+    print(f"Model {args.arch} {args.patch_size}x{args.patch_size} built.")
+    model.cuda()
+    utils.load_pretrained_weights(model, args.pretrained_weights, args.checkpoint_key, args.arch, args.patch_size)
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+    color_palette = []
+    for line in urlopen("https://raw.githubusercontent.com/Liusifei/UVC/master/libs/data/palette.txt"):
+        color_palette.append([int(i) for i in line.decode("utf-8").split('\n')[0].split(" ")])
+    color_palette = np.asarray(color_palette, dtype=np.uint8).reshape(-1,3)
+
+    video_list = open(os.path.join(args.data_path, "ImageSets/2017/val.txt")).readlines()
+    for i, video_name in enumerate(video_list):
+        video_name = video_name.strip()
+        print(f'[{i}/{len(video_list)}] Begin to segmentate video {video_name}.')
+        video_dir = os.path.join(args.data_path, "JPEGImages/480p/", video_name)
+        frame_list = read_frame_list(video_dir)
+        seg_path = frame_list[0].replace("JPEGImages", "Annotations").replace("jpg", "png")
+        first_seg, seg_ori = read_seg(seg_path, args.patch_size)
+        eval_video_tracking_davis(args, model, frame_list, video_dir, first_seg, seg_ori, color_palette)

PuzzleTuning/Counterpart PreTrain Methods/dino-main/hubconf.py

@@ -0,0 +1,151 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from torchvision.models.resnet import resnet50
+
+import vision_transformer as vits
+
+dependencies = ["torch", "torchvision"]
+
+
+def dino_vits16(pretrained=True, **kwargs):
+    """
+    ViT-Small/16x16 pre-trained with DINO.
+    Achieves 74.5% top-1 accuracy on ImageNet with k-NN classification.
+    """
+    model = vits.__dict__["vit_small"](patch_size=16, num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_vits8(pretrained=True, **kwargs):
+    """
+    ViT-Small/8x8 pre-trained with DINO.
+    Achieves 78.3% top-1 accuracy on ImageNet with k-NN classification.
+    """
+    model = vits.__dict__["vit_small"](patch_size=8, num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_vitb16(pretrained=True, **kwargs):
+    """
+    ViT-Base/16x16 pre-trained with DINO.
+    Achieves 76.1% top-1 accuracy on ImageNet with k-NN classification.
+    """
+    model = vits.__dict__["vit_base"](patch_size=16, num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_vitb8(pretrained=True, **kwargs):
+    """
+    ViT-Base/8x8 pre-trained with DINO.
+    Achieves 77.4% top-1 accuracy on ImageNet with k-NN classification.
+    """
+    model = vits.__dict__["vit_base"](patch_size=8, num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_resnet50(pretrained=True, **kwargs):
+    """
+    ResNet-50 pre-trained with DINO.
+    Achieves 75.3% top-1 accuracy on ImageNet linear evaluation benchmark (requires to train `fc`).
+    """
+    model = resnet50(pretrained=False, **kwargs)
+    model.fc = torch.nn.Identity()
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_resnet50_pretrain/dino_resnet50_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=False)
+    return model
+
+
+def dino_xcit_small_12_p16(pretrained=True, **kwargs):
+    """
+    XCiT-Small-12/16 pre-trained with DINO.
+    """
+    model = torch.hub.load('facebookresearch/xcit:main', "xcit_small_12_p16", num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_xcit_small_12_p16_pretrain/dino_xcit_small_12_p16_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_xcit_small_12_p8(pretrained=True, **kwargs):
+    """
+    XCiT-Small-12/8 pre-trained with DINO.
+    """
+    model = torch.hub.load('facebookresearch/xcit:main', "xcit_small_12_p8", num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_xcit_small_12_p8_pretrain/dino_xcit_small_12_p8_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_xcit_medium_24_p16(pretrained=True, **kwargs):
+    """
+    XCiT-Medium-24/16 pre-trained with DINO.
+    """
+    model = torch.hub.load('facebookresearch/xcit:main', "xcit_medium_24_p16", num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_xcit_medium_24_p16_pretrain/dino_xcit_medium_24_p16_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model
+
+
+def dino_xcit_medium_24_p8(pretrained=True, **kwargs):
+    """
+    XCiT-Medium-24/8 pre-trained with DINO.
+    """
+    model = torch.hub.load('facebookresearch/xcit:main', "xcit_medium_24_p8", num_classes=0, **kwargs)
+    if pretrained:
+        state_dict = torch.hub.load_state_dict_from_url(
+            url="https://dl.fbaipublicfiles.com/dino/dino_xcit_medium_24_p8_pretrain/dino_xcit_medium_24_p8_pretrain.pth",
+            map_location="cpu",
+        )
+        model.load_state_dict(state_dict, strict=True)
+    return model

PuzzleTuning/Counterpart PreTrain Methods/dino-main/main_dino.py

@@ -0,0 +1,515 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import os
+import sys
+import datetime
+import time
+import math
+import json
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+import torch.backends.cudnn as cudnn
+import torch.nn.functional as F
+from torchvision import datasets, transforms
+from torchvision import models as torchvision_models
+
+import utils
+import vision_transformer as vits
+from vision_transformer import DINOHead
+
+torchvision_archs = sorted(name for name in torchvision_models.__dict__
+    if name.islower() and not name.startswith("__")
+    and callable(torchvision_models.__dict__[name]))
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('DINO', add_help=False)
+
+    # Model parameters
+    parser.add_argument('--arch', default='vit_base', type=str,
+        help="""Name of architecture to train. For quick experiments with ViTs,
+        we recommend using vit_tiny or vit_small.""")
+    parser.add_argument('--patch_size', default=16, type=int, help="""Size in pixels
+        of input square patches - default 16 (for 16x16 patches). Using smaller
+        values leads to better performance but requires more memory. Applies only
+        for ViTs (vit_tiny, vit_small and vit_base). If <16, we recommend disabling
+        mixed precision training (--use_fp16 false) to avoid unstabilities.""")
+    parser.add_argument('--input_size', default=224, type=int)
+    parser.add_argument('--out_dim', default=65536, type=int, help="""Dimensionality of
+        the DINO head output. For complex and large datasets large values (like 65k) work well.""")
+    parser.add_argument('--norm_last_layer', default=True, type=utils.bool_flag,
+        help="""Whether or not to weight normalize the last layer of the DINO head.
+        Not normalizing leads to better performance but can make the training unstable.
+        In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""")
+    parser.add_argument('--momentum_teacher', default=0.996, type=float, help="""Base EMA
+        parameter for teacher update. The value is increased to 1 during training with cosine schedule.
+        We recommend setting a higher value with small batches: for example use 0.9995 with batch size of 256.""")
+    parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag,
+        help="Whether to use batch normalizations in projection head (Default: False)")
+
+    # Temperature teacher parameters
+    parser.add_argument('--warmup_teacher_temp', default=0.04, type=float,
+        help="""Initial value for the teacher temperature: 0.04 works well in most cases.
+        Try decreasing it if the training loss does not decrease.""")
+    parser.add_argument('--teacher_temp', default=0.04, type=float, help="""Final value (after linear warmup)
+        of the teacher temperature. For most experiments, anything above 0.07 is unstable. We recommend
+        starting with the default value of 0.04 and increase this slightly if needed.""")
+    parser.add_argument('--warmup_teacher_temp_epochs', default=0, type=int,
+        help='Number of warmup epochs for the teacher temperature (Default: 30).')
+
+    # Training/Optimization parameters
+    parser.add_argument('--use_fp16', type=utils.bool_flag, default=True, help="""Whether or not
+        to use half precision for training. Improves training time and memory requirements,
+        but can provoke instability and slight decay of performance. We recommend disabling
+        mixed precision if the loss is unstable, if reducing the patch size or if training with bigger ViTs.""")
+    parser.add_argument('--weight_decay', type=float, default=0.04, help="""Initial value of the
+        weight decay. With ViT, a smaller value at the beginning of training works well.""")
+    parser.add_argument('--weight_decay_end', type=float, default=0.4, help="""Final value of the
+        weight decay. We use a cosine schedule for WD and using a larger decay by
+        the end of training improves performance for ViTs.""")
+    parser.add_argument('--clip_grad', type=float, default=3.0, help="""Maximal parameter
+        gradient norm if using gradient clipping. Clipping with norm .3 ~ 1.0 can
+        help optimization for larger ViT architectures. 0 for disabling.""")
+    parser.add_argument('--batch_size_per_gpu', default=512, type=int,
+        help='Per-GPU batch-size : number of distinct images loaded on one GPU.')
+    parser.add_argument('--epochs', default=150, type=int, help='Number of epochs of training.')
+    parser.add_argument('--freeze_last_layer', default=1, type=int, help="""Number of epochs
+        during which we keep the output layer fixed. Typically doing so during
+        the first epoch helps training. Try increasing this value if the loss does not decrease.""")
+    parser.add_argument("--lr", default=1.5e-4, type=float, help="""Learning rate at the end of
+        linear warmup (highest LR used during training). The learning rate is linearly scaled
+        with the batch size, and specified here for a reference batch size of 512.""")
+    parser.add_argument("--warmup_epochs", default=20, type=int,
+        help="Number of epochs for the linear learning-rate warm up.")
+    parser.add_argument('--min_lr', type=float, default=1e-6, help="""Target LR at the
+        end of optimization. We use a cosine LR schedule with linear warmup.""")
+    parser.add_argument('--optimizer', default='adamw', type=str,
+        choices=['adamw', 'sgd', 'lars'], help="""Type of optimizer. We recommend using adamw with ViTs.""")
+    parser.add_argument('--drop_path_rate', type=float, default=0.1, help="stochastic depth rate")
+
+    # Multi-crop parameters
+    parser.add_argument('--global_crops_scale', type=float, nargs='+', default=(0.4, 1.),
+        help="""Scale range of the cropped image before resizing, relatively to the origin image.
+        Used for large global view cropping. When disabling multi-crop (--local_crops_number 0), we
+        recommand using a wider range of scale ("--global_crops_scale 0.14 1." for example)""")
+    parser.add_argument('--local_crops_number', type=int, default=8, help="""Number of small
+        local views to generate. Set this parameter to 0 to disable multi-crop training.
+        When disabling multi-crop we recommend to use "--global_crops_scale 0.14 1." """)
+    parser.add_argument('--local_crops_scale', type=float, nargs='+', default=(0.05, 0.4),
+        help="""Scale range of the cropped image before resizing, relatively to the origin image.
+        Used for small local view cropping of multi-crop.""")
+
+    # Misc
+    parser.add_argument('--data_path', default='/root/autodl-tmp/All', type=str,
+        help='Please specify path to the ImageNet training data.')
+    parser.add_argument('--basic_state_dict', default='/root/autodl-tmp/ViT_b16_224_Imagenet.pth', type=str,
+        help='Load in pretrained or un-pretrained model pth')
+    parser.add_argument('--output_dir', default="/home/CPIA/saved_models/DINO", type=str, help='Path to save logs and checkpoints.')
+    parser.add_argument('--saveckp_freq', default=50, type=int, help='Save checkpoint every x epochs.')
+    parser.add_argument('--seed', default=0, type=int, help='Random seed.')
+    parser.add_argument('--num_workers', default=32, type=int, help='Number of data loading workers per GPU.')
+    parser.add_argument("--dist_url", default="env://", type=str, help="""url used to set up
+        distributed training; see https://pytorch.org/docs/stable/distributed.html""")
+    parser.add_argument("--local_rank", default=0, type=int, help="Please ignore and do not set this argument.")
+    return parser
+
+
+def train_dino(args):
+    utils.init_distributed_mode(args)
+    utils.fix_random_seeds(args.seed)
+    print('ok')
+    print("git:\n  {}\n".format(utils.get_sha()))
+    print("\n".join("%s: %s" % (k, str(v)) for k, v in sorted(dict(vars(args)).items())))
+    cudnn.benchmark = True
+
+    # ============ preparing data ... ============
+    transform = DataAugmentationDINO(
+        args.global_crops_scale,
+        args.local_crops_scale,
+        args.local_crops_number,
+    )
+    dataset = datasets.ImageFolder(args.data_path, transform=transform)
+    sampler = torch.utils.data.DistributedSampler(dataset, shuffle=True)
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        sampler=sampler,
+        batch_size=args.batch_size_per_gpu,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+    )
+    print(f"Data loaded: there are {len(dataset)} images.")
+
+    # ============ building student and teacher networks ... ============
+    # we changed the name DeiT-S for ViT-S to avoid confusions
+    args.arch = args.arch.replace("deit", "vit")
+    # if the network is a Vision Transformer (i.e. vit_tiny, vit_small, vit_base)
+    if args.arch in vits.__dict__.keys():
+        student = vits.__dict__[args.arch](
+            patch_size=args.patch_size,
+            drop_path_rate=args.drop_path_rate,  # stochastic depth
+        )
+        if args.basic_state_dict is not None:  # Transfer-learning
+            try:
+                basic_state_dict = torch.load(args.basic_state_dict)
+                if 'model' in basic_state_dict:
+                    basic_state_dict = basic_state_dict['model']
+                else:
+                    pass
+                student.load_state_dict(basic_state_dict, False)
+
+            except:
+                print('erro in args.basic_state_dict:', args.basic_state_dict)
+                print('Student PreTuning Restart')  # 没倒进去
+
+            else:
+                print('Student PreTuning with Transfer-learning with:', args.basic_state_dict)
+
+        else:
+            print('Student PreTuning Restart')
+        teacher = vits.__dict__[args.arch](patch_size=args.patch_size)
+        if args.basic_state_dict is not None:  # Transfer-learning
+            try:
+                basic_state_dict = torch.load(args.basic_state_dict)
+                if 'model' in basic_state_dict:
+                    basic_state_dict = basic_state_dict['model']
+                else:
+                    pass
+                teacher.load_state_dict(basic_state_dict, False)
+
+            except:
+                print('erro in args.basic_state_dict:', args.basic_state_dict)
+                print('Teacher PreTuning Restart')  # 没倒进去
+
+            else:
+                print('Teacher PreTuning with Transfer-learning with:', args.basic_state_dict)
+
+        else:
+            print('Teacher PreTuning Restart')
+        embed_dim = student.embed_dim
+    # if the network is a XCiT
+    """elif args.arch in torch.hub.list("facebookresearch/xcit:main"):
+        student = torch.hub.load('facebookresearch/xcit:main', args.arch,
+                                 pretrained=False, drop_path_rate=args.drop_path_rate)
+        teacher = torch.hub.load('facebookresearch/xcit:main', args.arch, pretrained=False)
+        embed_dim = student.embed_dim
+    # otherwise, we check if the architecture is in torchvision models
+    elif args.arch in torchvision_models.__dict__.keys():
+        student = torchvision_models.__dict__[args.arch]()
+        teacher = torchvision_models.__dict__[args.arch]()
+        embed_dim = student.fc.weight.shape[1]
+    else:
+        print(f"Unknow architecture: {args.arch}")"""
+
+    # multi-crop wrapper handles forward with inputs of different resolutions
+    student = utils.MultiCropWrapper(student, DINOHead(
+        embed_dim,
+        args.out_dim,
+        use_bn=args.use_bn_in_head,
+        norm_last_layer=args.norm_last_layer,
+    ))
+    teacher = utils.MultiCropWrapper(
+        teacher,
+        DINOHead(embed_dim, args.out_dim, args.use_bn_in_head),
+    )
+    # move networks to gpu
+    student, teacher = student.cuda(), teacher.cuda()
+    # synchronize batch norms (if any)
+    if utils.has_batchnorms(student):
+        student = nn.SyncBatchNorm.convert_sync_batchnorm(student)
+        teacher = nn.SyncBatchNorm.convert_sync_batchnorm(teacher)
+
+        # we need DDP wrapper to have synchro batch norms working...
+        teacher = nn.parallel.DistributedDataParallel(teacher, device_ids=[args.gpu])
+        teacher_without_ddp = teacher.module
+    else:
+        # teacher_without_ddp and teacher are the same thing
+        teacher_without_ddp = teacher
+    student = nn.parallel.DistributedDataParallel(student, device_ids=[args.gpu])
+    # teacher and student start with the same weights
+    teacher_without_ddp.load_state_dict(student.module.state_dict())
+    # there is no backpropagation through the teacher, so no need for gradients
+    for p in teacher.parameters():
+        p.requires_grad = False
+    print(f"Student and Teacher are built: they are both {args.arch} network.")
+
+    # ============ preparing loss ... ============
+    dino_loss = DINOLoss(
+        args.out_dim,
+        args.local_crops_number + 2,  # total number of crops = 2 global crops + local_crops_number
+        args.warmup_teacher_temp,
+        args.teacher_temp,
+        args.warmup_teacher_temp_epochs,
+        args.epochs,
+    ).cuda()
+
+    # ============ preparing optimizer ... ============
+    params_groups = utils.get_params_groups(student)
+    if args.optimizer == "adamw":
+        optimizer = torch.optim.AdamW(params_groups)  # to use with ViTs
+    elif args.optimizer == "sgd":
+        optimizer = torch.optim.SGD(params_groups, lr=0, momentum=0.9)  # lr is set by scheduler
+    elif args.optimizer == "lars":
+        optimizer = utils.LARS(params_groups)  # to use with convnet and large batches
+    # for mixed precision training
+    fp16_scaler = None
+    if args.use_fp16:
+        fp16_scaler = torch.cuda.amp.GradScaler()
+
+    # ============ init schedulers ... ============
+    lr_schedule = utils.cosine_scheduler(
+        args.lr * (args.batch_size_per_gpu * utils.get_world_size()) / 256.,  # linear scaling rule
+        args.min_lr,
+        args.epochs, len(data_loader),
+        warmup_epochs=args.warmup_epochs,
+    )
+    wd_schedule = utils.cosine_scheduler(
+        args.weight_decay,
+        args.weight_decay_end,
+        args.epochs, len(data_loader),
+    )
+    # momentum parameter is increased to 1. during training with a cosine schedule
+    momentum_schedule = utils.cosine_scheduler(args.momentum_teacher, 1,
+                                               args.epochs, len(data_loader))
+    print(f"Loss, optimizer and schedulers ready.")
+
+    # ============ optionally resume training ... ============
+    to_restore = {"epoch": 0}
+    utils.restart_from_checkpoint(
+        os.path.join(args.output_dir, "checkpoint.pth"),
+        run_variables=to_restore,
+        student=student,
+        teacher=teacher,
+        optimizer=optimizer,
+        fp16_scaler=fp16_scaler,
+        dino_loss=dino_loss,
+    )
+    start_epoch = to_restore["epoch"]
+
+    start_time = time.time()
+    print("Starting DINO training !")
+    for epoch in range(start_epoch, args.epochs):
+        data_loader.sampler.set_epoch(epoch)
+
+        # ============ training one epoch of DINO ... ============
+        train_stats = train_one_epoch(student, teacher, teacher_without_ddp, dino_loss,
+            data_loader, optimizer, lr_schedule, wd_schedule, momentum_schedule,
+            epoch, fp16_scaler, args)
+
+        # ============ writing logs ... ============
+        save_dict = {
+            'student': student.state_dict(),
+            'teacher': teacher.state_dict(),
+            'optimizer': optimizer.state_dict(),
+            'epoch': epoch + 1,
+            'args': args,
+            'dino_loss': dino_loss.state_dict(),
+        }
+        if fp16_scaler is not None:
+            save_dict['fp16_scaler'] = fp16_scaler.state_dict()
+        utils.save_on_master(save_dict, os.path.join(args.output_dir, 'checkpoint.pth'))
+        if args.saveckp_freq and (epoch % args.saveckp_freq == 0 or epoch + 1 == args.epochs):
+            initial_setting = os.path.split(args.basic_state_dict)[1].split('.')[0]  # 'ViT_b16_224_Imagenet'
+            dataset_using = os.path.split(args.data_path)[1]
+            utils.save_on_master(save_dict, os.path.join(args.output_dir, f'dino_'+initial_setting + '_' + dataset_using+f'_checkpoint{epoch:04}.pth'))
+
+
+
+
+        log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
+                     'epoch': epoch}
+        if utils.is_main_process():
+            with (Path(args.output_dir) / "log.txt").open("a") as f:
+                f.write(json.dumps(log_stats) + "\n")
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+
+def train_one_epoch(student, teacher, teacher_without_ddp, dino_loss, data_loader,
+                    optimizer, lr_schedule, wd_schedule, momentum_schedule,epoch,
+                    fp16_scaler, args):
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Epoch: [{}/{}]'.format(epoch, args.epochs)
+    for it, (images, _) in enumerate(metric_logger.log_every(data_loader, 10, header)):
+        # update weight decay and learning rate according to their schedule
+        it = len(data_loader) * epoch + it  # global training iteration
+        for i, param_group in enumerate(optimizer.param_groups):
+            param_group["lr"] = lr_schedule[it]
+            if i == 0:  # only the first group is regularized
+                param_group["weight_decay"] = wd_schedule[it]
+
+        # move images to gpu
+        images = [im.cuda(non_blocking=True) for im in images]
+        # teacher and student forward passes + compute dino loss
+        with torch.cuda.amp.autocast(fp16_scaler is not None):
+            teacher_output = teacher(images[:2])  # only the 2 global views pass through the teacher
+            student_output = student(images)
+            loss = dino_loss(student_output, teacher_output, epoch)
+
+        if not math.isfinite(loss.item()):
+            print("Loss is {}, stopping training".format(loss.item()), force=True)
+            sys.exit(1)
+
+        # student update
+        optimizer.zero_grad()
+        param_norms = None
+        if fp16_scaler is None:
+            loss.backward()
+            if args.clip_grad:
+                param_norms = utils.clip_gradients(student, args.clip_grad)
+            utils.cancel_gradients_last_layer(epoch, student,
+                                              args.freeze_last_layer)
+            optimizer.step()
+        else:
+            fp16_scaler.scale(loss).backward()
+            if args.clip_grad:
+                fp16_scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                param_norms = utils.clip_gradients(student, args.clip_grad)
+            utils.cancel_gradients_last_layer(epoch, student,
+                                              args.freeze_last_layer)
+            fp16_scaler.step(optimizer)
+            fp16_scaler.update()
+
+        # EMA update for the teacher
+        with torch.no_grad():
+            m = momentum_schedule[it]  # momentum parameter
+            for param_q, param_k in zip(student.module.parameters(), teacher_without_ddp.parameters()):
+                param_k.data.mul_(m).add_((1 - m) * param_q.detach().data)
+
+        # logging
+        torch.cuda.synchronize()
+        metric_logger.update(loss=loss.item())
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        metric_logger.update(wd=optimizer.param_groups[0]["weight_decay"])
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+class DINOLoss(nn.Module):
+    def __init__(self, out_dim, ncrops, warmup_teacher_temp, teacher_temp,
+                 warmup_teacher_temp_epochs, nepochs, student_temp=0.1,
+                 center_momentum=0.9):
+        super().__init__()
+        self.student_temp = student_temp
+        self.center_momentum = center_momentum
+        self.ncrops = ncrops
+        self.register_buffer("center", torch.zeros(1, out_dim))
+        # we apply a warm up for the teacher temperature because
+        # a too high temperature makes the training instable at the beginning
+        self.teacher_temp_schedule = np.concatenate((
+            np.linspace(warmup_teacher_temp,
+                        teacher_temp, warmup_teacher_temp_epochs),
+            np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp
+        ))
+
+    def forward(self, student_output, teacher_output, epoch):
+        """
+        Cross-entropy between softmax outputs of the teacher and student networks.
+        """
+        student_out = student_output / self.student_temp
+        student_out = student_out.chunk(self.ncrops)
+
+        # teacher centering and sharpening
+        temp = self.teacher_temp_schedule[epoch]
+        teacher_out = F.softmax((teacher_output - self.center) / temp, dim=-1)
+        teacher_out = teacher_out.detach().chunk(2)
+
+        total_loss = 0
+        n_loss_terms = 0
+        for iq, q in enumerate(teacher_out):
+            for v in range(len(student_out)):
+                if v == iq:
+                    # we skip cases where student and teacher operate on the same view
+                    continue
+                loss = torch.sum(-q * F.log_softmax(student_out[v], dim=-1), dim=-1)
+                total_loss += loss.mean()
+                n_loss_terms += 1
+        total_loss /= n_loss_terms
+        self.update_center(teacher_output)
+        return total_loss
+
+    @torch.no_grad()
+    def update_center(self, teacher_output):
+        """
+        Update center used for teacher output.
+        """
+        batch_center = torch.sum(teacher_output, dim=0, keepdim=True)
+        dist.all_reduce(batch_center)
+        batch_center = batch_center / (len(teacher_output) * dist.get_world_size())
+
+        # ema update
+        self.center = self.center * self.center_momentum + batch_center * (1 - self.center_momentum)
+
+
+class DataAugmentationDINO(object):
+    def __init__(self, global_crops_scale, local_crops_scale, local_crops_number):
+        flip_and_color_jitter = transforms.Compose([
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.RandomApply(
+                [transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.2, hue=0.1)],
+                p=0.8
+            ),
+            transforms.RandomGrayscale(p=0.2),
+        ])
+        normalize = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ])
+
+        # first global crop
+        self.global_transfo1 = transforms.Compose([
+            transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Image.BICUBIC),
+            flip_and_color_jitter,
+            utils.GaussianBlur(1.0),
+            normalize,
+        ])
+        # second global crop
+        self.global_transfo2 = transforms.Compose([
+            transforms.RandomResizedCrop(224, scale=global_crops_scale, interpolation=Image.BICUBIC),
+            flip_and_color_jitter,
+            utils.GaussianBlur(0.1),
+            utils.Solarization(0.2),
+            normalize,
+        ])
+        # transformation for the local small crops
+        self.local_crops_number = local_crops_number
+        self.local_transfo = transforms.Compose([
+            transforms.RandomResizedCrop(96, scale=local_crops_scale, interpolation=Image.BICUBIC),
+            flip_and_color_jitter,
+            utils.GaussianBlur(p=0.5),
+            normalize,
+        ])
+
+    def __call__(self, image):
+        crops = []
+        crops.append(self.global_transfo1(image))
+        crops.append(self.global_transfo2(image))
+        for _ in range(self.local_crops_number):
+            crops.append(self.local_transfo(image))
+        return crops
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DINO', parents=[get_args_parser()])
+    args = parser.parse_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    train_dino(args)

PuzzleTuning/Counterpart PreTrain Methods/dino-main/run_with_submitit.py

@@ -0,0 +1,132 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+A script to run multinode training with submitit.
+Almost copy-paste from https://github.com/facebookresearch/deit/blob/main/run_with_submitit.py
+"""
+import argparse
+import os
+import uuid
+from pathlib import Path
+
+import main_dino
+import submitit
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("Submitit for DINO", parents=[main_dino.get_args_parser()])
+    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
+    parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request")
+    parser.add_argument("--timeout", default=2800, type=int, help="Duration of the job")
+
+    parser.add_argument("--partition", default="learnfair", type=str, help="Partition where to submit")
+    parser.add_argument("--use_volta32", action='store_true', help="Big models? Use this")
+    parser.add_argument('--comment', default="", type=str,
+                        help='Comment to pass to scheduler, e.g. priority message')
+    return parser.parse_args()
+
+
+def get_shared_folder() -> Path:
+    user = os.getenv("USER")
+    if Path("/checkpoint/").is_dir():
+        p = Path(f"/checkpoint/{user}/experiments")
+        p.mkdir(exist_ok=True)
+        return p
+    raise RuntimeError("No shared folder available")
+
+
+def get_init_file():
+    # Init file must not exist, but it's parent dir must exist.
+    os.makedirs(str(get_shared_folder()), exist_ok=True)
+    init_file = get_shared_folder() / f"{uuid.uuid4().hex}_init"
+    if init_file.exists():
+        os.remove(str(init_file))
+    return init_file
+
+
+class Trainer(object):
+    def __init__(self, args):
+        self.args = args
+
+    def __call__(self):
+        import main_dino
+
+        self._setup_gpu_args()
+        main_dino.train_dino(self.args)
+
+    def checkpoint(self):
+        import os
+        import submitit
+
+        self.args.dist_url = get_init_file().as_uri()
+        print("Requeuing ", self.args)
+        empty_trainer = type(self)(self.args)
+        return submitit.helpers.DelayedSubmission(empty_trainer)
+
+    def _setup_gpu_args(self):
+        import submitit
+        from pathlib import Path
+
+        job_env = submitit.JobEnvironment()
+        self.args.output_dir = Path(str(self.args.output_dir).replace("%j", str(job_env.job_id)))
+        self.args.gpu = job_env.local_rank
+        self.args.rank = job_env.global_rank
+        self.args.world_size = job_env.num_tasks
+        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+
+
+def main():
+    args = parse_args()
+    if args.output_dir == "":
+        args.output_dir = get_shared_folder() / "%j"
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    executor = submitit.AutoExecutor(folder=args.output_dir, slurm_max_num_timeout=30)
+
+    num_gpus_per_node = args.ngpus
+    nodes = args.nodes
+    timeout_min = args.timeout
+
+    partition = args.partition
+    kwargs = {}
+    if args.use_volta32:
+        kwargs['slurm_constraint'] = 'volta32gb'
+    if args.comment:
+        kwargs['slurm_comment'] = args.comment
+
+    executor.update_parameters(
+        mem_gb=40 * num_gpus_per_node,
+        gpus_per_node=num_gpus_per_node,
+        tasks_per_node=num_gpus_per_node,  # one task per GPU
+        cpus_per_task=10,
+        nodes=nodes,
+        timeout_min=timeout_min,  # max is 60 * 72
+        # Below are cluster dependent parameters
+        slurm_partition=partition,
+        slurm_signal_delay_s=120,
+        **kwargs
+    )
+
+    executor.update_parameters(name="dino")
+
+    args.dist_url = get_init_file().as_uri()
+
+    trainer = Trainer(args)
+    job = executor.submit(trainer)
+
+    print(f"Submitted job_id: {job.job_id}")
+    print(f"Logs and checkpoints will be saved at: {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

PuzzleTuning/Counterpart PreTrain Methods/dino-main/utils.py

@@ -0,0 +1,829 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Misc functions.
+
+Mostly copy-paste from torchvision references or other public repos like DETR:
+https://github.com/facebookresearch/detr/blob/master/util/misc.py
+"""
+import os
+import sys
+import time
+import math
+import random
+import datetime
+import subprocess
+from collections import defaultdict, deque
+
+import numpy as np
+import torch
+from torch import nn
+import torch.distributed as dist
+from PIL import ImageFilter, ImageOps
+
+
+class GaussianBlur(object):
+    """
+    Apply Gaussian Blur to the PIL image.
+    """
+    def __init__(self, p=0.5, radius_min=0.1, radius_max=2.):
+        self.prob = p
+        self.radius_min = radius_min
+        self.radius_max = radius_max
+
+    def __call__(self, img):
+        do_it = random.random() <= self.prob
+        if not do_it:
+            return img
+
+        return img.filter(
+            ImageFilter.GaussianBlur(
+                radius=random.uniform(self.radius_min, self.radius_max)
+            )
+        )
+
+
+class Solarization(object):
+    """
+    Apply Solarization to the PIL image.
+    """
+    def __init__(self, p):
+        self.p = p
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return ImageOps.solarize(img)
+        else:
+            return img
+
+
+def load_pretrained_weights(model, pretrained_weights, checkpoint_key, model_name, patch_size):
+    if os.path.isfile(pretrained_weights):
+        state_dict = torch.load(pretrained_weights, map_location="cpu")
+        if checkpoint_key is not None and checkpoint_key in state_dict:
+            print(f"Take key {checkpoint_key} in provided checkpoint dict")
+            state_dict = state_dict[checkpoint_key]
+        # remove `module.` prefix
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        # remove `backbone.` prefix induced by multicrop wrapper
+        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+        msg = model.load_state_dict(state_dict, strict=False)
+        print('Pretrained weights found at {} and loaded with msg: {}'.format(pretrained_weights, msg))
+    else:
+        print("Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate.")
+        url = None
+        if model_name == "vit_small" and patch_size == 16:
+            url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
+        elif model_name == "vit_small" and patch_size == 8:
+            url = "dino_deitsmall8_pretrain/dino_deitsmall8_pretrain.pth"
+        elif model_name == "vit_base" and patch_size == 16:
+            url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
+        elif model_name == "vit_base" and patch_size == 8:
+            url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
+        elif model_name == "xcit_small_12_p16":
+            url = "dino_xcit_small_12_p16_pretrain/dino_xcit_small_12_p16_pretrain.pth"
+        elif model_name == "xcit_small_12_p8":
+            url = "dino_xcit_small_12_p8_pretrain/dino_xcit_small_12_p8_pretrain.pth"
+        elif model_name == "xcit_medium_24_p16":
+            url = "dino_xcit_medium_24_p16_pretrain/dino_xcit_medium_24_p16_pretrain.pth"
+        elif model_name == "xcit_medium_24_p8":
+            url = "dino_xcit_medium_24_p8_pretrain/dino_xcit_medium_24_p8_pretrain.pth"
+        elif model_name == "resnet50":
+            url = "dino_resnet50_pretrain/dino_resnet50_pretrain.pth"
+        if url is not None:
+            print("Since no pretrained weights have been provided, we load the reference pretrained DINO weights.")
+            state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url)
+            model.load_state_dict(state_dict, strict=True)
+        else:
+            print("There is no reference weights available for this model => We use random weights.")
+
+
+def load_pretrained_linear_weights(linear_classifier, model_name, patch_size):
+    url = None
+    if model_name == "vit_small" and patch_size == 16:
+        url = "dino_deitsmall16_pretrain/dino_deitsmall16_linearweights.pth"
+    elif model_name == "vit_small" and patch_size == 8:
+        url = "dino_deitsmall8_pretrain/dino_deitsmall8_linearweights.pth"
+    elif model_name == "vit_base" and patch_size == 16:
+        url = "dino_vitbase16_pretrain/dino_vitbase16_linearweights.pth"
+    elif model_name == "vit_base" and patch_size == 8:
+        url = "dino_vitbase8_pretrain/dino_vitbase8_linearweights.pth"
+    elif model_name == "resnet50":
+        url = "dino_resnet50_pretrain/dino_resnet50_linearweights.pth"
+    if url is not None:
+        print("We load the reference pretrained linear weights.")
+        state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url)["state_dict"]
+        linear_classifier.load_state_dict(state_dict, strict=True)
+    else:
+        print("We use random linear weights.")
+
+
+def clip_gradients(model, clip):
+    norms = []
+    for name, p in model.named_parameters():
+        if p.grad is not None:
+            param_norm = p.grad.data.norm(2)
+            norms.append(param_norm.item())
+            clip_coef = clip / (param_norm + 1e-6)
+            if clip_coef < 1:
+                p.grad.data.mul_(clip_coef)
+    return norms
+
+
+def cancel_gradients_last_layer(epoch, model, freeze_last_layer):
+    if epoch >= freeze_last_layer:
+        return
+    for n, p in model.named_parameters():
+        if "last_layer" in n:
+            p.grad = None
+
+
+def restart_from_checkpoint(ckp_path, run_variables=None, **kwargs):
+    """
+    Re-start from checkpoint
+    """
+    if not os.path.isfile(ckp_path):
+        return
+    print("Found checkpoint at {}".format(ckp_path))
+
+    # open checkpoint file
+    checkpoint = torch.load(ckp_path, map_location="cpu")
+
+    # key is what to look for in the checkpoint file
+    # value is the object to load
+    # example: {'state_dict': model}
+    for key, value in kwargs.items():
+        if key in checkpoint and value is not None:
+            try:
+                msg = value.load_state_dict(checkpoint[key], strict=False)
+                print("=> loaded '{}' from checkpoint '{}' with msg {}".format(key, ckp_path, msg))
+            except TypeError:
+                try:
+                    msg = value.load_state_dict(checkpoint[key])
+                    print("=> loaded '{}' from checkpoint: '{}'".format(key, ckp_path))
+                except ValueError:
+                    print("=> failed to load '{}' from checkpoint: '{}'".format(key, ckp_path))
+        else:
+            print("=> key '{}' not found in checkpoint: '{}'".format(key, ckp_path))
+
+    # re load variable important for the run
+    if run_variables is not None:
+        for var_name in run_variables:
+            if var_name in checkpoint:
+                run_variables[var_name] = checkpoint[var_name]
+
+
+def cosine_scheduler(base_value, final_value, epochs, niter_per_ep, warmup_epochs=0, start_warmup_value=0):
+    warmup_schedule = np.array([])
+    warmup_iters = warmup_epochs * niter_per_ep
+    if warmup_epochs > 0:
+        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)
+
+    iters = np.arange(epochs * niter_per_ep - warmup_iters)
+    schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))
+
+    schedule = np.concatenate((warmup_schedule, schedule))
+    assert len(schedule) == epochs * niter_per_ep
+    return schedule
+
+
+def bool_flag(s):
+    """
+    Parse boolean arguments from the command line.
+    """
+    FALSY_STRINGS = {"off", "false", "0"}
+    TRUTHY_STRINGS = {"on", "true", "1"}
+    if s.lower() in FALSY_STRINGS:
+        return False
+    elif s.lower() in TRUTHY_STRINGS:
+        return True
+    else:
+        raise argparse.ArgumentTypeError("invalid value for a boolean flag")
+
+
+def fix_random_seeds(seed=31):
+    """
+    Fix random seeds.
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.6f} ({global_avg:.6f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.6f}')
+        data_time = SmoothedValue(fmt='{avg:.6f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}',
+                'max mem: {memory:.0f}'
+            ])
+        else:
+            log_msg = self.delimiter.join([
+                header,
+                '[{0' + space_fmt + '}/{1}]',
+                'eta: {eta}',
+                '{meters}',
+                'time: {time}',
+                'data: {data}'
+            ])
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print('{} Total time: {} ({:.6f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def get_sha():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+
+    def _run(command):
+        return subprocess.check_output(command, cwd=cwd).decode('ascii').strip()
+    sha = 'N/A'
+    diff = "clean"
+    branch = 'N/A'
+    try:
+        sha = _run(['git', 'rev-parse', 'HEAD'])
+        subprocess.check_output(['git', 'diff'], cwd=cwd)
+        diff = _run(['git', 'diff-index', 'HEAD'])
+        diff = "has uncommited changes" if diff else "clean"
+        branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD'])
+    except Exception:
+        pass
+    message = f"sha: {sha}, status: {diff}, branch: {branch}"
+    return message
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def init_distributed_mode(args):
+    # launched with torch.distributed.launch
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    # launched with submitit on a slurm cluster
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    # launched naively with `python main_dino.py`
+    # we manually add MASTER_ADDR and MASTER_PORT to env variables
+    elif torch.cuda.is_available():
+        print('Will run the code on one GPU.')
+        args.rank, args.gpu, args.world_size = 0, 0, 1
+        os.environ['MASTER_ADDR'] = '127.0.0.1'
+        os.environ['MASTER_PORT'] = '29500'
+    else:
+        print('Does not support training without GPU.')
+        sys.exit(1)
+
+    dist.init_process_group(
+        backend="nccl",
+        init_method=args.dist_url,
+        world_size=args.world_size,
+        rank=args.rank,
+    )
+
+    torch.cuda.set_device(args.gpu)
+    print('| distributed init (rank {}): {}'.format(
+        args.rank, args.dist_url), flush=True)
+    dist.barrier()
+    setup_for_distributed(args.rank == 0)
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.reshape(1, -1).expand_as(pred))
+    return [correct[:k].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+class LARS(torch.optim.Optimizer):
+    """
+    Almost copy-paste from https://github.com/facebookresearch/barlowtwins/blob/main/main.py
+    """
+    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, eta=0.001,
+                 weight_decay_filter=None, lars_adaptation_filter=None):
+        defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum,
+                        eta=eta, weight_decay_filter=weight_decay_filter,
+                        lars_adaptation_filter=lars_adaptation_filter)
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g['params']:
+                dp = p.grad
+
+                if dp is None:
+                    continue
+
+                if p.ndim != 1:
+                    dp = dp.add(p, alpha=g['weight_decay'])
+
+                if p.ndim != 1:
+                    param_norm = torch.norm(p)
+                    update_norm = torch.norm(dp)
+                    one = torch.ones_like(param_norm)
+                    q = torch.where(param_norm > 0.,
+                                    torch.where(update_norm > 0,
+                                                (g['eta'] * param_norm / update_norm), one), one)
+                    dp = dp.mul(q)
+
+                param_state = self.state[p]
+                if 'mu' not in param_state:
+                    param_state['mu'] = torch.zeros_like(p)
+                mu = param_state['mu']
+                mu.mul_(g['momentum']).add_(dp)
+
+                p.add_(mu, alpha=-g['lr'])
+
+
+class MultiCropWrapper(nn.Module):
+    """
+    Perform forward pass separately on each resolution input.
+    The inputs corresponding to a single resolution are clubbed and single
+    forward is run on the same resolution inputs. Hence we do several
+    forward passes = number of different resolutions used. We then
+    concatenate all the output features and run the head forward on these
+    concatenated features.
+    """
+    def __init__(self, backbone, head):
+        super(MultiCropWrapper, self).__init__()
+        # disable layers dedicated to ImageNet labels classification
+        backbone.fc, backbone.head = nn.Identity(), nn.Identity()
+        self.backbone = backbone
+        self.head = head
+
+    def forward(self, x):
+        # convert to list
+        if not isinstance(x, list):
+            x = [x]
+        idx_crops = torch.cumsum(torch.unique_consecutive(
+            torch.tensor([inp.shape[-1] for inp in x]),
+            return_counts=True,
+        )[1], 0)
+        start_idx, output = 0, torch.empty(0).to(x[0].device)
+        for end_idx in idx_crops:
+            _out = self.backbone(torch.cat(x[start_idx: end_idx]))
+            # The output is a tuple with XCiT model. See:
+            # https://github.com/facebookresearch/xcit/blob/master/xcit.py#L404-L405
+            if isinstance(_out, tuple):
+                _out = _out[0]
+            # accumulate outputs
+            output = torch.cat((output, _out))
+            start_idx = end_idx
+        # Run the head forward on the concatenated features.
+        return self.head(output)
+
+
+def get_params_groups(model):
+    regularized = []
+    not_regularized = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        # we do not regularize biases nor Norm parameters
+        if name.endswith(".bias") or len(param.shape) == 1:
+            not_regularized.append(param)
+        else:
+            regularized.append(param)
+    return [{'params': regularized}, {'params': not_regularized, 'weight_decay': 0.}]
+
+
+def has_batchnorms(model):
+    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+    for name, module in model.named_modules():
+        if isinstance(module, bn_types):
+            return True
+    return False
+
+
+class PCA():
+    """
+    Class to  compute and apply PCA.
+    """
+    def __init__(self, dim=256, whit=0.5):
+        self.dim = dim
+        self.whit = whit
+        self.mean = None
+
+    def train_pca(self, cov):
+        """
+        Takes a covariance matrix (np.ndarray) as input.
+        """
+        d, v = np.linalg.eigh(cov)
+        eps = d.max() * 1e-5
+        n_0 = (d < eps).sum()
+        if n_0 > 0:
+            d[d < eps] = eps
+
+        # total energy
+        totenergy = d.sum()
+
+        # sort eigenvectors with eigenvalues order
+        idx = np.argsort(d)[::-1][:self.dim]
+        d = d[idx]
+        v = v[:, idx]
+
+        print("keeping %.2f %% of the energy" % (d.sum() / totenergy * 100.0))
+
+        # for the whitening
+        d = np.diag(1. / d**self.whit)
+
+        # principal components
+        self.dvt = np.dot(d, v.T)
+
+    def apply(self, x):
+        # input is from numpy
+        if isinstance(x, np.ndarray):
+            if self.mean is not None:
+                x -= self.mean
+            return np.dot(self.dvt, x.T).T
+
+        # input is from torch and is on GPU
+        if x.is_cuda:
+            if self.mean is not None:
+                x -= torch.cuda.FloatTensor(self.mean)
+            return torch.mm(torch.cuda.FloatTensor(self.dvt), x.transpose(0, 1)).transpose(0, 1)
+
+        # input if from torch, on CPU
+        if self.mean is not None:
+            x -= torch.FloatTensor(self.mean)
+        return torch.mm(torch.FloatTensor(self.dvt), x.transpose(0, 1)).transpose(0, 1)
+
+
+def compute_ap(ranks, nres):
+    """
+    Computes average precision for given ranked indexes.
+    Arguments
+    ---------
+    ranks : zerro-based ranks of positive images
+    nres  : number of positive images
+    Returns
+    -------
+    ap    : average precision
+    """
+
+    # number of images ranked by the system
+    nimgranks = len(ranks)
+
+    # accumulate trapezoids in PR-plot
+    ap = 0
+
+    recall_step = 1. / nres
+
+    for j in np.arange(nimgranks):
+        rank = ranks[j]
+
+        if rank == 0:
+            precision_0 = 1.
+        else:
+            precision_0 = float(j) / rank
+
+        precision_1 = float(j + 1) / (rank + 1)
+
+        ap += (precision_0 + precision_1) * recall_step / 2.
+
+    return ap
+
+
+def compute_map(ranks, gnd, kappas=[]):
+    """
+    Computes the mAP for a given set of returned results.
+         Usage:
+           map = compute_map (ranks, gnd)
+                 computes mean average precsion (map) only
+           map, aps, pr, prs = compute_map (ranks, gnd, kappas)
+                 computes mean average precision (map), average precision (aps) for each query
+                 computes mean precision at kappas (pr), precision at kappas (prs) for each query
+         Notes:
+         1) ranks starts from 0, ranks.shape = db_size X #queries
+         2) The junk results (e.g., the query itself) should be declared in the gnd stuct array
+         3) If there are no positive images for some query, that query is excluded from the evaluation
+    """
+
+    map = 0.
+    nq = len(gnd) # number of queries
+    aps = np.zeros(nq)
+    pr = np.zeros(len(kappas))
+    prs = np.zeros((nq, len(kappas)))
+    nempty = 0
+
+    for i in np.arange(nq):
+        qgnd = np.array(gnd[i]['ok'])
+
+        # no positive images, skip from the average
+        if qgnd.shape[0] == 0:
+            aps[i] = float('nan')
+            prs[i, :] = float('nan')
+            nempty += 1
+            continue
+
+        try:
+            qgndj = np.array(gnd[i]['junk'])
+        except:
+            qgndj = np.empty(0)
+
+        # sorted positions of positive and junk images (0 based)
+        pos  = np.arange(ranks.shape[0])[np.in1d(ranks[:,i], qgnd)]
+        junk = np.arange(ranks.shape[0])[np.in1d(ranks[:,i], qgndj)]
+
+        k = 0;
+        ij = 0;
+        if len(junk):
+            # decrease positions of positives based on the number of
+            # junk images appearing before them
+            ip = 0
+            while (ip < len(pos)):
+                while (ij < len(junk) and pos[ip] > junk[ij]):
+                    k += 1
+                    ij += 1
+                pos[ip] = pos[ip] - k
+                ip += 1
+
+        # compute ap
+        ap = compute_ap(pos, len(qgnd))
+        map = map + ap
+        aps[i] = ap
+
+        # compute precision @ k
+        pos += 1 # get it to 1-based
+        for j in np.arange(len(kappas)):
+            kq = min(max(pos), kappas[j]); 
+            prs[i, j] = (pos <= kq).sum() / kq
+        pr = pr + prs[i, :]
+
+    map = map / (nq - nempty)
+    pr = pr / (nq - nempty)
+
+    return map, aps, pr, prs
+
+
+def multi_scale(samples, model):
+    v = None
+    for s in [1, 1/2**(1/2), 1/2]:  # we use 3 different scales
+        if s == 1:
+            inp = samples.clone()
+        else:
+            inp = nn.functional.interpolate(samples, scale_factor=s, mode='bilinear', align_corners=False)
+        feats = model(inp).clone()
+        if v is None:
+            v = feats
+        else:
+            v += feats
+    v /= 3
+    v /= v.norm()
+    return v

PuzzleTuning/Counterpart PreTrain Methods/dino-main/video_generation.py

@@ -0,0 +1,378 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import glob
+import sys
+import argparse
+import cv2
+
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision import transforms as pth_transforms
+import numpy as np
+from PIL import Image
+
+import utils
+import vision_transformer as vits
+
+
+FOURCC = {
+    "mp4": cv2.VideoWriter_fourcc(*"MP4V"),
+    "avi": cv2.VideoWriter_fourcc(*"XVID"),
+}
+DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+
+class VideoGenerator:
+    def __init__(self, args):
+        self.args = args
+        # self.model = None
+        # Don't need to load model if you only want a video
+        if not self.args.video_only:
+            self.model = self.__load_model()
+
+    def run(self):
+        if self.args.input_path is None:
+            print(f"Provided input path {self.args.input_path} is non valid.")
+            sys.exit(1)
+        else:
+            if self.args.video_only:
+                self._generate_video_from_images(
+                    self.args.input_path, self.args.output_path
+                )
+            else:
+                # If input path exists
+                if os.path.exists(self.args.input_path):
+                    # If input is a video file
+                    if os.path.isfile(self.args.input_path):
+                        frames_folder = os.path.join(self.args.output_path, "frames")
+                        attention_folder = os.path.join(
+                            self.args.output_path, "attention"
+                        )
+
+                        os.makedirs(frames_folder, exist_ok=True)
+                        os.makedirs(attention_folder, exist_ok=True)
+
+                        self._extract_frames_from_video(
+                            self.args.input_path, frames_folder
+                        )
+
+                        self._inference(
+                            frames_folder,
+                            attention_folder,
+                        )
+
+                        self._generate_video_from_images(
+                            attention_folder, self.args.output_path
+                        )
+
+                    # If input is a folder of already extracted frames
+                    if os.path.isdir(self.args.input_path):
+                        attention_folder = os.path.join(
+                            self.args.output_path, "attention"
+                        )
+
+                        os.makedirs(attention_folder, exist_ok=True)
+
+                        self._inference(self.args.input_path, attention_folder)
+
+                        self._generate_video_from_images(
+                            attention_folder, self.args.output_path
+                        )
+
+                # If input path doesn't exists
+                else:
+                    print(f"Provided input path {self.args.input_path} doesn't exists.")
+                    sys.exit(1)
+
+    def _extract_frames_from_video(self, inp: str, out: str):
+        vidcap = cv2.VideoCapture(inp)
+        self.args.fps = vidcap.get(cv2.CAP_PROP_FPS)
+
+        print(f"Video: {inp} ({self.args.fps} fps)")
+        print(f"Extracting frames to {out}")
+
+        success, image = vidcap.read()
+        count = 0
+        while success:
+            cv2.imwrite(
+                os.path.join(out, f"frame-{count:04}.jpg"),
+                image,
+            )
+            success, image = vidcap.read()
+            count += 1
+
+    def _generate_video_from_images(self, inp: str, out: str):
+        img_array = []
+        attention_images_list = sorted(glob.glob(os.path.join(inp, "attn-*.jpg")))
+
+        # Get size of the first image
+        with open(attention_images_list[0], "rb") as f:
+            img = Image.open(f)
+            img = img.convert("RGB")
+            size = (img.width, img.height)
+            img_array.append(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+
+        print(f"Generating video {size} to {out}")
+
+        for filename in tqdm(attention_images_list[1:]):
+            with open(filename, "rb") as f:
+                img = Image.open(f)
+                img = img.convert("RGB")
+                img_array.append(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+
+        out = cv2.VideoWriter(
+            os.path.join(out, "video." + self.args.video_format),
+            FOURCC[self.args.video_format],
+            self.args.fps,
+            size,
+        )
+
+        for i in range(len(img_array)):
+            out.write(img_array[i])
+        out.release()
+        print("Done")
+
+    def _inference(self, inp: str, out: str):
+        print(f"Generating attention images to {out}")
+
+        for img_path in tqdm(sorted(glob.glob(os.path.join(inp, "*.jpg")))):
+            with open(img_path, "rb") as f:
+                img = Image.open(f)
+                img = img.convert("RGB")
+
+            if self.args.resize is not None:
+                transform = pth_transforms.Compose(
+                    [
+                        pth_transforms.ToTensor(),
+                        pth_transforms.Resize(self.args.resize),
+                        pth_transforms.Normalize(
+                            (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
+                        ),
+                    ]
+                )
+            else:
+                transform = pth_transforms.Compose(
+                    [
+                        pth_transforms.ToTensor(),
+                        pth_transforms.Normalize(
+                            (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
+                        ),
+                    ]
+                )
+
+            img = transform(img)
+
+            # make the image divisible by the patch size
+            w, h = (
+                img.shape[1] - img.shape[1] % self.args.patch_size,
+                img.shape[2] - img.shape[2] % self.args.patch_size,
+            )
+            img = img[:, :w, :h].unsqueeze(0)
+
+            w_featmap = img.shape[-2] // self.args.patch_size
+            h_featmap = img.shape[-1] // self.args.patch_size
+
+            attentions = self.model.get_last_selfattention(img.to(DEVICE))
+
+            nh = attentions.shape[1]  # number of head
+
+            # we keep only the output patch attention
+            attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
+
+            # we keep only a certain percentage of the mass
+            val, idx = torch.sort(attentions)
+            val /= torch.sum(val, dim=1, keepdim=True)
+            cumval = torch.cumsum(val, dim=1)
+            th_attn = cumval > (1 - self.args.threshold)
+            idx2 = torch.argsort(idx)
+            for head in range(nh):
+                th_attn[head] = th_attn[head][idx2[head]]
+            th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
+            # interpolate
+            th_attn = (
+                nn.functional.interpolate(
+                    th_attn.unsqueeze(0),
+                    scale_factor=self.args.patch_size,
+                    mode="nearest",
+                )[0]
+                .cpu()
+                .numpy()
+            )
+
+            attentions = attentions.reshape(nh, w_featmap, h_featmap)
+            attentions = (
+                nn.functional.interpolate(
+                    attentions.unsqueeze(0),
+                    scale_factor=self.args.patch_size,
+                    mode="nearest",
+                )[0]
+                .cpu()
+                .numpy()
+            )
+
+            # save attentions heatmaps
+            fname = os.path.join(out, "attn-" + os.path.basename(img_path))
+            plt.imsave(
+                fname=fname,
+                arr=sum(
+                    attentions[i] * 1 / attentions.shape[0]
+                    for i in range(attentions.shape[0])
+                ),
+                cmap="inferno",
+                format="jpg",
+            )
+
+    def __load_model(self):
+        # build model
+        model = vits.__dict__[self.args.arch](
+            patch_size=self.args.patch_size, num_classes=0
+        )
+        for p in model.parameters():
+            p.requires_grad = False
+        model.eval()
+        model.to(DEVICE)
+
+        if os.path.isfile(self.args.pretrained_weights):
+            state_dict = torch.load(self.args.pretrained_weights, map_location="cpu")
+            if (
+                self.args.checkpoint_key is not None
+                and self.args.checkpoint_key in state_dict
+            ):
+                print(
+                    f"Take key {self.args.checkpoint_key} in provided checkpoint dict"
+                )
+                state_dict = state_dict[self.args.checkpoint_key]
+            state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+            # remove `backbone.` prefix induced by multicrop wrapper
+            state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+            msg = model.load_state_dict(state_dict, strict=False)
+            print(
+                "Pretrained weights found at {} and loaded with msg: {}".format(
+                    self.args.pretrained_weights, msg
+                )
+            )
+        else:
+            print(
+                "Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate."
+            )
+            url = None
+            if self.args.arch == "vit_small" and self.args.patch_size == 16:
+                url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
+            elif self.args.arch == "vit_small" and self.args.patch_size == 8:
+                url = "dino_deitsmall8_300ep_pretrain/dino_deitsmall8_300ep_pretrain.pth"  # model used for visualizations in our paper
+            elif self.args.arch == "vit_base" and self.args.patch_size == 16:
+                url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
+            elif self.args.arch == "vit_base" and self.args.patch_size == 8:
+                url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
+            if url is not None:
+                print(
+                    "Since no pretrained weights have been provided, we load the reference pretrained DINO weights."
+                )
+                state_dict = torch.hub.load_state_dict_from_url(
+                    url="https://dl.fbaipublicfiles.com/dino/" + url
+                )
+                model.load_state_dict(state_dict, strict=True)
+            else:
+                print(
+                    "There is no reference weights available for this model => We use random weights."
+                )
+        return model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("Generation self-attention video")
+    parser.add_argument(
+        "--arch",
+        default="vit_small",
+        type=str,
+        choices=["vit_tiny", "vit_small", "vit_base"],
+        help="Architecture (support only ViT atm).",
+    )
+    parser.add_argument(
+        "--patch_size", default=8, type=int, help="Patch resolution of the self.model."
+    )
+    parser.add_argument(
+        "--pretrained_weights",
+        default="",
+        type=str,
+        help="Path to pretrained weights to load.",
+    )
+    parser.add_argument(
+        "--checkpoint_key",
+        default="teacher",
+        type=str,
+        help='Key to use in the checkpoint (example: "teacher")',
+    )
+    parser.add_argument(
+        "--input_path",
+        required=True,
+        type=str,
+        help="""Path to a video file if you want to extract frames
+            or to a folder of images already extracted by yourself.
+            or to a folder of attention images.""",
+    )
+    parser.add_argument(
+        "--output_path",
+        default="./",
+        type=str,
+        help="""Path to store a folder of frames and / or a folder of attention images.
+            and / or a final video. Default to current directory.""",
+    )
+    parser.add_argument(
+        "--threshold",
+        type=float,
+        default=0.6,
+        help="""We visualize masks
+        obtained by thresholding the self-attention maps to keep xx percent of the mass.""",
+    )
+    parser.add_argument(
+        "--resize",
+        default=None,
+        type=int,
+        nargs="+",
+        help="""Apply a resize transformation to input image(s). Use if OOM error.
+        Usage (single or W H): --resize 512, --resize 720 1280""",
+    )
+    parser.add_argument(
+        "--video_only",
+        action="store_true",
+        help="""Use this flag if you only want to generate a video and not all attention images.
+            If used, --input_path must be set to the folder of attention images. Ex: ./attention/""",
+    )
+    parser.add_argument(
+        "--fps",
+        default=30.0,
+        type=float,
+        help="FPS of input / output video. Automatically set if you extract frames from a video.",
+    )
+    parser.add_argument(
+        "--video_format",
+        default="mp4",
+        type=str,
+        choices=["mp4", "avi"],
+        help="Format of generated video (mp4 or avi).",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    vg = VideoGenerator(args)
+    vg.run()

PuzzleTuning/Counterpart PreTrain Methods/dino-main/vision_transformer.py

@@ -0,0 +1,291 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+from utils import trunc_normal_
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    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):
+    """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)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=256):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            self.mlp = nn.Linear(in_dim, bottleneck_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.last_layer(x)
+        return x

PuzzleTuning/Counterpart PreTrain Methods/dino-main/visualize_attention.py

@@ -0,0 +1,213 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import sys
+import argparse
+import cv2
+import random
+import colorsys
+import requests
+from io import BytesIO
+
+import skimage.io
+from skimage.measure import find_contours
+import matplotlib.pyplot as plt
+from matplotlib.patches import Polygon
+import torch
+import torch.nn as nn
+import torchvision
+from torchvision import transforms as pth_transforms
+import numpy as np
+from PIL import Image
+
+import utils
+import vision_transformer as vits
+
+
+def apply_mask(image, mask, color, alpha=0.5):
+    for c in range(3):
+        image[:, :, c] = image[:, :, c] * (1 - alpha * mask) + alpha * mask * color[c] * 255
+    return image
+
+
+def random_colors(N, bright=True):
+    """
+    Generate random colors.
+    """
+    brightness = 1.0 if bright else 0.7
+    hsv = [(i / N, 1, brightness) for i in range(N)]
+    colors = list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv))
+    random.shuffle(colors)
+    return colors
+
+
+def display_instances(image, mask, fname="test", figsize=(5, 5), blur=False, contour=True, alpha=0.5):
+    fig = plt.figure(figsize=figsize, frameon=False)
+    ax = plt.Axes(fig, [0., 0., 1., 1.])
+    ax.set_axis_off()
+    fig.add_axes(ax)
+    ax = plt.gca()
+
+    N = 1
+    mask = mask[None, :, :]
+    # Generate random colors
+    colors = random_colors(N)
+
+    # Show area outside image boundaries.
+    height, width = image.shape[:2]
+    margin = 0
+    ax.set_ylim(height + margin, -margin)
+    ax.set_xlim(-margin, width + margin)
+    ax.axis('off')
+    masked_image = image.astype(np.uint32).copy()
+    for i in range(N):
+        color = colors[i]
+        _mask = mask[i]
+        if blur:
+            _mask = cv2.blur(_mask,(10,10))
+        # Mask
+        masked_image = apply_mask(masked_image, _mask, color, alpha)
+        # Mask Polygon
+        # Pad to ensure proper polygons for masks that touch image edges.
+        if contour:
+            padded_mask = np.zeros((_mask.shape[0] + 2, _mask.shape[1] + 2))
+            padded_mask[1:-1, 1:-1] = _mask
+            contours = find_contours(padded_mask, 0.5)
+            for verts in contours:
+                # Subtract the padding and flip (y, x) to (x, y)
+                verts = np.fliplr(verts) - 1
+                p = Polygon(verts, facecolor="none", edgecolor=color)
+                ax.add_patch(p)
+    ax.imshow(masked_image.astype(np.uint8), aspect='auto')
+    fig.savefig(fname)
+    print(f"{fname} saved.")
+    return
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('Visualize Self-Attention maps')
+    parser.add_argument('--arch', default='vit_small', type=str,
+        choices=['vit_tiny', 'vit_small', 'vit_base'], help='Architecture (support only ViT atm).')
+    parser.add_argument('--patch_size', default=8, type=int, help='Patch resolution of the model.')
+    parser.add_argument('--pretrained_weights', default='', type=str,
+        help="Path to pretrained weights to load.")
+    parser.add_argument("--checkpoint_key", default="teacher", type=str,
+        help='Key to use in the checkpoint (example: "teacher")')
+    parser.add_argument("--image_path", default=None, type=str, help="Path of the image to load.")
+    parser.add_argument("--image_size", default=(480, 480), type=int, nargs="+", help="Resize image.")
+    parser.add_argument('--output_dir', default='.', help='Path where to save visualizations.')
+    parser.add_argument("--threshold", type=float, default=None, help="""We visualize masks
+        obtained by thresholding the self-attention maps to keep xx% of the mass.""")
+    args = parser.parse_args()
+
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    # build model
+    model = vits.__dict__[args.arch](patch_size=args.patch_size, num_classes=0)
+    for p in model.parameters():
+        p.requires_grad = False
+    model.eval()
+    model.to(device)
+    if os.path.isfile(args.pretrained_weights):
+        state_dict = torch.load(args.pretrained_weights, map_location="cpu")
+        if args.checkpoint_key is not None and args.checkpoint_key in state_dict:
+            print(f"Take key {args.checkpoint_key} in provided checkpoint dict")
+            state_dict = state_dict[args.checkpoint_key]
+        # remove `module.` prefix
+        state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()}
+        # remove `backbone.` prefix induced by multicrop wrapper
+        state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()}
+        msg = model.load_state_dict(state_dict, strict=False)
+        print('Pretrained weights found at {} and loaded with msg: {}'.format(args.pretrained_weights, msg))
+    else:
+        print("Please use the `--pretrained_weights` argument to indicate the path of the checkpoint to evaluate.")
+        url = None
+        if args.arch == "vit_small" and args.patch_size == 16:
+            url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
+        elif args.arch == "vit_small" and args.patch_size == 8:
+            url = "dino_deitsmall8_300ep_pretrain/dino_deitsmall8_300ep_pretrain.pth"  # model used for visualizations in our paper
+        elif args.arch == "vit_base" and args.patch_size == 16:
+            url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
+        elif args.arch == "vit_base" and args.patch_size == 8:
+            url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
+        if url is not None:
+            print("Since no pretrained weights have been provided, we load the reference pretrained DINO weights.")
+            state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/" + url)
+            model.load_state_dict(state_dict, strict=True)
+        else:
+            print("There is no reference weights available for this model => We use random weights.")
+
+    # open image
+    if args.image_path is None:
+        # user has not specified any image - we use our own image
+        print("Please use the `--image_path` argument to indicate the path of the image you wish to visualize.")
+        print("Since no image path have been provided, we take the first image in our paper.")
+        response = requests.get("https://dl.fbaipublicfiles.com/dino/img.png")
+        img = Image.open(BytesIO(response.content))
+        img = img.convert('RGB')
+    elif os.path.isfile(args.image_path):
+        with open(args.image_path, 'rb') as f:
+            img = Image.open(f)
+            img = img.convert('RGB')
+    else:
+        print(f"Provided image path {args.image_path} is non valid.")
+        sys.exit(1)
+    transform = pth_transforms.Compose([
+        pth_transforms.Resize(args.image_size),
+        pth_transforms.ToTensor(),
+        pth_transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+    ])
+    img = transform(img)
+
+    # make the image divisible by the patch size
+    w, h = img.shape[1] - img.shape[1] % args.patch_size, img.shape[2] - img.shape[2] % args.patch_size
+    img = img[:, :w, :h].unsqueeze(0)
+
+    w_featmap = img.shape[-2] // args.patch_size
+    h_featmap = img.shape[-1] // args.patch_size
+
+    attentions = model.get_last_selfattention(img.to(device))
+
+    nh = attentions.shape[1] # number of head
+
+    # we keep only the output patch attention
+    attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
+
+    if args.threshold is not None:
+        # we keep only a certain percentage of the mass
+        val, idx = torch.sort(attentions)
+        val /= torch.sum(val, dim=1, keepdim=True)
+        cumval = torch.cumsum(val, dim=1)
+        th_attn = cumval > (1 - args.threshold)
+        idx2 = torch.argsort(idx)
+        for head in range(nh):
+            th_attn[head] = th_attn[head][idx2[head]]
+        th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
+        # interpolate
+        th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=args.patch_size, mode="nearest")[0].cpu().numpy()
+
+    attentions = attentions.reshape(nh, w_featmap, h_featmap)
+    attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=args.patch_size, mode="nearest")[0].cpu().numpy()
+
+    # save attentions heatmaps
+    os.makedirs(args.output_dir, exist_ok=True)
+    torchvision.utils.save_image(torchvision.utils.make_grid(img, normalize=True, scale_each=True), os.path.join(args.output_dir, "img.png"))
+    for j in range(nh):
+        fname = os.path.join(args.output_dir, "attn-head" + str(j) + ".png")
+        plt.imsave(fname=fname, arr=attentions[j], format='png')
+        print(f"{fname} saved.")
+
+    if args.threshold is not None:
+        image = skimage.io.imread(os.path.join(args.output_dir, "img.png"))
+        for j in range(nh):
+            display_instances(image, th_attn[j], fname=os.path.join(args.output_dir, "mask_th" + str(args.threshold) + "_head" + str(j) +".png"), blur=False)

PuzzleTuning/Counterpart PreTrain Methods/gcmae/FINETUNE.md

@@ -0,0 +1,13 @@
+## Fine tune GCMAE 
+
+```
+python main_finetune.py \
+    --data_path path/to/data \
+    --nb_classes 9 \
+    --output_dir path/to/ouput/dir \
+    --log_dir path/to/log/dir \
+    --batch_size 128 \
+    --model vit_base_patch16 \
+    --epochs 50 \
+    --finetune path/to/pth/path \
+```

PuzzleTuning/Counterpart PreTrain Methods/gcmae/LICENSE

@@ -0,0 +1,399 @@
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+   wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More_considerations
+     for the public: 
+   wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

PuzzleTuning/Counterpart PreTrain Methods/gcmae/LINPROBE.md

@@ -0,0 +1,15 @@
+## Linear probe GCMAE 
+
+```
+python main_linprobe.py \
+    --data_path_train path/to/train/data \
+    --data_path_val path/to/val/data \
+    --nb_classes 2 \
+    --output_dir path/to/ouput/dir \
+    --log_dir path/to/log/dir \
+    --batch_size 512 \
+    --model vit_base_patch16 \
+    --epochs 90 \
+    --finetune path/to/pth/path
+```
+

PuzzleTuning/Counterpart PreTrain Methods/gcmae/PRETRAIN.md

@@ -0,0 +1,21 @@
+## Pre-training GCMAE
+
+To pre-train ViT-Base (recommended default)
+```
+python main_pretrain.py \
+    --data_path path/to/data \
+    --data_val_path path/to/data \
+    --output_dir path/to/ouput/dir \
+    --log_dir path/to/log/dir \
+    --batch_size 128 \
+    --model gcmae_vit_base_patch16 \
+    --norm_pix_loss \
+    --mask_ratio 0.5 \
+    --epochs 80 \
+    --warmup_epochs 40 \
+    --blr 1e-3 --weight_decay 0.05 \
+    --low_dim 768 \
+    --nce_k 8192 \
+    --nce_t 0.07 \
+    --nce_m 0.5 \
+```

PuzzleTuning/Counterpart PreTrain Methods/gcmae/README.md

@@ -0,0 +1,15 @@
+# GCMAE
+
+The original repo of GCMAE could be found [here](https://github.com/StarUniversus/gcmae)
+
+To install environments:
+```bash
+pip install -r requirements.txt
+```
+
+
+To start pretraining:
+```bash
+# You need to alter the script according to your directories
+bash pretrain.sh
+```

PuzzleTuning/Counterpart PreTrain Methods/gcmae/README_origin.md

@@ -0,0 +1,39 @@
+# GCMAE
+
+<p align="center">
+  <img src="./GCMAE.png" width="800">
+</p>
+
+The official implementation of the paper [Global Contrast Masked Autoencoders Are Powerful Pathological Representation Learners](https://arxiv.org/abs/2205.09048)
+
+```
+@article{li2022gcmae,
+  author = {Quan, Hao and Li, Xingyu and Chen, Weixing and Bai, Qun and Zou, Mingchen and Yang, Ruijie and Zheng, Tingting and Qi, Ruiqun and Gao, Xinghua and Cui, Xiaoyu},
+  title = {Global Contrast Masked Autoencoders Are Powerful Pathological Representation Learners},
+  journal={arXiv:2205.09048},
+  year = {2022},
+}
+```
+## Abstract
+Based on digital whole slide scanning technique, artificial intelligence algorithms represented by deep learning have achieved remarkable results in the field of computational pathology. Compared with other medical images such as Computed Tomography (CT) or Magnetic Resonance Imaging (MRI), pathological images are more difficult to annotate, thus there is an extreme lack of data sets that can be used for supervised learning. In this study, a self-supervised learning (SSL) model, Global Contrast Masked Autoencoders (GCMAE), is proposed, which has the ability to represent both global and local domain-specific features of whole slide image (WSI), as well as excellent cross-data transfer ability. The Camelyon16 and NCTCRC datasets are used to evaluate the performance of our model. When dealing with transfer learning tasks with different data sets, the experimental results show that GCMAE has better linear classification accuracy than MAE, which can reach **81.10%** and **89.22%** respectively. Our method outperforms the previous state of-the-art algorithm and even surpass supervised learning (improved by **3.86%** on NCTCRC data sets).
+
+## Installation
+This repo is a modification on the [mae repo](https://github.com/facebookresearch/mae). Installation and preparation follow that repo.
+
+## Usage
+
+* [PRETRAIN](PRETRAIN.md)
+
+* [LINPROBE](LINPROBE.md)
+
+* [FINETUNE](FINETUNE.md)
+
+* [ Visual GCMAE feature representation](VISUAL.md)
+
+## Dataset
+  * [Camelyon16](https://pan.baidu.com/s/1N0fqJR9u8yq-y6ZY0mSoUw?pwd=noms)
+  * [NCT-CRC-HE-100K](https://zenodo.org/record/1214456)
+  * [BreakHis](https://web.inf.ufpr.br/vri/databases/breast-cancer-histopathological-database-breakhis/)
+
+## License
+Distributed under the CC-BY-NC 4.0 License. See [LICENSE](LICENSE) for more information.

PuzzleTuning/Counterpart PreTrain Methods/gcmae/VISUAL.md

@@ -0,0 +1,10 @@
+## Visual GCMAE feature representation
+
+```
+python tsne.py \
+    --batch_size 128 \
+    --model vit_base_patch16 \
+    --finetune path/to/pth \
+    --save_path path/to/png \ 
+    --data_path_val path/to/data \
+```

PuzzleTuning/Counterpart PreTrain Methods/gcmae/engine_finetune.py

@@ -0,0 +1,183 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# MAE: https://github.com/facebookresearch/mae
+# --------------------------------------------------------
+
+import math
+import sys
+from typing import Iterable, Optional
+
+import torch
+
+from timm.data import Mixup
+from timm.utils import accuracy
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+from sklearn.metrics import roc_auc_score
+from sklearn.metrics import precision_score, recall_score, f1_score, confusion_matrix
+import torch.nn.functional as F
+import numpy
+import numpy as np
+from torchmetrics import Specificity, AUROC
+
+
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler, max_norm: float = 0,
+                    mixup_fn: Optional[Mixup] = None, log_writer=None,
+                    args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    for data_iter_step, (samples, targets) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+
+        if mixup_fn is not None:
+            samples, targets = mixup_fn(samples, targets)
+
+        with torch.cuda.amp.autocast():
+            outputs = model(samples)
+            loss = criterion(outputs, targets)
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, clip_grad=max_norm,
+                    parameters=model.parameters(), create_graph=False,
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(loss=loss_value)
+        min_lr = 10.
+        max_lr = 0.
+        for group in optimizer.param_groups:
+            min_lr = min(min_lr, group["lr"])
+            max_lr = max(max_lr, group["lr"])
+
+        metric_logger.update(lr=max_lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', max_lr, epoch_1000x)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+def confusion_m(y_true, y_pred):
+    tn, fp, fn, tp = confusion_matrix(y_true, y_pred).ravel()
+    return tn, fp, fn, tp
+
+def compute_metrics_binary(probs, preds, targets):
+    auc = roc_auc_score(targets, probs) * 100  
+    precision = precision_score(targets, preds) * 100
+    recall = recall_score(targets, preds) * 100
+    f1 = f1_score(targets, preds) * 100
+    tn, fp, fn, tp = confusion_m(targets, preds)
+    specificity = (tn / float(tn+fp)) * 100
+    return auc, precision, recall, f1, specificity
+
+def compute_metrics_multiclass(probs, preds, targets, nb_classes):
+    preds_tensor, probs_tensor, targets_tensor = torch.tensor(preds), torch.tensor(probs), torch.tensor(targets)
+    auroc = AUROC(average='macro', num_classes=nb_classes)
+    auc = auroc(probs_tensor, targets_tensor) * 100
+    precision = precision_score(targets, preds, average='macro') * 100
+    recall = recall_score(targets, preds, average='macro') * 100
+    f1 = f1_score(targets, preds, average='macro') * 100
+    speci = Specificity(average='macro', num_classes=nb_classes)
+    specificity = speci(preds_tensor, targets_tensor) * 100
+    return auc, precision, recall, f1, specificity
+
+@torch.no_grad()
+def evaluate(data_loader, model, device, nb_classes):
+    criterion = torch.nn.CrossEntropyLoss()
+    m = torch.nn.Softmax(dim=1)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    # switch to evaluation mode
+    model.eval()
+    probs = []
+    targets = []
+    preds = []
+
+    for batch in metric_logger.log_every(data_loader, 10, header):
+        images = batch[0]
+        target = batch[-1]
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        # compute output
+        with torch.cuda.amp.autocast():
+            output = model(images)
+            loss = criterion(output, target)
+        output = m(output)
+        score, pred = output.topk(1, 1, True, True)
+        if nb_classes == 2:
+            prob = output[:, 1]
+        elif nb_classes > 2:
+            prob = output
+            
+        probs.extend(prob.detach().cpu().numpy())
+        targets.extend(target.detach().cpu().numpy())
+        preds.extend(pred.tolist())
+        if nb_classes < 5:
+            acc1, acc5 = accuracy(output, target, topk=(1, 1))
+        else:
+            acc1, acc5 = accuracy(output, target, topk=(1, 5))
+
+        batch_size = images.shape[0]
+        metric_logger.update(loss=loss.item())
+        metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+        metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print('* Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f} loss {losses.global_avg:.3f}'
+          .format(top1=metric_logger.acc1, top5=metric_logger.acc5, losses=metric_logger.loss))
+    if nb_classes == 2:
+        print("binary class metrics!")
+        auc, precision, recall, f1, specificity = compute_metrics_binary(probs, preds, targets)
+
+    elif nb_classes > 2:
+        print("multi_class metrics!")
+        auc, precision, recall, f1, specificity = compute_metrics_multiclass(probs, preds, targets, nb_classes)
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, auc, precision, recall, f1, specificity

PuzzleTuning/Counterpart PreTrain Methods/gcmae/engine_pretrain.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# BEiT: https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+from test_npid import NN, kNN
+
+def train_one_epoch(model: torch.nn.Module,
+                    data_loader: Iterable,
+                    optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    log_writer=None,
+                    args=None,
+                    lemniscate=None,
+                    ):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    for data_iter_step, (samples, _, index) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        samples = samples.to(device, non_blocking=True)
+        index = index.to(device, non_blocking=True)
+        with torch.cuda.amp.autocast():
+            loss_mae, _, _, loss_npid, _= model(samples, mask_ratio=args.mask_ratio, index = index, is_train=True)
+            loss = loss_mae + 0.1 * loss_npid
+
+        loss_value = loss.item()
+
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+
+        metric_logger.update(loss_all=loss_value)
+        metric_logger.update(loss_mae=loss_mae.item())
+        metric_logger.update(loss_npid=loss_npid.item())
+
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('train_loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('origin_loss/train_loss_mae', loss_mae.item(), epoch_1000x)
+            log_writer.add_scalar('origin_loss/train_loss_npid', loss_npid.item(), epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+
+    # pred1 = NN(epoch, model, lemniscate, data_loader, data_loader_val)
+    # log_writer.add_scalar('NN_ac', pred1, epoch)
+    # if args.output_dir and (epoch % 20 == 0 or epoch + 1 == args.epochs):
+
+    #     top1 = kNN(0, model, lemniscate, data_loader, data_loader_val, 200, args.nce_t)
+    #     log_writer.add_scalar('KNN_top1', top1, epoch)
+
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    # return {k: meter.global_avg for k, meter in metric_logger.meters.items()}, pred1
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

PuzzleTuning/Counterpart PreTrain Methods/gcmae/lib/NCEAverage.py

@@ -0,0 +1,94 @@
+import torch
+from torch.autograd import Function
+from torch import nn
+from .alias_multinomial import AliasMethod
+import math
+
+class NCEFunction(Function):
+    @staticmethod
+    def forward(self, x, y, memory, idx, params):
+        K = int(params[0].item())
+        T = params[1].item()
+        Z = params[2].item()
+
+        momentum = params[3].item()
+        batchSize = x.size(0)
+        outputSize = memory.size(0)
+        inputSize = memory.size(1)
+
+        # sample positives & negatives
+        idx.select(1,0).copy_(y.detach())
+
+        # sample correspoinding weights
+        weight = torch.index_select(memory, 0, idx.view(-1))
+        weight.resize_(batchSize, K+1, inputSize)
+
+        # inner product
+        with torch.cuda.amp.autocast(enabled = False):
+            out = torch.bmm(weight, x.detach().reshape(batchSize, inputSize, 1))
+        #print(out)
+        out.div_(T).exp_() # batchSize * self.K+1
+        #x.detach().resize_(batchSize, inputSize)
+        #print(out)
+        if Z < 0:
+            params[2] = out.mean() * outputSize
+            Z = params[2].item()
+            print("normalization constant Z is set to {:.1f}".format(Z))
+
+        out.div_(Z).resize_(batchSize, K+1)
+
+        self.save_for_backward(x, memory, y, weight, out, params)
+
+        return out
+
+    @staticmethod
+    def backward(self, gradOutput):
+        x, memory, y, weight, out, params = self.saved_tensors
+        K = int(params[0].item())
+        T = params[1].item()
+        Z = params[2].item()
+        momentum = params[3].item()
+        batchSize = gradOutput.size(0)
+        
+        # gradients d Pm / d linear = exp(linear) / Z
+        gradOutput.detach().mul_(out.detach())
+        # add temperature
+        gradOutput.detach().div_(T)
+
+        
+        
+        # gradient of linear
+        with torch.cuda.amp.autocast(enabled = False):
+            gradInput = torch.bmm(gradOutput.detach().reshape(batchSize, 1, K+1), weight)
+        gradInput.resize_as_(x)
+
+        # update the non-parametric detach()
+        weight_pos = weight.select(1, 0).resize_as_(x)
+        weight_pos.mul_(momentum)
+        weight_pos.add_(torch.mul(x.detach(), 1-momentum))
+        w_norm = weight_pos.pow(2).sum(1, keepdim=True).pow(0.5)
+        updated_weight = weight_pos.div(w_norm)
+        memory.index_copy_(0, y, updated_weight)
+        
+        return gradInput, None, None, None, None
+
+class NCEAverage(nn.Module):
+
+    def __init__(self, inputSize, outputSize, K, T=0.07, momentum=0.5, Z=None):
+        super(NCEAverage, self).__init__()
+        self.nLem = outputSize
+        self.unigrams = torch.ones(self.nLem)
+        self.multinomial = AliasMethod(self.unigrams)
+        self.multinomial.cuda()
+        self.K = K
+
+        self.register_buffer('params',torch.tensor([K, T, -1, momentum]))
+        stdv = 1. / math.sqrt(inputSize/3)
+        self.register_buffer('memory', torch.rand(outputSize, inputSize).mul_(2*stdv).add_(-stdv))
+ 
+    def forward(self, x, y):
+        batchSize = x.size(0)
+        idx = self.multinomial.draw(batchSize * (self.K+1)).view(batchSize, -1)
+        out = NCEFunction.apply(x, y, self.memory, idx, self.params)
+        return out
+

PuzzleTuning/Counterpart PreTrain Methods/gcmae/lib/NCECriterion.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+
+eps = 1e-7
+
+class NCECriterion(nn.Module):
+
+    def __init__(self, nLem):
+        super(NCECriterion, self).__init__()
+        self.nLem = nLem
+
+    def forward(self, x, targets):
+        batchSize = x.size(0)
+        K = x.size(1)-1
+        Pnt = 1 / float(self.nLem)
+        Pns = 1 / float(self.nLem)
+        
+        # eq 5.1 : P(origin=model) = Pmt / (Pmt + k*Pnt) 
+        Pmt = x.select(1,0)
+        Pmt_div = Pmt.add(K * Pnt + eps)
+        lnPmt = torch.div(Pmt, Pmt_div)
+        
+        # eq 5.2 : P(origin=noise) = k*Pns / (Pms + k*Pns)
+        Pon_div = x.narrow(1,1,K).add(K * Pns + eps)
+        Pon = Pon_div.clone().fill_(K * Pns)
+        lnPon = torch.div(Pon, Pon_div)
+     
+        # equation 6 in ref. A
+        lnPmt.log_()
+        lnPon.log_()
+        
+        lnPmtsum = lnPmt.sum(0)
+        lnPonsum = lnPon.view(-1, 1).sum(0)
+        
+        loss = - (lnPmtsum + lnPonsum) / batchSize
+        
+        return loss
+

PuzzleTuning/Counterpart PreTrain Methods/gcmae/lib/__init__.py

@@ -0,0 +1 @@
+# nothing

PuzzleTuning/Counterpart PreTrain Methods/gcmae/lib/alias_multinomial.py

@@ -0,0 +1,64 @@
+import torch
+import numpy as np
+
+class AliasMethod(object):
+    '''
+        From: https://hips.seas.harvard.edu/blog/2013/03/03/the-alias-method-efficient-sampling-with-many-discrete-outcomes/
+    '''
+    def __init__(self, probs):
+
+        if probs.sum() > 1:
+            probs.div_(probs.sum())
+        K = len(probs)
+        self.prob = torch.zeros(K)
+        self.alias = torch.LongTensor([0]*K)
+
+        # Sort the data into the outcomes with probabilities
+        # that are larger and smaller than 1/K.
+        smaller = []
+        larger = []
+        for kk, prob in enumerate(probs):
+            self.prob[kk] = K*prob
+            if self.prob[kk] < 1.0:
+                smaller.append(kk)
+            else:
+                larger.append(kk)
+
+        # Loop though and create little binary mixtures that
+        # appropriately allocate the larger outcomes over the
+        # overall uniform mixture.
+        while len(smaller) > 0 and len(larger) > 0:
+            small = smaller.pop()
+            large = larger.pop()
+
+            self.alias[small] = large
+            self.prob[large] = (self.prob[large] - 1.0) + self.prob[small]
+
+            if self.prob[large] < 1.0:
+                smaller.append(large)
+            else:
+                larger.append(large)
+
+        for last_one in smaller+larger:
+            self.prob[last_one] = 1
+
+    def cuda(self): 
+        self.prob = self.prob.cuda()
+        self.alias = self.alias.cuda()
+
+    def draw(self, N):
+        '''
+            Draw N samples from multinomial
+        '''
+        K = self.alias.size(0)
+
+        kk = torch.zeros(N, dtype=torch.long, device=self.prob.device).random_(0, K)
+        prob = self.prob.index_select(0, kk)
+        alias = self.alias.index_select(0, kk)
+        # b is whether a random number is greater than q
+        b = torch.bernoulli(prob)
+        oq = kk.mul(b.long())
+        oj = alias.mul((1-b).long())
+
+        return oq + oj
+

PuzzleTuning/Counterpart PreTrain Methods/gcmae/lib/normalize.py

@@ -0,0 +1,14 @@
+import torch
+from torch.autograd import Variable
+from torch import nn
+
+class Normalize(nn.Module):
+
+    def __init__(self, power=2):
+        super(Normalize, self).__init__()
+        self.power = power
+    
+    def forward(self, x):
+        norm = x.pow(self.power).sum(1, keepdim=True).pow(1./self.power)
+        out = x.div(norm)
+        return out