import ml_collections

def get_dogs_config():
    config = ml_collections.ConfigDict()
    config.Name = "StanfordDogs"
    config.Num_Classes = 120
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_flowers_config():
    config = ml_collections.ConfigDict()
    config.Name = "OxfordFlowers"
    config.Num_Classes = 102
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_CUB_config():
    config = ml_collections.ConfigDict()
    config.Name = "CUB_200_2011"
    config.Num_Classes = 200
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_Cars_config():
    config = ml_collections.ConfigDict()
    config.Name = "StanfordCars"
    config.Num_Classes = 196
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_NABirds_config():
    config = ml_collections.ConfigDict()
    config.Name = "NABirds"
    config.Num_Classes = 555
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

DATA_CONFIGS ={
    "CUB_200_2011": get_CUB_config(),
    "NABirds": get_NABirds_config(),
    "OxfordFlowers": get_flowers_config(),
    "StanfordCars": get_Cars_config(),
    "StanfordDogs": get_dogs_config()
}
Data_process/FGVC_config.py
DavidYanAnDe/ARC
#!/usr/bin/env python3

"""Data_process loader."""
import logging
import torch

from torch.utils.data.sampler import RandomSampler, SequentialSampler

from Data_process.json_dataset import CUB200Dataset, CarsDataset, DogsDataset, FlowersDataset, NabirdsDataset

logger = logging.getLogger(__name__)

_DATASET_CATALOG = {
    "CUB_200_2011": CUB200Dataset,
    'OxfordFlowers': FlowersDataset,
    'StanfordCars': CarsDataset,
    'StanfordDogs': DogsDataset,
    "NABirds": NabirdsDataset,
}


def _construct_loader(cfg, split, batch_size, shuffle, drop_last, data_path):
    """Constructs the data loader for the given dataset."""
    dataset_name = cfg.Name

    assert (
        dataset_name in _DATASET_CATALOG.keys()
    ), "Dataset '{}' not supported".format(dataset_name)
    dataset = _DATASET_CATALOG[dataset_name](cfg, split, data_path)

    # Create a sampler for training
    if "train" in split:
        sampler = RandomSampler(dataset)
    else:
        sampler = SequentialSampler(dataset)
    # Create a loader
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(False if sampler else shuffle),
        sampler=sampler,
        num_workers=cfg.Num_workers,
        pin_memory=cfg.pin_memory,
        drop_last=drop_last,
    )
    return loader


def construct_train_loader(cfg, batch_size, data_path):
    """Train loader wrapper."""
    drop_last = False

    return _construct_loader(
        cfg=cfg,
        split="train",
        batch_size=batch_size,
        shuffle=True,
        drop_last=drop_last,
        data_path=data_path
    )

def construct_trainval_loader(cfg, batch_size, data_path):
    """Train loader wrapper."""
    drop_last = False
    return _construct_loader(
        cfg=cfg,
        split="trainval",
        batch_size=batch_size,
        shuffle=True,
        drop_last=drop_last,
        data_path=data_path
    )

def construct_test_loader(cfg, batch_size, data_path):
    """Test loader wrapper."""
    return _construct_loader(
        cfg=cfg,
        split="test",
        batch_size= batch_size,
        shuffle=False,
        drop_last=False,
        data_path=data_path
    )


def construct_val_loader(cfg, batch_size, data_path):
    """Validation loader wrapper."""
    return _construct_loader(
        cfg=cfg,
        split="val",
        batch_size=batch_size,
        shuffle=False,
        drop_last=False,
        data_path=data_path
    )

Data_process/FGVC_loader.py
DavidYanAnDe/ARC
import ml_collections


def get_caltech101_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 102
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_cifar_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 100
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


def get_clevr_count_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 8
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_clevr_dist_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 6
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_diabetic_retinopathy_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 5
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_dmlab_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 6
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_dsprites_loc_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 16
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_dsprites_ori_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 16
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_dtd_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 47
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_eurosat_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 10
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_kitti_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 4
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_oxford_flowers102_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 102
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_oxford_iiit_pet_config():
    config = ml_collections.ConfigDict()

    config.Num_Classes = 37
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_patch_camelyon_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 2
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_resisc45_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 45
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_smallnorb_azi_config():
    config = ml_collections.ConfigDict()

    config.Num_Classes = 18
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_smallnorb_ele_config():
    config = ml_collections.ConfigDict()

    config.Num_Classes = 9
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_sun397_config():
    config = ml_collections.ConfigDict()

    config.Num_Classes = 397
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config

def get_svhn_config():
    config = ml_collections.ConfigDict()
    config.Num_Classes = 10
    config.CropSize = 224

    config.Num_workers = 4
    config.pin_memory = True

    config.NUM_GPUS = 1
    return config


DATA_CONFIGS ={
    "caltech101": get_caltech101_config(),
    "cifar": get_cifar_config(),
    "clevr_count": get_clevr_count_config(),
    "clevr_dist": get_clevr_dist_config(),
    "diabetic_retinopathy": get_diabetic_retinopathy_config(),
    "dmlab": get_dmlab_config(),
    "dsprites_loc": get_dsprites_loc_config(),
    "dsprites_ori": get_dsprites_ori_config(),
    "dtd": get_dtd_config(),
    "eurosat": get_eurosat_config(),
    "kitti": get_kitti_config(),
    "oxford_flowers102": get_oxford_flowers102_config(),
    "oxford_iiit_pet": get_oxford_iiit_pet_config(),
    "patch_camelyon": get_patch_camelyon_config(),
    "resisc45": get_resisc45_config(),
    "smallnorb_azi": get_smallnorb_azi_config(),
    "smallnorb_ele": get_smallnorb_ele_config(),
    "sun397": get_sun397_config(),
    "svhn": get_svhn_config()
}
Data_process/VTAB_config.py
DavidYanAnDe/ARC
# coding=utf-8
# Copyright 2022 Gen Luo. All rights reserved.
#
# 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.utils.data as data
from PIL import Image
import os
import os.path
from torchvision import transforms
import torch
from timm.data import create_transform

def default_loader(path):
    return Image.open(path).convert('RGB')


def default_flist_reader(flist):
    """
    flist format: impath label\nimpath label\n ...(same to caffe's filelist)
    """
    imlist = []
    with open(flist, 'r') as rf:
        for line in rf.readlines():
            impath, imlabel = line.strip().split()
            imlist.append((impath, int(imlabel)))

    return imlist


class ImageFilelist(data.Dataset):
    def __init__(self, root, flist, transform=None, target_transform=None,
                 flist_reader=default_flist_reader, loader=default_loader):
        self.root = root
        self.imlist = flist_reader(flist)
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

    def __getitem__(self, index):
        impath, target = self.imlist[index]
        img = self.loader(os.path.join(self.root, impath))
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        return len(self.imlist)


def get_data(data_path, name, evaluate=True, batch_size=32,few_shot=False,shot=1,seed=1,mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    if few_shot:
        root_train = './data/few_shot/' + name + '/images/train/'
        root_val = './data/few_shot/' + name + '/images/test/'
        trainval_flist='./data/few_shot/' +name+'/train_meta.list.num_shot_%d.seed_%d'%(shot,seed)
        train_flist='./data/few_shot/' +name+'/train_meta.list.num_shot_%d.seed_%d'%(shot,seed)
        val_flist='./data/few_shot/' +name+'test_meta.list'
        test_flist='./data/few_shot/' +name+'test_meta.list'
        train_transform = create_transform(
            input_size=224,
            is_training=True,
            color_jitter=0.4,
            auto_augment='rand-m9-mstd0.5-inc1',
            interpolation='bicubic',
            re_prob=0.25,
            re_mode='pixel',
            re_count=1,
        )
    else:
        root=os.path.join(data_path, name)
        root_train = root
        root_val = root
        trainval_flist=root + "/train800val200.txt"
        train_flist=root +"/train800.txt"
        val_flist=root + "/val200.txt"
        test_flist=root + "/test.txt"
        train_transform = transforms.Compose([
                transforms.Resize((224, 224), interpolation=3),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std)])
    val_transform = transforms.Compose([
                transforms.Resize((224, 224), interpolation=3),
                transforms.ToTensor(),
                transforms.Normalize(mean=mean, std=std)])
    if evaluate:
        train_loader = torch.utils.data.DataLoader(
            ImageFilelist(root=root_train, flist=trainval_flist,
                transform=train_transform),
            batch_size=batch_size, shuffle=True, drop_last=False,
            num_workers=4, pin_memory=True)

        val_loader = torch.utils.data.DataLoader(
            ImageFilelist(root=root_val, flist=test_flist,
                transform=val_transform),
            batch_size=256, shuffle=False,
            num_workers=4, pin_memory=True)
    else:
        train_loader = torch.utils.data.DataLoader(
            ImageFilelist(root=root_train, flist=train_flist,
                transform=train_transform),
            batch_size=batch_size, shuffle=True, drop_last=False,
            num_workers=4, pin_memory=True)

        val_loader = torch.utils.data.DataLoader(
            ImageFilelist(root=root_val, flist=val_flist,
                transform=val_transform),
            batch_size=256, shuffle=False,
            num_workers=4, pin_memory=True)
    return train_loader, val_loader


Data_process/VTAB_loader.py
DavidYanAnDe/ARC

"""JSON dataset: support CUB, NABrids, Flower, Dogs and Cars"""

import logging

import os
import torch
import torch.utils.data
import torchvision as tv
import numpy as np
from collections import Counter

from Data_process.tranforms import get_transforms
from Data_process.json_process import read_json

logger = logging.getLogger(__name__)


class JSONDataset(torch.utils.data.Dataset):
    def __init__(self, cfg, split, data_path):
        assert split in {
            "train",
            "val",
            "test",
            "trainval"
        }, "Split '{}' not supported for {} dataset".format(
            split, cfg.Name)
        logger.info("Constructing {} dataset {}...".format(
            cfg.Name, split))

        self.cfg = cfg
        self._split = split
        self.name = cfg.Name
        self.data_dir = os.path.join(data_path, self.name)
        # self.data_percentage = cfg.DATA.PERCENTAGE
        self._construct_imdb(cfg)
        self.transform = get_transforms(split, cfg.CropSize)

    def get_anno(self):
        # /root/Data_process/CUB  /train.json
        anno_path = os.path.join(self.data_dir, "{}.json".format(self._split))
        # if "train" in self._split:
        #     if self.data_percentage < 1.0:
        #         anno_path = os.path.join(
        #             self.data_dir,
        #             "{}_{}.json".format(self._split, self.data_percentage)
        #         )
        assert os.path.exists(anno_path), "{} dir not found".format(anno_path)

        return read_json(anno_path)

    def get_imagedir(self):
        raise NotImplementedError()

    def _construct_imdb(self, cfg):
        """Constructs the imdb."""

        img_dir = self.get_imagedir()
        assert os.path.exists(img_dir), "{} dir not found".format(img_dir)

        anno = self.get_anno()
        # Map class ids to contiguous ids
        # {0,3,5,7,8} 120
        self._class_ids = sorted(list(set(anno.values())))

        self._class_id_cont_id = {v: i for i, v in enumerate(self._class_ids)}
        # {0:0, 3:1, 5:2 7:3, 8:4}


        # Construct the image db
        self._imdb = []
        for img_name, cls_id in anno.items():
            # 对标签进行了重新排序的处理
            cont_id = self._class_id_cont_id[cls_id]
            im_path = os.path.join(img_dir, img_name)
            self._imdb.append({"im_path": im_path, "class": cont_id})

       # print(len(self._imdb[]))
        print(len(self._class_ids))
        logger.info("Number of images: {}".format(len(self._imdb)))
        logger.info("Number of classes: {}".format(len(self._class_ids)))

    def get_info(self):
        num_imgs = len(self._imdb)
        return num_imgs, self.get_class_num()

    def get_class_num(self):
        return self.cfg.Num_Classes
        # return len(self._class_ids)

    def get_class_weights(self, weight_type):
        """get a list of class weight, return a list float"""
        if "train" not in self._split:
            raise ValueError(
                "only getting training class distribution, " + \
                "got split {} instead".format(self._split)
            )

        cls_num = self.get_class_num()
        if weight_type == "none":
            return [1.0] * cls_num

        id2counts = Counter(self._class_ids)
        assert len(id2counts) == cls_num
        num_per_cls = np.array([id2counts[i] for i in self._class_ids])

        if weight_type == 'inv':
            mu = -1.0
        elif weight_type == 'inv_sqrt':
            mu = -0.5
        weight_list = num_per_cls ** mu
        weight_list = np.divide(
            weight_list, np.linalg.norm(weight_list, 1)) * cls_num
        return weight_list.tolist()

    def __getitem__(self, index):
        # Load the image
        im = tv.datasets.folder.default_loader(self._imdb[index]["im_path"])
        label = self._imdb[index]["class"]
        im = self.transform(im)
        if self._split == "train":
            index = index
        else:
            index = f"{self._split}{index}"
        sample = {
            "image": im,
            "label": label,
            # "id": index
        }
        return im, label

    def __len__(self):
        return len(self._imdb)


class CUB200Dataset(JSONDataset):
    """CUB_200 dataset."""

    def __init__(self, cfg, split, data_path):
        super(CUB200Dataset, self).__init__(cfg, split, data_path)

    def get_imagedir(self):
        return os.path.join(self.data_dir, "images")


class CarsDataset(JSONDataset):
    """stanford-cars dataset."""

    def __init__(self, cfg, split, data_path):
        super(CarsDataset, self).__init__(cfg, split, data_path)

    def get_imagedir(self):
        return self.data_dir


class DogsDataset(JSONDataset):
    """stanford-dogs dataset."""

    def __init__(self, cfg, split, data_path):
        super(DogsDataset, self).__init__(cfg, split, data_path)

    def get_imagedir(self):
        return self.data_dir


class FlowersDataset(JSONDataset):
    """flowers dataset."""

    def __init__(self, cfg, split, data_path):
        super(FlowersDataset, self).__init__(cfg, split, data_path)

    def get_imagedir(self):
        return self.data_dir


class NabirdsDataset(JSONDataset):
    """Nabirds dataset."""

    def __init__(self, cfg, split, data_path):
        super(NabirdsDataset, self).__init__(cfg, split, data_path)

    def get_imagedir(self):
        return os.path.join(self.data_dir, "images")
Data_process/json_dataset.py
DavidYanAnDe/ARC
#!/usr/bin/env python3
"""
a bunch of helper functions for read and write data
"""
import os
import json
import numpy as np
import time
import pandas as pd

from typing import List, Union
from PIL import Image, ImageFile
Image.MAX_IMAGE_PIXELS = None


def save_or_append_df(out_path, df):
    if os.path.exists(out_path):
        previous_df = pd.read_pickle(out_path)
        df = pd.concat([previous_df, df], ignore_index=True)
    df.to_pickle(out_path)
    print(f"Saved output at {out_path}")


class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, bytes):
            return str(obj, encoding='utf-8')
        elif isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            # return super(MyEncoder, self).default(obj)

            raise TypeError(
                "Unserializable object {} of type {}".format(obj, type(obj))
            )


def write_json(data: Union[list, dict], outfile: str) -> None:
    json_dir, _ = os.path.split(outfile)
    if json_dir and not os.path.exists(json_dir):
        os.makedirs(json_dir)

    with open(outfile, 'w') as f:
        json.dump(data, f, cls=JSONEncoder, ensure_ascii=False, indent=2)


def read_json(filename: str) -> Union[list, dict]:
    """read json files"""
    with open(filename, "rb") as fin:
        data = json.load(fin, encoding="utf-8")
    return data


def pil_loader(path: str) -> Image.Image:
    """load an image from path, and suppress warning"""
    # to avoid crashing for truncated (corrupted images)
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    # open path as file to avoid ResourceWarning
    # (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')
Data_process/json_process.py
DavidYanAnDe/ARC
'''

Image preprocessing steps

'''

import torchvision as tv

def get_transforms(split, size):
    # 使用了标准的mean和std
    normalize = tv.transforms.Normalize(
        mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]
    )

    if size == 448:
        resize_dim = 512
        crop_dim = 448
    elif size == 224:
        resize_dim = 256
        crop_dim = 224
    elif size == 384:
        resize_dim = 438
        crop_dim = 384
    if split == "train":
        transform = tv.transforms.Compose(
            [
                tv.transforms.Resize(resize_dim),
                tv.transforms.RandomCrop(crop_dim),
                tv.transforms.RandomHorizontalFlip(0.5),
                # tv.transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
                # tv.transforms.RandomHorizontalFlip(),
                tv.transforms.ToTensor(),
                normalize,
            ]
        )
    else:
        transform = tv.transforms.Compose(
            [
                tv.transforms.Resize(resize_dim),
                tv.transforms.CenterCrop(crop_dim),
                tv.transforms.ToTensor(),
                normalize,
            ]
        )


    return transform


Data_process/tranforms.py
DavidYanAnDe/ARC
# coding=utf-8
from __future__ import absolute_import, division, print_function

import logging
import argparse
import os
import random
import numpy as np
import torch

from Model.Model_Config import CONFIGS
from Model.ARC_ViT import ARCVisionTransformer
from Utils.tools import count_parameters, AverageMeter

from Data_process.FGVC_config import DATA_CONFIGS
from Data_process.FGVC_loader import construct_test_loader, construct_train_loader

from tqdm import tqdm
from Utils.scheduler import WarmupLinearSchedule, WarmupCosineSchedule
from Utils.Frozen_weight import ViT_ARC_Frozen

logger = logging.getLogger(__name__)


def save_model(args, model):
    model_to_save = model.Module if hasattr(model, "module") else model
    model_checkpoint = os.path.join(args.output_dir, "%s_checkpoint.bin" % args.name)
    torch.save(model_to_save.state_dict(), model_checkpoint)
    logger.info("Saved model checkpoint to [DIR: %s]", args.output_dir)


def simple_accuracy(preds, labels):
    return (preds == labels).mean()


def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def setup(args):
    # get configure information
    config = CONFIGS[args.model_type]
    config.transformer["mlp_drop_rate"] = args.vit_drop
    config.adapter_dropout = args.adapt_drop

    # 初始化transformer模型
    num_classes = DATA_CONFIGS[args.dataset].Num_Classes
    model = ARCVisionTransformer(config, args.img_size, num_classes=num_classes, zero_head=True,
                                 tuning_mode=args.tuning_mode)
    model.load_from(np.load(args.pretrained_dir))

    model.to(args.device)

    # frozen weights
    ViT_ARC_Frozen(model)

    num_params = count_parameters(model)
    logger.info("{}".format(config))
    logger.info("Training HypeParameters %s", args)
    logger.info("Total Parameter: \t%2.2fM" % num_params)

    for name, para in model.named_parameters():
        if para.requires_grad == True:
            print(name)

    return model, num_params


def valid(args, model, test_loader):
    eval_losses = AverageMeter()

    logger.info("***** Running Validation *****")
    logger.info("  Num steps = %d", len(test_loader))
    logger.info("  Batch size = %d", args.eval_batch_size)

    model.eval()
    all_preds, all_label = [], []
    epoch_iterator = tqdm(test_loader,
                          desc="Testing (X / X steps) (loss = X.X)",
                          bar_format='{l_bar}-{r_bar}',
                          dynamic_ncols=True)
    loss_fct = torch.nn.CrossEntropyLoss()
    for step, batch in enumerate(epoch_iterator):
        batch = tuple(t.to(args.device) for t in batch)
        x, label = batch
        with torch.no_grad():
            logits = model(x)
            eval_loss = loss_fct(logits, label)
            eval_losses.update(eval_loss.item())
            prediction = torch.argmax(logits, dim=-1)

        if len(all_preds) == 0:
            all_preds.append(prediction.detach().cpu().numpy())
            all_label.append(label.detach().cpu().numpy())
        else:
            all_preds[0] = np.append(
                all_preds[0], prediction.detach().cpu().numpy(), axis=0)
            all_label[0] = np.append(
                all_label[0], label.detach().cpu().numpy(), axis=0)

        epoch_iterator.set_description("Validating... (loss=%2.5f)" % eval_losses.val)

    all_preds = all_preds[0]
    all_label = all_label[0]
    accuracy = simple_accuracy(all_preds, all_label)

    logger.info("\n")
    logger.info("Validation Results")
    logger.info("Valid Loss: %2.5f" % eval_losses.avg)
    logger.info("Valid Accuracy: %2.5f" % accuracy)

    return accuracy


def train(args, model):
    # output dir
    os.makedirs(args.output_dir, exist_ok=True)

    data_cfg = DATA_CONFIGS[args.dataset]
    train_loader = construct_train_loader(data_cfg, args.train_batch_size, data_path=args.data_path)
    test_loader = construct_test_loader(data_cfg, args.eval_batch_size, data_path=args.data_path)

    epoch_steps = len(train_loader)
    total_steps = epoch_steps * args.num_epochs
    warm_steps = epoch_steps * args.warmup_epochs
    evl_steps = args.eval_every * epoch_steps

    optimizer = torch.optim.AdamW(model.parameters(),
                                  lr=args.learning_rate,
                                  weight_decay=args.weight_decay)

    if args.decay_type == "cosine":
        schedule = WarmupCosineSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)
    else:
        schedule = WarmupLinearSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)

    logger.info("***** Running training *****")
    logger.info("  Total optimization steps = %d", total_steps)
    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size)
    logger.info("  Total train batch size = %d", args.train_batch_size * args.gradient_accumulation_steps)
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)

    # 开始训练，首先将梯度给初始化为0
    model.zero_grad()
    set_seed(args)  # 需要进行设定唯一的seed，这里相当于重复现
    # 用于记录历史的Loss值
    losses = AverageMeter()
    # 全局的步骤和最好的准确率
    global_step, best_acc = 0, 0

    # 开始训练的过程
    while True:
        model.train()
        # 用于显示训练进度的进度条类，可以需要的时候再单独学习
        epoch_iterator = tqdm(train_loader,
                              desc="Training (X / X steps) (loss = X.X)",
                              bar_format='{l_bar}-{r_bar}',
                              dynamic_ncols=True)
        for step, batch in enumerate(epoch_iterator):
            batch = tuple(t.to(args.device) for t in batch)
            x, label = batch

            loss = model(x, label)

            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            loss.backward()

            if (step + 1) % args.gradient_accumulation_steps == 0:
                losses.update(loss.item())
                # 用于添加norm的阶段信息，类似于加了一个惩罚项，就是所谓的L1或者L2约束
                torch.nn.utils.clip_grad_norm(model.parameters(), args.max_grad_norm)
                optimizer.step()
                schedule.step()
                optimizer.zero_grad()
                global_step += 1

                epoch_iterator.set_description(
                    "Training (%d / %d Steps) (loss=%2.5f)" % (global_step, total_steps, losses.val))

                if (global_step + 1) % evl_steps == 0:
                    accuracy = valid(args, model, test_loader)
                    if best_acc < accuracy:
                        # save_model(args, model)
                        best_acc = accuracy
                    model.train()

            if (global_step + 1) % total_steps == 0:
                break

        losses.reset()
        if (global_step + 1) % total_steps == 0:
            break

    logger.info("Best Accuracy: \t%f" % best_acc)
    logger.info("End Training!")
    return best_acc


def main():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("--data_path", default="/root/autodl-tmp/Data",
                        help="where is the data store, should be changed according to local file")

    parser.add_argument("--name", default="ViT-B_16_cifar_att_mlp_independent_50dim_bias0_repadapterStructure_error",
                        help="Name of this run. Used for monitoring.")

    parser.add_argument("--dataset", choices=["StanfordDogs", "OxfordFlowers", "CUB_200_2011", "StanfordCars",
                                              "NABirds"], default="StanfordDogs",
                        help="Which downstream task.")

    parser.add_argument("--tuning_mode", choices=["ARC_att", "ARC"],
                        default="ARC_att", help="Which downstream task.")

    parser.add_argument("--model_type", choices=["ViT-B_16", "ViT-L_16", "ViT-H_14"],
                        default="ViT-B_16",
                        help="Which variant to use.")
    parser.add_argument("--pretrained_dir", type=str, default="ViT-B_16.npz",
                        help="Where to search for pretrained ViT models.")
    parser.add_argument("--output_dir", default="output", type=str,
                        help="The output directory where checkpoints will be written.")

    parser.add_argument("--img_size", default=224, type=int,
                        help="Resolution size")
    parser.add_argument("--train_batch_size", default=32, type=int,
                        help="Total batch size for training.")
    parser.add_argument("--eval_batch_size", default=256, type=int,
                        help="Total batch size for eval.")
    parser.add_argument("--eval_every", default=1, type=int,
                        help="Run prediction on validation set every so many steps."
                             "Will always run one evaluation at the end of training.")

    parser.add_argument("--learning_rate", default=0.005, type=float,
                        help="The initial learning rate for SGD.")
    parser.add_argument("--weight_decay", default=0.01, type=float,
                        help="Weight decay if we apply some.")
    parser.add_argument("--num_epochs", default=100, type=int,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--decay_type", choices=["cosine", "linear"], default="cosine",
                        help="How to decay the learning rate.")
    parser.add_argument("--warmup_epochs", default=10, type=int,
                        help="Step of training to perform learning rate warmup for.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")

    parser.add_argument("--vit_drop", default=0, type=float,
                        help="Max gradient norm.")
    parser.add_argument("--adapt_drop", default=0.1, type=float,
                        help="Max gradient norm.")

    args = parser.parse_args()

    # Setup CUDA, GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = torch.cuda.device_count()
    args.device = device

    # Setup logging
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
    logger.warning("Process device: %s, n_gpu: %s" % (args.device, args.n_gpu))

    # Set seed
    set_seed(args)

    # Model & Tokenizer Setup
    model, num_params = setup(args)

    # Training
    best_acc = train(args, model)

    return best_acc, num_params

if __name__ == "__main__":
    main()
FGVC_ARC_train.py
DavidYanAnDe/ARC
import copy
import logging
import math

import torch.nn as nn
from torch.nn.functional import gelu, relu
from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm
import torch
from torch.nn.modules.utils import _pair
from Utils.tools import np2th
from os.path import join as pjoin

logger = logging.getLogger(__name__)

# activation function define
ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu}

ATTENTION_Q = "MultiHeadDotProductAttention_1/query"
ATTENTION_K = "MultiHeadDotProductAttention_1/key"
ATTENTION_V = "MultiHeadDotProductAttention_1/value"
ATTENTION_OUT = "MultiHeadDotProductAttention_1/out"
FNN_0 = "MlpBlock_3/Dense_0"
FNN_1 = "MlpBlock_3/Dense_1"
ATTENTION_NORM = "LayerNorm_0"
MLP_NORM = "LayerNorm_2"

# adapter module
class ARC_adapter(nn.Module):
    def __init__(self, adapter_dim, hidden_dim, dropout=0.0, position='att'):
        super(ARC_adapter, self).__init__()
        self.adapter_rescale = nn.Parameter(torch.empty(1, adapter_dim))
        self.adapter_bias = nn.Parameter(torch.empty(hidden_dim))
        self.dropout = Dropout(dropout)

        if position == 'att':
            nn.init.zeros_(self.adapter_rescale)
        else:
            nn.init.zeros_(self.adapter_rescale)
        nn.init.zeros_(self.adapter_bias)

    def forward(self, x, down_projection, up_projection):
        adapter_output = torch.matmul(x, down_projection * self.adapter_rescale)
        adapter_output = self.dropout(adapter_output)
        adapter_output = torch.matmul(adapter_output, up_projection) + self.adapter_bias
        output = adapter_output + x

        return output

# embedding layer in ViT
class Embedding(nn.Module):
    def __init__(self, config, img_size, in_channel=3):
        super(Embedding, self).__init__()

        # image_size(n,n)    patch_size(t, t)
        image_size = _pair(img_size)
        patch_size = config.patches['size']
        num_patches = (image_size[0] // patch_size[0]) * (image_size[1] // patch_size[1])

        # 将[H,W,3]转换为[224/16, 224/16, hidden_size]
        self.projective = Conv2d(in_channels=in_channel,
                                 out_channels=config.hidden_size,
                                 kernel_size=patch_size,
                                 stride=patch_size)

        # 可学习的cls, cls维度[1, 1, hidden_dim]
        self.cls = nn.Parameter(torch.zeros(1, 1, config.hidden_size))

        # 可学习的position_embedding, [1, patch_num+1, hidden_dim]
        self.position_embedding = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))

        # 输入encode前有一个dropout
        self.dropout = Dropout(config.transformer["mlp_drop_rate"])

    def forward(self, input_image):
        # [B, hidden_size, patch_size, patch_size]
        projective_embedding = self.projective(input_image)
        # flatten to [B, num_patches, hidden_size]
        projective_embedding = projective_embedding.flatten(2).transpose(-1, -2)

        # add num_patches to class_token. [B, num_patches, hidden_size]and[1, 1, hidden_size]
        batch_size = input_image.shape[0]
        cls = self.cls.expand(batch_size, -1, -1)
        # add cls token, feature dimensions become [B, num_patches+1, hidden_size]
        embedding_features = torch.concat([cls, projective_embedding], dim=1)

        output_features = embedding_features + self.position_embedding
        output_features = self.dropout(output_features)
        return output_features


# Attention block in encoder layer
class Attention(nn.Module):
    def __init__(self, config):
        super(Attention, self).__init__()

        hidden_dim = config.hidden_size

        self.head_num = config.transformer["num_heads"]

        self.head_hidden_dim = hidden_dim // self.head_num
        self.all_head_size = self.head_hidden_dim * self.head_num

        # q,k,v matrices
        self.query = Linear(hidden_dim, self.all_head_size)
        self.key = Linear(hidden_dim, self.all_head_size)
        self.value = Linear(hidden_dim, self.all_head_size)

        # softmax for attention matrix，
        self.softmax = Softmax(dim=-1)
        # projection matrix after attention
        self.projective = Linear(hidden_dim, hidden_dim)

        self.attn_dropout = Dropout(config.transformer["attention_drop_rate"])
        self.proj_dropout = Dropout(config.transformer["attention_drop_rate"])

    def transpose_head(self, mix_metric):
        # new_x_shape size [batch_size,n,H,n/H]
        new_x_shape = mix_metric.size()[:-1] + (self.head_num, self.head_hidden_dim)
        # '*' change shape only, donot change memory
        mix_metric = mix_metric.view(*new_x_shape)

        return mix_metric.permute(0, 2, 1, 3)

    def forward(self, hidden_states):
        mix_query_layer = self.query(hidden_states)
        mix_key_layer = self.key(hidden_states)
        mix_value_layer = self.value(hidden_states)

        # split q,k,v change dimensions from [Batch_size, num_patches, hidden_dim] to [Batch_size, head_num, num_patches, head_dim]
        query_layer = self.transpose_head(mix_query_layer)
        key_layer = self.transpose_head(mix_key_layer)
        value_layer = self.transpose_head(mix_value_layer)

        # calculate attention score
        attention_metric = torch.matmul(query_layer, key_layer.transpose(-2, -1))
        attention_metric = attention_metric / math.sqrt(self.head_hidden_dim)
        attention_probs = self.softmax(attention_metric)

        attention_probs = self.attn_dropout(attention_probs)

        attention_features = torch.matmul(attention_probs, value_layer)

        # reconstruct attention_features进，recover the [batch_size,n,d] format
        attention_features = attention_features.permute(0, 2, 1, 3).contiguous()
        new_output_features_shape = attention_features.size()[:-2] + (self.all_head_size,)
        attention_feature = attention_features.view(*new_output_features_shape)

        # add one projection before output
        output_features = self.projective(attention_feature)
        output_features = self.proj_dropout(output_features)

        return output_features


# MLP block of encoder layer
class MLP(nn.Module):
    def __init__(self, config):
        super(MLP, self).__init__()
        # two linear
        self.fc1 = Linear(config.hidden_size, config.transformer["mlp_dim"])
        # activation function
        self.activation = ACT2FN["gelu"]
        self.fc2 = Linear(config.transformer["mlp_dim"], config.hidden_size)
        self.dropout = Dropout(config.transformer["mlp_drop_rate"])

    def forward(self, input_features):
        output_features = self.fc1(input_features)
        output_features = self.activation(output_features)
        output_features = self.dropout(output_features)
        output_features = self.fc2(output_features)
        output_features = self.dropout(output_features)

        return output_features

class Encoder(nn.Module):
    def __init__(self, config, tuning_mode):
        super(Encoder, self).__init__()

        self.hidden_size = config.hidden_size
        self.attention = Attention(config)
        self.mlp = MLP(config)
        # layer norm在输入attention和mlp前
        self.att_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)
        self.fnn_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)

        self.att_adapter = None
        self.mlp_adapter = None
        if tuning_mode == 'ARC_att':
            self.att_adapter = ARC_adapter(adapter_dim=config.adapter_dim, hidden_dim=self.hidden_size,
                                           dropout=config.adapter_dropout, position='att')
        elif tuning_mode == 'ARC':
            self.att_adapter = ARC_adapter(adapter_dim=config.adapter_dim, hidden_dim=self.hidden_size,
                                           dropout=config.adapter_dropout, position='att')
            self.mlp_adapter = ARC_adapter(adapter_dim=config.adapter_dim, hidden_dim=self.hidden_size,
                                           dropout=config.adapter_dropout, position='mlp')

    def forward(self, hidden_states, att_projection=None, mlp_projection=None):
        # 残差网络的特征
        res_att_feature = hidden_states
        # 输入前先进行layer norm
        hidden_states = self.att_layer_norm(hidden_states)

        if self.att_adapter is not None and att_projection is not None:
            hidden_states = self.att_adapter(hidden_states, att_projection, att_projection.t())

        attention_output = self.attention(hidden_states)
        attention_output = attention_output + res_att_feature

        res_fnn_feature = attention_output
        mlp_output = self.fnn_layer_norm(attention_output)

        if self.mlp_adapter is not None and mlp_projection is not None:
            mlp_output = self.mlp_adapter(mlp_output, mlp_projection, mlp_projection.t())
        
        output = self.mlp(mlp_output)

        output = output + res_fnn_feature
        return output

    def load_from(self, weights, n_block):
        ROOT = f"Transformer/encoderblock_{n_block}"
        with torch.no_grad():
            query_weight = np2th(weights[pjoin(ROOT, ATTENTION_Q, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()
            key_weight = np2th(weights[pjoin(ROOT, ATTENTION_K, "kernel")]).view(self.hidden_size, self.hidden_size).t()
            value_weight = np2th(weights[pjoin(ROOT, ATTENTION_V, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()
            out_weight = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()

            query_bias = np2th(weights[pjoin(ROOT, ATTENTION_Q, "bias")]).view(-1)
            key_bias = np2th(weights[pjoin(ROOT, ATTENTION_K, "bias")]).view(-1)
            value_bias = np2th(weights[pjoin(ROOT, ATTENTION_V, "bias")]).view(-1)
            out_bias = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "bias")]).view(-1)

            self.attention.query.weight.copy_(query_weight)
            self.attention.key.weight.copy_(key_weight)
            self.attention.value.weight.copy_(value_weight)
            self.attention.projective.weight.copy_(out_weight)
            self.attention.query.bias.copy_(query_bias)
            self.attention.key.bias.copy_(key_bias)
            self.attention.value.bias.copy_(value_bias)
            self.attention.projective.bias.copy_(out_bias)

            mlp_weight_0 = np2th(weights[pjoin(ROOT, FNN_0, "kernel")]).t()
            mlp_weight_1 = np2th(weights[pjoin(ROOT, FNN_1, "kernel")]).t()
            mlp_bias_0 = np2th(weights[pjoin(ROOT, FNN_0, "bias")]).t()
            mlp_bias_1 = np2th(weights[pjoin(ROOT, FNN_1, "bias")]).t()

            self.mlp.fc1.weight.copy_(mlp_weight_0)
            self.mlp.fc2.weight.copy_(mlp_weight_1)
            self.mlp.fc1.bias.copy_(mlp_bias_0)
            self.mlp.fc2.bias.copy_(mlp_bias_1)

            self.att_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "scale")]))
            self.att_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "bias")]))
            self.fnn_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "scale")]))
            self.fnn_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "bias")]))


class Transformer(nn.Module):
    def __init__(self, config, tuning_mode):
        super(Transformer, self).__init__()
        # encoder层的数量
        self.layer_num = config.transformer["num_layers"]
        # encoder按照存储每一层
        self.layer = nn.ModuleList()
        # 添加每一层
        for _ in range(self.layer_num):
            encoder = Encoder(config, tuning_mode)
            # 这里需要深度拷贝
            self.layer.append(copy.deepcopy(encoder))

        # 最后有一个layer_norm输出
        self.layer_norm = LayerNorm(config.hidden_size, eps=1e-6)

        self.att_down_projection=None
        self.mlp_down_projection=None
        if tuning_mode == 'ARC_att':
            self.att_down_projection = nn.Parameter(torch.empty(config.hidden_size, config.adapter_dim))
            nn.init.xavier_uniform_(self.att_down_projection)
        elif tuning_mode == 'ARC':
            self.att_down_projection = nn.Parameter(torch.empty(config.hidden_size, config.adapter_dim))
            self.mlp_down_projection = nn.Parameter(torch.empty(config.hidden_size, config.adapter_dim))
            nn.init.xavier_uniform_(self.mlp_down_projection)
            nn.init.xavier_uniform_(self.att_down_projection)

    def forward(self, input_feature):
        for layer_encoder in self.layer:
            input_feature = layer_encoder(input_feature, self.att_down_projection, self.mlp_down_projection)

        output = self.layer_norm(input_feature)
        return output

# head layer in ViT
class HeadMLP(nn.Module):
    def __init__(self, config, num_classes):
        super(HeadMLP, self).__init__()
        self.head = Linear(config.hidden_size, num_classes)

    def forward(self, input_features):
        cls_token = input_features[:, 0]

        output = self.head(cls_token)
        return output


class ARCVisionTransformer(nn.Module):
    def __init__(self, config, img_size=224, num_classes=21843, zero_head=False, tuning_mode='ARC'):
        super(ARCVisionTransformer, self).__init__()
        self.zero_head = zero_head
        self.num_classes = num_classes

        self.embedding = Embedding(config, img_size)
        self.transformer = Transformer(config, tuning_mode)

        self.head_mlp = HeadMLP(config, num_classes)

    def forward(self, image, labels=None):
        image_projective = self.embedding(image)
        output_representation = self.transformer(image_projective)
        # get prediction results from model
        prediction = self.head_mlp(output_representation)

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(prediction.view(-1, self.num_classes), labels.view(-1))
            return loss
        else:
            return prediction

    # load pre-trained parameters
    def load_from(self, weights):
        with torch.no_grad():
            if self.zero_head:
                nn.init.zeros_(self.head_mlp.head.weight)
                nn.init.zeros_(self.head_mlp.head.bias)
            else:
                self.head_mlp.head.weight.copy_(np2th(weights["head/kernel"]).t())
                self.head_mlp.head.weight.copy_(np2th(weights["head/bias"]).t())

            self.embedding.projective.weight.copy_(np2th(weights["embedding/kernel"], conv=True))
            self.embedding.projective.bias.copy_(np2th(weights["embedding/bias"]))
            self.embedding.cls.copy_(np2th(weights["cls"]))

            posemb = np2th(weights["Transformer/posembed_input/pos_embedding"])
            posemb_new = self.embedding.position_embedding
            if posemb.size() == posemb_new.size():
                self.embedding.position_embedding.copy_(posemb)

            self.transformer.layer_norm.weight.copy_(np2th(weights["Transformer/encoder_norm/scale"]))
            self.transformer.layer_norm.bias.copy_(np2th(weights["Transformer/encoder_norm/bias"]))

            for bname, encoder in self.transformer.named_children():
                for uname, unit in encoder.named_children():
                    unit.load_from(weights, n_block=uname)
Model/ARC_ViT.py
DavidYanAnDe/ARC
import copy
import logging
import math

import torch.nn as nn
from torch.nn.functional import gelu, relu
from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm
import torch
from torch.nn.modules.utils import _pair
from Utils.tools import np2th
from os.path import join as pjoin

logger = logging.getLogger(__name__)

# 定义MLP需要用到的激活函数
ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu}

ATTENTION_Q = "MultiHeadDotProductAttention_1/query"
ATTENTION_K = "MultiHeadDotProductAttention_1/key"
ATTENTION_V = "MultiHeadDotProductAttention_1/value"
ATTENTION_OUT = "MultiHeadDotProductAttention_1/out"
FNN_0 = "MlpBlock_3/Dense_0"
FNN_1 = "MlpBlock_3/Dense_1"
ATTENTION_NORM = "LayerNorm_0"
MLP_NORM = "LayerNorm_2"


# 图像输入的embedding层
class Embedding(nn.Module):
    def __init__(self, config, img_size, in_channel=3):
        super(Embedding, self).__init__()

        # image_size(n,n)    patch_size(t, t)
        image_size = _pair(img_size)
        patch_size = config.patches['size']
        num_patches = (image_size[0] // patch_size[0]) * (image_size[1] // patch_size[1])

        # 将[H,W,3]转换为[224/16, 224/16, hidden_size]
        self.projective = Conv2d(in_channels=in_channel,
                                 out_channels=config.hidden_size,
                                 kernel_size=patch_size,
                                 stride=patch_size)

        # 可学习的cls, cls维度[1, 1, hidden_dim]
        self.cls = nn.Parameter(torch.zeros(1, 1, config.hidden_size))

        # 可学习的position_embedding, [1, patch_num+1, hidden_dim]
        self.position_embedding = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))

        # 输入encode前有一个dropout
        self.dropout = Dropout(config.transformer["mlp_drop_rate"])

    def forward(self, input_image):
        # [B, hidden_size, patch_size, patch_size]
        projective_embedding = self.projective(input_image)
        # 展开，并且调换维度，变成[B, num_patches, hidden_size]
        projective_embedding = projective_embedding.flatten(2).transpose(-1, -2)

        # num_patches, 加入class_token. [B, num_patches, hidden_size]和[1, 1, hidden_size]
        # 首先将cls变为[B, 1, hidden_size]
        batch_size = input_image.shape[0]
        cls = self.cls.expand(batch_size, -1, -1)
        # 拼接cls和projective_embedding, 变为[B, num_patches+1, hidden_size]
        embedding_features = torch.concat([cls, projective_embedding], dim=1)

        output_features = embedding_features + self.position_embedding
        output_features = self.dropout(output_features)
        return output_features


# encoder中的自注意力层
class Attention(nn.Module):
    def __init__(self, config):
        super(Attention, self).__init__()
        hidden_dim = config.hidden_size

        self.head_num = config.transformer["num_heads"]

        self.head_hidden_dim = hidden_dim // self.head_num
        self.all_head_size = self.head_hidden_dim * self.head_num


        self.query = Linear(hidden_dim, self.all_head_size)
        self.key = Linear(hidden_dim, self.all_head_size)
        self.value = Linear(hidden_dim, self.all_head_size)


        self.softmax = Softmax(dim=-1)

        self.projective = Linear(hidden_dim, hidden_dim)

        self.attn_dropout = Dropout(config.transformer["attention_drop_rate"])
        self.proj_dropout = Dropout(config.transformer["attention_drop_rate"])

    def transpose_head(self, mix_metric):
        new_x_shape = mix_metric.size()[:-1] + (self.head_num, self.head_hidden_dim)

        mix_metric = mix_metric.view(*new_x_shape)

        return mix_metric.permute(0, 2, 1, 3)

    def forward(self, hidden_states):

        mix_query_layer = self.query(hidden_states)
        mix_key_layer = self.key(hidden_states)
        mix_value_layer = self.value(hidden_states)


        query_layer = self.transpose_head(mix_query_layer)
        key_layer = self.transpose_head(mix_key_layer)
        value_layer = self.transpose_head(mix_value_layer)


        attention_metric = torch.matmul(query_layer, key_layer.transpose(-2, -1))
        attention_metric = attention_metric / math.sqrt(self.head_hidden_dim)
        attention_probs = self.softmax(attention_metric)

        attention_probs = self.attn_dropout(attention_probs)
        attention_features = torch.matmul(attention_probs, value_layer)

        attention_features = attention_features.permute(0, 2, 1, 3).contiguous()
        new_output_features_shape = attention_features.size()[:-2] + (self.all_head_size,)
        attention_feature = attention_features.view(*new_output_features_shape)


        output_features = self.projective(attention_feature)
        output_features = self.proj_dropout(output_features)

        return output_features


class MLP(nn.Module):
    def __init__(self, config):
        super(MLP, self).__init__()

        self.fc1 = Linear(config.hidden_size, config.transformer["mlp_dim"])

        self.activation = ACT2FN["gelu"]
        self.fc2 = Linear(config.transformer["mlp_dim"], config.hidden_size)

        self.dropout = Dropout(config.transformer["mlp_drop_rate"])

    def forward(self, input_features):

        output_features = self.fc1(input_features)
        output_features = self.activation(output_features)
        output_features = self.dropout(output_features)
        output_features = self.fc2(output_features)
        output_features = self.dropout(output_features)

        return output_features


class Encoder(nn.Module):
    def __init__(self, config):
        super(Encoder, self).__init__()

        self.hidden_size = config.hidden_size
        self.attention = Attention(config)
        self.mlp = MLP(config)

        self.att_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)
        self.fnn_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)

        self.att_scale = nn.Parameter(torch.empty(1, config.adapter_dim))
        self.att_scale_bias = nn.Parameter(torch.empty(self.hidden_size))

        self.mlp_scale = nn.Parameter(torch.empty(1, config.adapter_dim))
        self.mlp_scale_bias = nn.Parameter(torch.empty(self.hidden_size))

        nn.init.zeros_(self.att_scale)
        nn.init.zeros_(self.att_scale_bias)
        nn.init.xavier_uniform_(self.mlp_scale)
        nn.init.zeros_(self.mlp_scale_bias)

        self.adapter_drop = nn.Dropout(config.adapter_dropout)

    def forward(self, hidden_states, left_u_att, right_v_att, left_u_mlp, right_v_mlp):
        res_att_feature = hidden_states
        hidden_states = self.att_layer_norm(hidden_states)
        
        attention_uv_output = torch.matmul(hidden_states, left_u_att * self.att_scale)
        attention_uv_output = self.adapter_drop(attention_uv_output)
        attention_uv_output = torch.matmul(attention_uv_output, right_v_att) + self.att_scale_bias
        attention_uv_output = hidden_states + attention_uv_output
        
        attention_output = self.attention(attention_uv_output)
        attention_output = attention_output + res_att_feature

        res_fnn_feature = attention_output
        fnn_output = self.fnn_layer_norm(attention_output)
        
        mlp_uv_output = torch.matmul(fnn_output, left_u_mlp * self.mlp_scale)
        mlp_uv_output = self.adapter_drop(mlp_uv_output)
        mlp_uv_output = torch.matmul(mlp_uv_output, right_v_mlp) + self.mlp_scale_bias
        mlp_uv_output = mlp_uv_output + fnn_output
        
        fnn_output = self.mlp(mlp_uv_output)
        fnn_output = fnn_output + res_fnn_feature
        return fnn_output

    def load_from(self, weights, n_block):
        ROOT = f"Transformer/encoderblock_{n_block}"
        with torch.no_grad():
            query_weight = np2th(weights[pjoin(ROOT, ATTENTION_Q, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()
            key_weight = np2th(weights[pjoin(ROOT, ATTENTION_K, "kernel")]).view(self.hidden_size, self.hidden_size).t()
            value_weight = np2th(weights[pjoin(ROOT, ATTENTION_V, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()
            out_weight = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "kernel")]).view(self.hidden_size,
                                                                                   self.hidden_size).t()

            query_bias = np2th(weights[pjoin(ROOT, ATTENTION_Q, "bias")]).view(-1)
            key_bias = np2th(weights[pjoin(ROOT, ATTENTION_K, "bias")]).view(-1)
            value_bias = np2th(weights[pjoin(ROOT, ATTENTION_V, "bias")]).view(-1)
            out_bias = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "bias")]).view(-1)

            self.attention.query.weight.copy_(query_weight)
            self.attention.key.weight.copy_(key_weight)
            self.attention.value.weight.copy_(value_weight)
            self.attention.projective.weight.copy_(out_weight)
            self.attention.query.bias.copy_(query_bias)
            self.attention.key.bias.copy_(key_bias)
            self.attention.value.bias.copy_(value_bias)
            self.attention.projective.bias.copy_(out_bias)

            mlp_weight_0 = np2th(weights[pjoin(ROOT, FNN_0, "kernel")]).t()
            mlp_weight_1 = np2th(weights[pjoin(ROOT, FNN_1, "kernel")]).t()
            mlp_bias_0 = np2th(weights[pjoin(ROOT, FNN_0, "bias")]).t()
            mlp_bias_1 = np2th(weights[pjoin(ROOT, FNN_1, "bias")]).t()

            self.mlp.fc1.weight.copy_(mlp_weight_0)
            self.mlp.fc2.weight.copy_(mlp_weight_1)
            self.mlp.fc1.bias.copy_(mlp_bias_0)
            self.mlp.fc2.bias.copy_(mlp_bias_1)

            self.att_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "scale")]))
            self.att_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "bias")]))
            self.fnn_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "scale")]))
            self.fnn_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "bias")]))



class Transformer(nn.Module):
    def __init__(self, config):
        super(Transformer, self).__init__()
        self.layer_num = config.transformer["num_layers"]
        self.layer = nn.ModuleList()
        for _ in range(self.layer_num):
            encoder = Encoder(config)
            self.layer.append(copy.deepcopy(encoder))

        self.layer_norm = LayerNorm(config.hidden_size, eps=1e-6)

        self.left_u_att = nn.Parameter(torch.empty(config.hidden_size, config.adapter_dim))
        self.left_u_mlp = nn.Parameter(torch.empty(config.hidden_size, config.adapter_dim))

        nn.init.xavier_uniform_(self.left_u_mlp)
        nn.init.xavier_uniform_(self.left_u_att)


    def forward(self, input_feature):
        for layer_encoder in self.layer:
            input_feature = layer_encoder(input_feature, self.left_u_att, self.left_u_att.t(), self.left_u_mlp,
                                          self.left_u_mlp.t())

        output = self.layer_norm(input_feature)
        return output

class HeadMLP(nn.Module):
    def __init__(self, config, num_classes):
        super(HeadMLP, self).__init__()
        self.head = Linear(config.hidden_size, num_classes)

    def forward(self, input_features):
        cls_token = input_features[:, 0]

        output = self.head(cls_token)
        return output


class ARCVisionTransformer(nn.Module):
    def __init__(self, config, img_size=224, num_classes=21843, zero_head=False):
        super(ARCVisionTransformer, self).__init__()
        self.zero_head = zero_head
        self.num_classes = num_classes

        self.embedding = Embedding(config, img_size)
        self.transformer = Transformer(config)

        self.head_mlp = HeadMLP(config, num_classes)

    def forward(self, image, labels=None):
        image_projective = self.embedding(image)
        output_representation = self.transformer(image_projective)
        prediction = self.head_mlp(output_representation)

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(prediction.view(-1, self.num_classes), labels.view(-1))
            return loss
        else:
            return prediction

    def load_from(self, weights):
        with torch.no_grad():
            if self.zero_head:
                nn.init.zeros_(self.head_mlp.head.weight)
                nn.init.zeros_(self.head_mlp.head.bias)
            else:
                self.head_mlp.head.weight.copy_(np2th(weights["head/kernel"]).t())
                self.head_mlp.head.weight.copy_(np2th(weights["head/bias"]).t())

            self.embedding.projective.weight.copy_(np2th(weights["embedding/kernel"], conv=True))
            self.embedding.projective.bias.copy_(np2th(weights["embedding/bias"]))
            self.embedding.cls.copy_(np2th(weights["cls"]))

            posemb = np2th(weights["Transformer/posembed_input/pos_embedding"])
            posemb_new = self.embedding.position_embedding
            if posemb.size() == posemb_new.size():
                self.embedding.position_embedding.copy_(posemb)

            self.transformer.layer_norm.weight.copy_(np2th(weights["Transformer/encoder_norm/scale"]))
            self.transformer.layer_norm.bias.copy_(np2th(weights["Transformer/encoder_norm/bias"]))

            for bname, encoder in self.transformer.named_children():
                for uname, unit in encoder.named_children():
                    unit.load_from(weights, n_block=uname)

Model/ARC_new.py
DavidYanAnDe/ARC
""" Swin Transformer
A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`
    - https://arxiv.org/pdf/2103.14030
Code/weights from https://github.com/microsoft/Swin-Transformer, original copyright/license info below
Modifications and additions for timm hacked together by / Copyright 2021, Ross Wightman
"""
# --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------
import logging
import math
from copy import deepcopy
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint

from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.fx_features import register_notrace_function
from timm.models.helpers import build_model_with_cfg, named_apply 
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.layers import _assert
from timm.models.registry import register_model

from timm.models.vision_transformer import checkpoint_filter_fn, get_init_weights_vit 

_logger = logging.getLogger(__name__)


def _cfg(url='', **kwargs):
    return {
        'url': url,
        'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': None,
        'crop_pct': .9, 'interpolation': 'bicubic', 'fixed_input_size': True,
        'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD,
        'first_conv': 'patch_embed.proj', 'classifier': 'head',
        **kwargs
    }


default_cfgs = {
    # patch models (my experiments)
    'swin_base_patch4_window12_384': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22kto1k.pth',
        input_size=(3, 384, 384), crop_pct=1.0),

    'swin_base_patch4_window7_224': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22kto1k.pth',
    ),

    'swin_large_patch4_window12_384': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22kto1k.pth',
        input_size=(3, 384, 384), crop_pct=1.0),

    'swin_large_patch4_window7_224': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22kto1k.pth',
    ),

    'swin_small_patch4_window7_224': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth',
    ),

    'swin_tiny_patch4_window7_224': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth',
    ),

    'swin_base_patch4_window12_384_in22k': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth',
        input_size=(3, 384, 384), crop_pct=1.0, num_classes=21841),

    'swin_base_patch4_window7_224_in22k': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth',
        num_classes=21841),

    'swin_large_patch4_window12_384_in22k': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth',
        input_size=(3, 384, 384), crop_pct=1.0, num_classes=21841),

    'swin_large_patch4_window7_224_in22k': _cfg(
        url='https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth',
        num_classes=21841),

}



class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0., tuning_mode=None):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)
        drop_probs = to_2tuple(drop)

        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
        self.act = act_layer()
        self.drop1 = nn.Dropout(drop_probs[0])
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
        self.drop2 = nn.Dropout(drop_probs[1])

            
        self.tuning_mode = tuning_mode

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop1(x)
        x = self.fc2(x)
        
        x = self.drop2(x)
        
        return x

        
def window_partition(x, window_size: int):
    """
    Args:
        x: (B, H, W, C)
        window_size (int): window size
    Returns:
        windows: (num_windows*B, window_size, window_size, C)
    """
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows


@register_notrace_function  # reason: int argument is a Proxy
def window_reverse(windows, window_size: int, H: int, W: int):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size (int): Window size
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x


class WindowAttention(nn.Module):
    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.
    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        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, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0., tuning_mode=None):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH

        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_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.window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * self.window_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)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

        self.tuning_mode = tuning_mode

    def forward(self, x, mask: Optional[torch.Tensor] = None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, 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.unbind(0)  # make torchscript happy (cannot use tensor as tuple)

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_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


class SwinTransformerBlock(nn.Module):
    r""" Swin Transformer Block.
    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resulotion.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        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
        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, attn_u, mlp_u, svd_dim = 50, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0., adp_drop = 0.5,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, tuning_mode='ARC'):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        
        self.attn_u = attn_u
        self.mlp_u = mlp_u
        
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, qkv_bias=qkv_bias,
            attn_drop=attn_drop, proj_drop=drop, tuning_mode=tuning_mode)

        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, tuning_mode=tuning_mode)

        if self.shift_size > 0:
            # calculate attention mask for SW-MSA
            H, W = self.input_resolution
            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
            h_slices = (slice(0, -self.window_size),
                        slice(-self.window_size, -self.shift_size),
                        slice(-self.shift_size, None))
            w_slices = (slice(0, -self.window_size),
                        slice(-self.window_size, -self.shift_size),
                        slice(-self.shift_size, None))
            cnt = 0
            for h in h_slices:
                for w in w_slices:
                    img_mask[:, h, w, :] = cnt
                    cnt += 1

            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

        self.tuning_mode = tuning_mode

        if('ARC' == tuning_mode):
            
            self.attn_lambda = nn.Parameter(torch.empty(1, svd_dim))
            self.attn_lambda_bias = nn.Parameter(torch.empty(self.dim))

            nn.init.zeros_(self.attn_lambda)
            nn.init.zeros_(self.attn_lambda_bias)

            self.mlp_lambda = nn.Parameter(torch.empty(1, svd_dim))
            self.mlp_lambda_bias = nn.Parameter(torch.empty(self.dim))

            nn.init.xavier_uniform_(self.mlp_lambda)
            nn.init.zeros_(self.mlp_lambda_bias)

            self.attn_drop = nn.Dropout(adp_drop)
            self.mlp_drop = nn.Dropout(adp_drop)
            
    def ARC_adapter(self, x, u, adapter_drop, lambda_, bias):
        hidden_states = x
        
        attention_uv_output = torch.matmul(hidden_states, u * lambda_)
        attention_uv_output = adapter_drop(attention_uv_output)
        attention_uv_output = torch.matmul(attention_uv_output, u.t()) + bias
        attention_uv_output = hidden_states + attention_uv_output
        
        return attention_uv_output
        

    def forward(self, x):
        H, W = self.input_resolution
        B, L, C = x.shape

        shortcut = x
        
        x = self.norm1(x)
        if('ARC' == self.tuning_mode):
            
            x = self.ARC_adapter(x, self.attn_u,  self.attn_drop, self.attn_lambda, self.attn_lambda_bias)

        x = x.view(B, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_x = x

        # partition windows
        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C

        # W-MSA/SW-MSA
        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x
        x = x.view(B, H * W, C)

        # FFN
        x = shortcut + self.drop_path(x)
        
        if('ARC' == self.tuning_mode):
            
            x = self.ARC_adapter(x, self.mlp_u, self.mlp_drop, self.mlp_lambda, self.mlp_lambda_bias)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        
        else:
            
            x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x


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, tuning_mode=None):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(4 * dim)
        self.tuning_mode = tuning_mode


    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 = x.view(B, H, W, C)

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)

        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
        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
        return flops


class BasicLayer(nn.Module):
    """ A basic Swin Transformer layer 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.
        window_size (int): Local window 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
        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, window_size, svd_dim = 50,
                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,
                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, tuning_mode='ARC'):

        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.use_checkpoint = use_checkpoint
        if('ARC' == tuning_mode):
            
            self.attn_u = nn.Parameter(torch.empty(dim, svd_dim))
            nn.init.xavier_uniform_(self.attn_u)
            
            self.mlp_u = nn.Parameter(torch.empty(dim, svd_dim))
            nn.init.xavier_uniform_(self.mlp_u)

        # build blocks
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(
                dim=dim, input_resolution=input_resolution, num_heads=num_heads, 
                attn_u = self.attn_u,mlp_u = self.mlp_u, svd_dim = svd_dim,
                window_size=window_size, shift_size=0 if (i % 2 == 0) else window_size // 2, mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias, drop=drop, attn_drop=attn_drop,
                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, norm_layer=norm_layer, tuning_mode=tuning_mode)
            for i in range(depth)])

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer, tuning_mode=tuning_mode)
        else:
            self.downsample = None

    def forward(self, x):
        for blk in self.blocks:
            if not torch.jit.is_scripting() and 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}"


class PatchEmbed(nn.Module):
    """ 2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True, tuning_mode=None):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten
        self.norm_layer = norm_layer

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()


    def forward(self, x):
        B, C, H, W = x.shape
        _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
        _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")

        x = self.proj(x) 
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC

        x = self.norm(x)
        return x

class SwinTransformer(nn.Module):
    r""" Swin Transformer
        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
          https://arxiv.org/pdf/2103.14030
    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 Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window 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
        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), svd_dim = 50, 
                 window_size=7, mlp_ratio=4., qkv_bias=True,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
                 use_checkpoint=False, weight_init='', tuning_mode='ARC', **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, tuning_mode=tuning_mode)
        num_patches = self.patch_embed.num_patches
        self.patch_grid = self.patch_embed.grid_size

        # 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)
        else:
            self.absolute_pos_embed = None

        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
        self.tuning_mode = tuning_mode
        tuning_mode_list = [[tuning_mode] * depths[i_layer] for i_layer in range(self.num_layers)]

        # build layers
        layers = []
        for i_layer in range(self.num_layers):
            layers += [BasicLayer(
                dim=int(embed_dim * 2 ** i_layer),
                input_resolution=(self.patch_grid[0] // (2 ** i_layer), self.patch_grid[1] // (2 ** i_layer)),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                svd_dim = svd_dim,
                window_size=window_size,
                mlp_ratio=self.mlp_ratio,
                qkv_bias=qkv_bias,
                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, 
                tuning_mode=tuning_mode)
            ]
        self.layers = nn.Sequential(*layers)

        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()

        if weight_init != 'skip':
            self.init_weights(weight_init)


    @torch.jit.ignore
    def init_weights(self, mode=''):
        assert mode in ('jax', 'jax_nlhb', 'moco', '')
        if self.absolute_pos_embed is not None:
            trunc_normal_(self.absolute_pos_embed, std=.02)
        head_bias = -math.log(self.num_classes) if 'nlhb' in mode else 0.
        named_apply(get_init_weights_vit(mode, head_bias=head_bias), self)

    @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 get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

    def forward_features(self, x):
        x = self.patch_embed(x)
        if self.absolute_pos_embed is not None:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)
        x = self.layers(x)
        x = self.norm(x)

        x = self.avgpool(x.transpose(1, 2))  # B C 1
        x = torch.flatten(x, 1)
        return x

    def forward(self, x, labels=None):
        x = self.forward_features(x)
        x = self.head(x)
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(x.view(-1, self.num_classes), labels.view(-1))
            return loss
        else:
            return x


def _create_swin_transformer(variant, pretrained=False, **kwargs):  
    model = build_model_with_cfg(
        SwinTransformer, variant, pretrained,
        pretrained_filter_fn=checkpoint_filter_fn,
        **kwargs)

    return model



@register_model
def swin_base_patch4_window12_384(pretrained=False, **kwargs):
    """ Swin-B @ 384x384, pretrained ImageNet-22k, fine tune 1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)
    return _create_swin_transformer('swin_base_patch4_window12_384', pretrained=pretrained, **model_kwargs)


@register_model
def swin_base_patch4_window7_224(pretrained=False, **kwargs):
    """ Swin-B @ 224x224, pretrained ImageNet-22k, fine tune 1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)
    return _create_swin_transformer('swin_base_patch4_window7_224', pretrained=pretrained, **model_kwargs)


@register_model
def swin_large_patch4_window12_384(pretrained=False, **kwargs):
    """ Swin-L @ 384x384, pretrained ImageNet-22k, fine tune 1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)
    return _create_swin_transformer('swin_large_patch4_window12_384', pretrained=pretrained, **model_kwargs)


@register_model
def swin_large_patch4_window7_224(pretrained=False, **kwargs):
    """ Swin-L @ 224x224, pretrained ImageNet-22k, fine tune 1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)
    return _create_swin_transformer('swin_large_patch4_window7_224', pretrained=pretrained, **model_kwargs)


@register_model
def swin_small_patch4_window7_224(pretrained=False, **kwargs):
    """ Swin-S @ 224x224, trained ImageNet-1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 18, 2), num_heads=(3, 6, 12, 24), **kwargs)
    return _create_swin_transformer('swin_small_patch4_window7_224', pretrained=pretrained, **model_kwargs)


@register_model
def swin_tiny_patch4_window7_224(pretrained=False, **kwargs):
    """ Swin-T @ 224x224, trained ImageNet-1k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=96, depths=(2, 2, 6, 2), num_heads=(3, 6, 12, 24), **kwargs)
    return _create_swin_transformer('swin_tiny_patch4_window7_224', pretrained=pretrained, **model_kwargs)


@register_model
def swin_base_patch4_window12_384_in22k(pretrained=False, **kwargs):
    """ Swin-B @ 384x384, trained ImageNet-22k
    """
    model_kwargs = dict(
        patch_size=4, window_size=12, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)
    return _create_swin_transformer('swin_base_patch4_window12_384_in22k', pretrained=pretrained, **model_kwargs)


@register_model
def swin_base_patch4_window7_224_in22k(pretrained=False, **kwargs):
    """ Swin-B @ 224x224, trained ImageNet-22k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=128, depths=(2, 2, 18, 2), num_heads=(4, 8, 16, 32), **kwargs)
    return _create_swin_transformer('swin_base_patch4_window7_224_in22k', pretrained=pretrained, **model_kwargs)


@register_model
def swin_large_patch4_window12_384_in22k(pretrained=False, **kwargs):
    """ Swin-L @ 384x384, trained ImageNet-22k
    """
    model_kwargs = dict(
        patch_size=4, window_size=12, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)
    return _create_swin_transformer('swin_large_patch4_window12_384_in22k', pretrained=pretrained, **model_kwargs)


@register_model
def swin_large_patch4_window7_224_in22k(pretrained=False, **kwargs):
    """ Swin-L @ 224x224, trained ImageNet-22k
    """
    model_kwargs = dict(
        patch_size=4, window_size=7, embed_dim=192, depths=(2, 2, 18, 2), num_heads=(6, 12, 24, 48), **kwargs)
    return _create_swin_transformer('swin_large_patch4_window7_224_in22k', pretrained=pretrained, **model_kwargs)
Model/ARC_swin_b.py
DavidYanAnDe/ARC
import ml_collections


# 获取VIT_B_16模型
def get_b16_config():
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (16, 16)})
    config.hidden_size = 768
    config.transformer = ml_collections.ConfigDict()
    config.transformer.mlp_dim = 3072
    config.transformer.num_heads = 12
    config.transformer.num_layers = 12
    config.transformer.attention_drop_rate = 0.0
    config.transformer.mlp_drop_rate = 0.0
    # 按照什么输出进行分类
    config.classifier = 'token'
    config.adapter_dim = 50
    config.adapter_dropout = 0.0

    return config


def get_l16_config():
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (16, 16)})
    config.hidden_size = 1024
    config.transformer = ml_collections.ConfigDict()
    config.transformer.mlp_dim = 4096
    config.transformer.num_heads = 16
    config.transformer.num_layers = 24
    config.transformer.attention_drop_rate = 0.0
    config.transformer.mlp_drop_rate = 0.0
    # 按照什么输出进行分类
    config.classifier = 'token'
    config.adapter_dim = 50
    config.adapter_dropout = 0.0

    return config


def get_h14_config():
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (16, 16)})
    config.hidden_size = 1280
    config.transformer = ml_collections.ConfigDict()
    config.transformer.mlp_dim = 5120
    config.transformer.num_heads = 16
    config.transformer.num_layers = 32
    config.transformer.attention_drop_rate = 0.0
    config.transformer.mlp_drop_rate = 0.0
    # 按照什么输出进行分类
    config.classifier = 'token'
    config.adapter_dim = 50
    config.adapter_dropout = 0.0

    return config


CONFIGS = {
    "ViT-B_16": get_b16_config(),
    "ViT-L_16": get_l16_config(),
    "ViT-H_14": get_h14_config(),
}
Model/Model_Config.py
DavidYanAnDe/ARC
import copy
import logging
import math

import torch.nn as nn
from torch.nn.functional import gelu, relu
from torch.nn import  CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm
import torch
from torch.nn.modules.utils import _pair
from Utils.tools import np2th
from os.path import join as pjoin


logger = logging.getLogger(__name__)

# 定义MLP需要用到的激活函数
ACT2FN = {"gelu": torch.nn.functional.gelu, "relu": torch.nn.functional.relu}

ATTENTION_Q = "MultiHeadDotProductAttention_1/query"
ATTENTION_K = "MultiHeadDotProductAttention_1/key"
ATTENTION_V = "MultiHeadDotProductAttention_1/value"
ATTENTION_OUT = "MultiHeadDotProductAttention_1/out"
FNN_0 = "MlpBlock_3/Dense_0"
FNN_1 = "MlpBlock_3/Dense_1"
ATTENTION_NORM = "LayerNorm_0"
MLP_NORM = "LayerNorm_2"

# 图像输入的embedding层
class Embedding(nn.Module):
    def __init__(self, config, img_size, in_channel = 3):
        super(Embedding, self).__init__()

        # image_size(n,n)    patch_size(t, t)
        image_size = _pair(img_size)
        patch_size = config.patches['size']
        num_patches = (image_size[0] // patch_size[0]) * (image_size[1] // patch_size[1])

        # 将[H,W,3]转换为[224/16, 224/16, hidden_size]
        self.projective = Conv2d(in_channels=in_channel,
                                    out_channels=config.hidden_size,
                                    kernel_size=patch_size,
                                    stride=patch_size)

        # 可学习的cls, cls维度[1, 1, hidden_dim]
        self.cls = nn.Parameter(torch.zeros(1, 1, config.hidden_size))

        # 可学习的position_embedding, [1, patch_num+1, hidden_dim]
        self.position_embedding = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))

        # 输入encode前有一个dropout
        self.dropout = Dropout(config.transformer["drop_rate"])

    def forward(self, input_image):
        # [B, hidden_size, patch_size, patch_size]
        projective_embedding = self.projective(input_image)
        # 展开，并且调换维度，变成[B, num_patches, hidden_size]
        projective_embedding = projective_embedding.flatten(2).transpose(-1, -2)

        # num_patches, 加入class_token. [B, num_patches, hidden_size]和[1, 1, hidden_size]
        # 首先将cls变为[B, 1, hidden_size]
        batch_size = input_image.shape[0]
        cls = self.cls.expand(batch_size, -1, -1)
        # 拼接cls和projective_embedding, 变为[B, num_patches+1, hidden_size]
        embedding_features = torch.concat([cls, projective_embedding], dim=1)

        output_features = embedding_features + self.position_embedding
        output_features = self.dropout(output_features)
        return output_features

# encoder中的自注意力层
class Attention(nn.Module):
    def __init__(self, config):
        super(Attention, self).__init__()

        hidden_dim = config.hidden_size
        # 一共多少个注意力头
        self.head_num = config.transformer["num_heads"]
        # 每个头多少维度
        self.head_hidden_dim = hidden_dim // self.head_num
        self.all_head_size = self.head_hidden_dim * self.head_num

        # q,k,v三个参数矩阵
        self.query = Linear(hidden_dim, self.all_head_size)
        self.key = Linear(hidden_dim, self.all_head_size)
        self.value = Linear(hidden_dim, self.all_head_size)

        # 计算注意力头的归一化操作，对每一行进行归一化，也就是列归一化，
        self.softmax = Softmax(dim=-1)
        # 自注意力后需要一个project矩阵，进行特征的转换
        self.projective = Linear(hidden_dim, hidden_dim)
        # dropout在attention输出后
        self.attn_dropout = Dropout(config.transformer["attention_drop_rate"])
        self.proj_dropout = Dropout(config.transformer["attention_drop_rate"])

    def transpose_head(self, mix_metric):
        # size可以看成是一种tuple,tuple相加是维度的相加
        # new_x_shape可以看作是 [batch_size,n,H,n/H] （四维数组）
        new_x_shape = mix_metric.size()[:-1] + (self.head_num, self.head_hidden_dim)
        # *只改变形状，不改变内存地址
        mix_metric = mix_metric.view(*new_x_shape)
        # permute调整顺序
        return mix_metric.permute(0, 2, 1, 3)

    def forward(self, hidden_states):
        # hidden_states维度 [Batch_size, num_patches, hidden_dim]，计算出混合q,k,v
        mix_query_layer = self.query(hidden_states)
        mix_key_layer = self.key(hidden_states)
        mix_value_layer = self.value(hidden_states)

        # 分离注意力头，从[Batch_size, num_patches, hidden_dim]到[Batch_size, head_num, num_patches, head_dim]
        query_layer = self.transpose_head(mix_query_layer)
        key_layer = self.transpose_head(mix_key_layer)
        value_layer = self.transpose_head(mix_value_layer)

        # 计算attention score，注意需要归一化以及softmax
        attention_metric = torch.matmul(query_layer, key_layer.transpose(-2,-1))
        attention_metric = attention_metric / math.sqrt(self.head_hidden_dim)
        attention_probs = self.softmax(attention_metric)
        # dropout注意力矩阵
        attention_probs = self.attn_dropout(attention_probs)
        # 计算注意力矩阵m和value的乘积，[B, head_num, patch_num, head_dim]
        attention_features = torch.matmul(attention_probs, value_layer)

        # 对attention_features进行重新构建，恢复[batch_size,n,d]的格式
        attention_features = attention_features.permute(0, 2, 1, 3).contiguous()
        new_output_features_shape = attention_features.size()[:-2] + (self.all_head_size, )
        attention_feature = attention_features.view(*new_output_features_shape)

        # 输出前需要加一层映射
        output_features = self.projective(attention_feature)
        output_features = self.proj_dropout(output_features)

        return output_features

# encoder里面的MLP层
class MLP(nn.Module):
    def __init__(self, config):
        super(MLP, self).__init__()
        # 两层MLP
        self.fc1 = Linear(config.hidden_size, config.transformer["mlp_dim"])
        # 中间有一层激活函数
        self.activation = ACT2FN["gelu"]
        self.fc2 = Linear(config.transformer["mlp_dim"], config.hidden_size)
        # MLP输出前需要加一层dropout
        self.dropout = Dropout(config.transformer["drop_rate"])

    def forward(self, input_features):
        # 逐层输出计算特征
        output_features = self.fc1(input_features)
        output_features = self.activation(output_features)
        output_features = self.dropout(output_features)
        output_features = self.fc2(output_features)
        output_features = self.dropout(output_features)

        return output_features

class Encoder(nn.Module):
    def __init__(self, config):
        super(Encoder, self).__init__()

        self.hidden_size = config.hidden_size
        self.attention = Attention(config)
        self.mlp = MLP(config)
        # layer norm在输入attention和mlp前
        self.att_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)
        self.fnn_layer_norm = LayerNorm(self.hidden_size, eps=1e-6)

    def forward(self, hidden_states):
        # 残差网络的特征
        res_att_feature = hidden_states
        # 输入前先进行layer norm
        hidden_states = self.att_layer_norm(hidden_states)
        attention_output = self.attention(hidden_states)
        attention_output = attention_output + res_att_feature

        res_fnn_feature = attention_output
        fnn_output = self.fnn_layer_norm(attention_output)
        fnn_output = self.mlp(fnn_output)
        fnn_output = fnn_output + res_fnn_feature
        return fnn_output

    def load_from(self, weights, n_block):
        ROOT = f"Transformer/encoderblock_{n_block}"
        with torch.no_grad():
            query_weight = np2th(weights[pjoin(ROOT, ATTENTION_Q, "kernel")]).view(self.hidden_size, self.hidden_size).t()
            key_weight = np2th(weights[pjoin(ROOT, ATTENTION_K, "kernel")]).view(self.hidden_size, self.hidden_size).t()
            value_weight = np2th(weights[pjoin(ROOT, ATTENTION_V, "kernel")]).view(self.hidden_size, self.hidden_size).t()
            out_weight = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "kernel")]).view(self.hidden_size, self.hidden_size).t()

            query_bias = np2th(weights[pjoin(ROOT, ATTENTION_Q, "bias")]).view(-1)
            key_bias = np2th(weights[pjoin(ROOT, ATTENTION_K, "bias")]).view(-1)
            value_bias = np2th(weights[pjoin(ROOT, ATTENTION_V, "bias")]).view(-1)
            out_bias = np2th(weights[pjoin(ROOT, ATTENTION_OUT, "bias")]).view(-1)

            self.attention.query.weight.copy_(query_weight)
            self.attention.key.weight.copy_(key_weight)
            self.attention.value.weight.copy_(value_weight)
            self.attention.projective.weight.copy_(out_weight)
            self.attention.query.bias.copy_(query_bias)
            self.attention.key.bias.copy_(key_bias)
            self.attention.value.bias.copy_(value_bias)
            self.attention.projective.bias.copy_(out_bias)

            mlp_weight_0 = np2th(weights[pjoin(ROOT, FNN_0, "kernel")]).t()
            mlp_weight_1 = np2th(weights[pjoin(ROOT, FNN_1, "kernel")]).t()
            mlp_bias_0 = np2th(weights[pjoin(ROOT, FNN_0, "bias")]).t()
            mlp_bias_1 = np2th(weights[pjoin(ROOT, FNN_1, "bias")]).t()

            self.mlp.fc1.weight.copy_(mlp_weight_0)
            self.mlp.fc2.weight.copy_(mlp_weight_1)
            self.mlp.fc1.bias.copy_(mlp_bias_0)
            self.mlp.fc2.bias.copy_(mlp_bias_1)

            self.att_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "scale")]))
            self.att_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, ATTENTION_NORM, "bias")]))
            self.fnn_layer_norm.weight.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "scale")]))
            self.fnn_layer_norm.bias.copy_(np2th(weights[pjoin(ROOT, MLP_NORM, "bias")]))

class Transformer(nn.Module):
    def __init__(self, config):
        super(Transformer, self).__init__()
        # encoder层的数量
        self.layer_num = config.transformer["num_layers"]
        # encoder按照存储每一层
        self.layer = nn.ModuleList()
        # 添加每一层
        for _ in range(self.layer_num):
            encoder = Encoder(config)
            # 这里需要深度拷贝
            self.layer.append(copy.deepcopy(encoder))

        # 最后有一个layer_norm输出
        self.layer_norm = LayerNorm(config.hidden_size, eps=1e-6)
    def forward(self, input_feature):
        for layer_encoder in self.layer:
            input_feature = layer_encoder(input_feature)

        output = self.layer_norm(input_feature)
        return output

class HeadMLP(nn.Module):
    def __init__(self, config, num_classes):
        super(HeadMLP, self).__init__()
        self.head = Linear(config.hidden_size, num_classes)

    def forward(self, input_features):
        # input维度[B, num_patches, hidden_dim], 获取cls的表征就行，也就是[B, 0, hidden_dim]
        # cls_token维度[B, hidden_dim]
        cls_token = input_features[:, 0]

        # output维度[B, num_classes]
        output = self.head(cls_token)
        return output

class VisionTransformer(nn.Module):
    def __init__(self, config, img_size=224, num_classes=21843, zero_head=False):
        super(VisionTransformer, self).__init__()
        self.zero_head = zero_head
        self.num_classes = num_classes

        self.embedding = Embedding(config, img_size)
        self.transformer = Transformer(config)

        self.head_mlp = HeadMLP(config, num_classes)

    def forward(self, image, labels=None):
        image_projective = self.embedding(image)
        output_representation = self.transformer(image_projective)
        # 只需要使用到cls这一个token，因此需要提取出来
        prediction = self.head_mlp(output_representation)

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(prediction.view(-1, self.num_classes), labels.view(-1))
            return loss
        else:
            return prediction

    # 加载预训练模型
    def load_from(self, weights):
        with torch.no_grad():
            # 初始化头部，或者加载预训练参数
            if self.zero_head:
                nn.init.zeros_(self.head_mlp.head.weight)
                nn.init.zeros_(self.head_mlp.head.bias)
            else:
                self.head_mlp.head.weight.copy_(np2th(weights["head/kernel"]).t())
                self.head_mlp.head.weight.copy_(np2th(weights["head/bias"]).t())

            # 加载embedding参数
            self.embedding.projective.weight.copy_(np2th(weights["embedding/kernel"], conv=True))
            self.embedding.projective.bias.copy_(np2th(weights["embedding/bias"]))
            self.embedding.cls.copy_(np2th(weights["cls"]))

            posemb = np2th(weights["Transformer/posembed_input/pos_embedding"])
            posemb_new = self.embedding.position_embedding
            if posemb.size() == posemb_new.size():
                self.embedding.position_embedding.copy_(posemb)

            # 加载transformer最后输出的参数
            self.transformer.layer_norm.weight.copy_(np2th(weights["Transformer/encoder_norm/scale"]))
            self.transformer.layer_norm.bias.copy_(np2th(weights["Transformer/encoder_norm/bias"]))

            for bname, encoder in self.transformer.named_children():
                for uname, unit in encoder.named_children():
                    unit.load_from(weights, n_block=uname)


Model/VIT.py
DavidYanAnDe/ARC
import torch

'''
    frozen the original parameters pf pre-trained models.
    only active the learnable parameters of ARC modules and head layer.
'''

def ViT_ARC_Frozen(model):
    for name, para in model.named_parameters():
        if "head" in name:
            para.requires_grad = True
        elif "adapter" in name:
            para.requires_grad = True
        elif "down_projection" in name:
            para.requires_grad = True
        else:
            para.requires_grad = False

def Swin_ARC_Frozen(model):
    for name, para in model.named_parameters():
        if "head" in name:
            para.requires_grad = True
        elif "attn_u" in name:
            para.requires_grad = True
        elif "mlp_u" in name:
            para.requires_grad = True
        elif "lambda" in name:
            para.requires_grad = True
        else:
            para.requires_grad = False

Utils/Frozen_weight.py
DavidYanAnDe/ARC
import logging
import math

from torch.optim.lr_scheduler import LambdaLR

logger = logging.getLogger(__name__)

class ConstantLRSchedule(LambdaLR):
    """ Constant learning rate schedule.
        用于创建常量的learning rate学习器
    """
    def __init__(self, optimizer, last_epoch=-1):
        super(ConstantLRSchedule, self).__init__(optimizer, lambda _: 1.0, last_epoch=last_epoch)


class WarmupConstantSchedule(LambdaLR):
    """ Linear warmup and then constant.
        Linearly increases learning rate schedule from 0 to 1 over `warmup_steps` training steps.
        Keeps learning rate schedule equal to 1. after warmup_steps.
        自主创建的学习器，具有warnup功能
    """
    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
        self.warmup_steps = warmup_steps
        super(WarmupConstantSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return float(step) / float(max(1.0, self.warmup_steps))
        return 1.


class WarmupLinearSchedule(LambdaLR):
    """ Linear warmup and then linear decay.
        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
        Linearly decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps.
    """
    def __init__(self, optimizer, warmup_steps, t_total, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.t_total = t_total
        super(WarmupLinearSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return float(step) / float(max(1, self.warmup_steps))
        return max(0.0, float(self.t_total - step) / float(max(1.0, self.t_total - self.warmup_steps)))


class WarmupCosineSchedule(LambdaLR):
    """ Linear warmup and then cosine decay.
        Linearly increases learning rate from 0 to 1 over `warmup_steps` training steps.
        Decreases learning rate from 1. to 0. over remaining `t_total - warmup_steps` steps following a cosine curve.
        If `cycles` (default=0.5) is different from default, learning rate follows cosine function after warmup.
    """
    def __init__(self, optimizer, warmup_steps, t_total, cycles=.5, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.t_total = t_total
        self.cycles = cycles
        super(WarmupCosineSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return float(step) / float(max(1.0, self.warmup_steps))
        # progress after warmup
        progress = float(step - self.warmup_steps) / float(max(1, self.t_total - self.warmup_steps))
        return max(0.000001, 0.5 * (1. + math.cos(math.pi * float(self.cycles) * 2.0 * progress)))
Utils/scheduler.py
DavidYanAnDe/ARC
import torch
import os

def count_parameters(model):
    params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return params/1000000

# 将numpy类型的权重转化为torch tensor
def np2th(weights, conv=False):
    """Possibly convert HWIO to OIHW."""
    if conv:
        weights = weights.transpose([3, 2, 0, 1])
    return torch.from_numpy(weights)

# 用于统计计算结果的平均值
class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
Utils/tools.py
DavidYanAnDe/ARC
# coding=utf-8
from __future__ import absolute_import, division, print_function

import logging
import argparse
import os
import random
import numpy as np
import torch

from Model.Model_Config import CONFIGS

from Model.ARC_swin_b import swin_base_patch4_window7_224
from Utils.tools import count_parameters, AverageMeter

from tqdm import tqdm
from Utils.scheduler import WarmupLinearSchedule, WarmupCosineSchedule


from Utils.Frozen_weight import Swin_ARC_Frozen
from Data_process.VTAB_config import DATA_CONFIGS
from Data_process.VTAB_loader import get_data

logger = logging.getLogger(__name__)

def save_model(args, model):
    model_to_save = model.Module if hasattr(model, "module") else model
    model_checkpoint = os.path.join(args.output_dir, "%s_checkpoint.bin" % args.name)
    torch.save(model_to_save.state_dict(), model_checkpoint)
    logger.info("Saved model checkpoint to [DIR: %s]", args.output_dir)

def simple_accuracy(preds, labels):
    return (preds == labels).mean()

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def setup(args):
    data_cfg = DATA_CONFIGS[args.dataset]

    num_classes = data_cfg.Num_Classes
    model = swin_base_patch4_window7_224(num_classes=num_classes,pretrained=True)

    model.to(args.device)

    Swin_ARC_Frozen(model)

    num_params = count_parameters(model)
    # logger.info("{}".format(config))
    logger.info("Training HypeParameters %s", args)
    logger.info("Total Parameter: \t%2.2fM" % num_params)

    for name, para in model.named_parameters():
        if para.requires_grad == True:
            print(name)

    return model,num_params

def valid(args, model, test_loader):
    eval_losses = AverageMeter()

    logger.info("***** Running Validation *****")
    logger.info("  Num steps = %d", len(test_loader))
    logger.info("  Batch size = %d", args.eval_batch_size)

    model.eval()
    all_preds, all_label = [], []
    epoch_iterator = tqdm(test_loader,
                      desc="Testing (X / X steps) (loss = X.X)",
                      bar_format='{l_bar}-{r_bar}',
                      dynamic_ncols=True)
    loss_fct = torch.nn.CrossEntropyLoss()
    for step, batch in enumerate(epoch_iterator):
        batch = tuple(t.to(args.device) for t in batch)
        x, label = batch
        with torch.no_grad():
            logits = model(x)

            eval_loss = loss_fct(logits, label)
            eval_losses.update(eval_loss.item())

            prediction = torch.argmax(logits, dim=-1)

        if len(all_preds) == 0:
            all_preds.append(prediction.detach().cpu().numpy())
            all_label.append(label.detach().cpu().numpy())
        else:
            all_preds[0] = np.append(
                all_preds[0], prediction.detach().cpu().numpy(), axis=0)
            all_label[0] = np.append(
                all_label[0], label.detach().cpu().numpy(), axis=0)

        epoch_iterator.set_description("Validating... (loss=%2.5f)" % eval_losses.val)

    # 计算平均准确率
    all_preds = all_preds[0]
    all_label = all_label[0]
    accuracy = simple_accuracy(all_preds, all_label)

    logger.info("\n")
    logger.info("Validation Results")
    logger.info("Valid Loss: %2.5f" % eval_losses.avg)
    logger.info("Valid Accuracy: %2.5f" % accuracy)

    return accuracy


def train(args, model):
    os.makedirs(args.output_dir, exist_ok=True)

    train_loader, test_loader = get_data(data_path=args.data_path, name=args.dataset, batch_size=args.train_batch_size)

    epoch_steps = len(train_loader)
    total_steps = epoch_steps * args.num_epochs
    warm_steps = epoch_steps * args.warmup_epochs
    evl_steps = args.eval_every * epoch_steps

    optimizer = torch.optim.AdamW(model.parameters(),
                                lr=args.learning_rate,
                                weight_decay=args.weight_decay)

    if args.decay_type == "cosine":
        schedule = WarmupCosineSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)
    else:
        schedule = WarmupLinearSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)

    logger.info("***** Running training *****")
    logger.info("  Total optimization steps = %d", total_steps)
    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size)
    logger.info("  Total train batch size = %d", args.train_batch_size * args.gradient_accumulation_steps )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)


    model.zero_grad()
    set_seed(args)
    losses = AverageMeter()
    global_step, best_acc = 0, 0


    while True:
        model.train()
        epoch_iterator = tqdm(train_loader,
                                desc="Training (X / X steps) (loss = X.X)",
                                bar_format='{l_bar}-{r_bar}',
                                dynamic_ncols=True)
        for step, batch in enumerate(epoch_iterator):
            batch = tuple(t.to(args.device) for t in batch)
            x, label = batch

            loss = model(x, label)

            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            loss.backward()

            if (step + 1) % args.gradient_accumulation_steps == 0:
                losses.update(loss.item())

                torch.nn.utils.clip_grad_norm(model.parameters(), args.max_grad_norm)
                optimizer.step()
                schedule.step()
                optimizer.zero_grad()
                global_step += 1

                epoch_iterator.set_description(
                    "Training (%d / %d Steps) (loss=%2.5f)" % (global_step, total_steps, losses.val))

                if (global_step + 1) % evl_steps == 0:
                    accuracy = valid(args, model, test_loader)
                    if best_acc < accuracy:
                        save_model(args, model)
                        best_acc = accuracy
                    model.train()

            if (global_step + 1) % total_steps == 0:
                break

        losses.reset()
        if (global_step + 1) % total_steps == 0:
            break

    logger.info("Best Accuracy: \t%f" % best_acc)
    logger.info("End Training!")
    return best_acc

def main(config=None):
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("--data_path",  default= "/root/autodl-tmp/VTAB_Data", help="where is the data store, should be changed according to local file")

    parser.add_argument("--name", default="ViT-B_16_cifar_att_mlp_independent_50dim_bias0_repadapterStructure_error",
                        help="Name of this run. Used for monitoring.")
    parser.add_argument("--dataset", choices=["caltech101", "cifar", "smallnorb_ele", "sun397", "svhn", "resisc45", "dmlab"
                                              "oxford_iiit_pet","oxford_flowers102", 'dtd', "kitti", "eurosat","clevr_count"
                                              ,"dsprites_ori"],
                        default= "kitti",  help="Which downstream task.")


    parser.add_argument("--tuning_mode", choices=["ARC_att", "ARC"],
                        default= "ARC",  help="tuning mode,can be ARC_att or ARC")

    parser.add_argument("--model_type", choices=["ViT-B_16", "ViT-L_16", "ViT-H_14"],
                        default="ViT-B_16",
                        help="Which variant to use.")
    parser.add_argument("--pretrained_dir", type=str, default="ViT-B_16.npz",
                        help="Where to search for pretrained ViT models.")
    parser.add_argument("--output_dir", default="output", type=str,
                        help="The output directory where checkpoints will be written.")

    parser.add_argument("--img_size", default=224, type=int,
                        help="Resolution size")
    parser.add_argument("--train_batch_size", default=32, type=int,
                        help="Total batch size for training.")
    parser.add_argument("--eval_batch_size", default=256, type=int,
                        help="Total batch size for eval.")
    parser.add_argument("--eval_every", default=1, type=int,
                        help="Run prediction on validation set every so many steps."
                             "Will always run one evaluation at the end of training.")

    parser.add_argument("--learning_rate", default=0.01, type=float,
                        help="The initial learning rate for SGD.")
    parser.add_argument("--weight_decay", default=0, type=float,
                        help="Weight decay if we apply some.")
    parser.add_argument("--num_epochs", default=100, type=int,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--decay_type", choices=["cosine", "linear"], default="cosine",
                        help="How to decay the learning rate.")
    parser.add_argument("--warmup_epochs", default=10, type=int,
                        help="Step of training to perform learning rate warmup for.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")

    parser.add_argument("--vit_drop", default=0, type=float,
                        help="dropout rate of ViT.")
    parser.add_argument("--adapt_drop", default=0.1, type=float,
                        help="dropout rate of adapter.")

    args = parser.parse_args()


    if config:
        args.dataset = config["dataset"]
        name = "SwinB_VTAB_dim50_ARCatt_" + args.dataset + "_lr%f" % config["lr"] + "_wd%f" % config["wd"] + "_vdrop0.0_adrop%f" %  config["a_drop"]
        args.name = name
        args.learning_rate = config["lr"]
        args.weight_decay = config["wd"]
        args.a_drop = config["a_drop"]


    # Setup CUDA, GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = torch.cuda.device_count()
    args.device = device

    # Setup logging
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
    logger.warning("Process device: %s, n_gpu: %s" % (args.device, args.n_gpu))

    # Set seed
    set_seed(args)

    # Model & Tokenizer Setup
    model, para_number = setup(args)

    # Training
    best_acc = train(args, model)
    
    
    return best_acc, para_number

if __name__ == "__main__":
    config_list = []

    cifar_config ={"dataset" : "cifar",
        "lr" : 0.005,
        "wd" : 0.01,
        "a_drop" : 0.1}
    config_list.append(cifar_config)
    
    caltech101_config ={"dataset" : "caltech101",
        "lr" : 0.003,
        "wd" : 0.05,
        "a_drop" : 0.1}
    config_list.append(caltech101_config)
    
    dtd_config ={"dataset" : "dtd",
        "lr" : 5e-3,
        "wd" : 5e-2,
        "a_drop" : 0.8}
    config_list.append(dtd_config)
    
    oxford_flowers102_config ={"dataset" : "oxford_flowers102",
        "lr" : 0.005,
        "wd" : 0.00005,
        "a_drop" : 0.5}
    config_list.append(oxford_flowers102_config)
    
    oxford_iiit_pet_config ={"dataset" : "oxford_iiit_pet",
        "lr" : 0.01,
        "wd" : 0.05,
        "a_drop" : 0.1}
    config_list.append(oxford_iiit_pet_config)
    
    svhn_config ={"dataset" : "svhn",
        "lr" : 0.01,
        "wd" : 0.05,
        "a_drop" : 0.1}
    config_list.append(svhn_config)
    
    sun397_config ={"dataset" : "sun397",
        "lr" : 0.005,
        "wd" : 0.00005,
        "a_drop" : 0.1}
    config_list.append(sun397_config)
    
    eurosat_config ={"dataset" : "eurosat",
        "lr" : 0.003,
        "wd" : 0,
        "a_drop" : 0.1}
    config_list.append(eurosat_config)
    
    resisc45_config ={"dataset" : "resisc45",
        "lr" : 0.01,
        "wd" : 0.0,
        "a_drop" : 0.5}
    config_list.append(resisc45_config)
    
    patch_camelyon_config ={"dataset" : "patch_camelyon",
        "lr" : 0.005,
        "wd" : 0.00005,
        "a_drop" : 0.1}
    config_list.append(patch_camelyon_config)
    
    print(os.path.exists("Train/output"))
    
    diabetic_retinopathy_config ={"dataset" : "diabetic_retinopathy",
        "lr" : 0.005,
        "wd" : 0.00005,
        "a_drop" : 0.1}
    config_list.append(diabetic_retinopathy_config)

    clevr_count_config ={"dataset" : "clevr_count",
        "lr" : 0.002,
        "wd" : 0.05,
        "a_drop" : 0.5}
    config_list.append(clevr_count_config)

    clevr_dist_config ={"dataset" : "clevr_dist",
        "lr" : 0.001,
        "wd" : 0.05,
        "a_drop" : 0.1}
    config_list.append(clevr_dist_config)

    dmlab_config ={"dataset" : "dmlab",
        "lr" : 0.005,
        "wd" : 0.01,
        "a_drop" : 0.1}
    config_list.append(dmlab_config)

    kitti_config ={"dataset" : "kitti",
        "lr" : 0.01,
        "wd" : 0.0,
        "a_drop" : 0.1}
    config_list.append(kitti_config)

    smallnorb_azi_config ={"dataset" : "smallnorb_azi",
        "lr" : 0.01,
        "wd" : 0.005,
        "a_drop" : 0.1}
    config_list.append(smallnorb_azi_config)

    smallnorb_ele_config ={"dataset" : "smallnorb_ele",
        "lr" : 0.001,
        "wd" : 0.05,
        "a_drop" : 0.1}
    config_list.append(smallnorb_ele_config)

    dsprites_loc_config ={"dataset" : "dsprites_loc",
        "lr" : 0.01,
        "wd" : 0.0,
        "a_drop" : 0.1}
    config_list.append(dsprites_loc_config)

    dsprites_ori_config ={"dataset" : "dsprites_ori",
        "lr" : 0.005,
        "wd" : 0.0,
        "a_drop" : 0.5}
    config_list.append(dsprites_ori_config)

    results_list = []

    for config in config_list:
        print(config)
        result = {}
        result["dataset"] = config["dataset"]
        acc, para_num = main(config)
        result["acc"] = acc
        result["para_num"] = para_num
        results_list.append(result)
        
    count = 0
    for result in results_list:
        print(result)
        count = count + result["para_num"]
        
    print(count/19)
train_ARC_swin.py
DavidYanAnDe/ARC
# coding=utf-8
from __future__ import absolute_import, division, print_function

import logging
import argparse
import os
import random
import numpy as np
import torch

from Model.Model_Config import CONFIGS
from Model.ARC_new import ARCVisionTransformer
from Utils.tools import count_parameters, AverageMeter

from Data_process.VTAB_loader import get_data
from Data_process.VTAB_config import DATA_CONFIGS

from tqdm import tqdm
from Utils.scheduler import WarmupLinearSchedule, WarmupCosineSchedule
from Utils.Frozen_weight import ViT_ARC_Frozen

logger = logging.getLogger(__name__)

def save_model(args, model):
    model_to_save = model.Module if hasattr(model, "module") else model
    model_checkpoint = os.path.join(args.output_dir, "%s_checkpoint.bin" % args.name)
    torch.save(model_to_save.state_dict(), model_checkpoint)
    logger.info("Saved model checkpoint to [DIR: %s]", args.output_dir)

def simple_accuracy(preds, labels):
    return (preds == labels).mean()

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def setup(args):
    # get configure information
    config = CONFIGS[args.model_type]
    config.transformer["mlp_drop_rate"] = args.vit_drop
    config.adapter_dropout = args.adapt_drop

    num_classes = DATA_CONFIGS[args.dataset].Num_Classes
    model = ARCVisionTransformer(config, args.img_size, num_classes=num_classes, zero_head=True)
    model.load_from(np.load(args.pretrained_dir))
    
    model.to(args.device)

    # frozen weights
    ViT_ARC_Frozen(model)

    num_params = count_parameters(model)
    logger.info("{}".format(config))
    logger.info("Training HypeParameters %s", args)
    logger.info("Total Parameter: \t%2.2fM" % num_params)

    for name, para in model.named_parameters():
        if para.requires_grad == True:
            print(name)

    return model, num_params

def valid(args, model, test_loader):
    eval_losses = AverageMeter()

    logger.info("***** Running Validation *****")
    logger.info("  Num steps = %d", len(test_loader))
    logger.info("  Batch size = %d", args.eval_batch_size)

    model.eval()
    all_preds, all_label = [], []
    epoch_iterator = tqdm(test_loader,
                      desc="Testing (X / X steps) (loss = X.X)",
                      bar_format='{l_bar}-{r_bar}',
                      dynamic_ncols=True)
    loss_fct = torch.nn.CrossEntropyLoss()
    for step, batch in enumerate(epoch_iterator):
        batch = tuple(t.to(args.device) for t in batch)
        x, label = batch
        with torch.no_grad():
            logits = model(x)
            eval_loss = loss_fct(logits, label)
            eval_losses.update(eval_loss.item())
            prediction = torch.argmax(logits, dim=-1)

        if len(all_preds) == 0:
            all_preds.append(prediction.detach().cpu().numpy())
            all_label.append(label.detach().cpu().numpy())
        else:
            all_preds[0] = np.append(
                all_preds[0], prediction.detach().cpu().numpy(), axis=0)
            all_label[0] = np.append(
                all_label[0], label.detach().cpu().numpy(), axis=0)

        epoch_iterator.set_description("Validating... (loss=%2.5f)" % eval_losses.val)

    all_preds = all_preds[0]
    all_label = all_label[0]
    accuracy = simple_accuracy(all_preds, all_label)

    logger.info("\n")
    logger.info("Validation Results")
    logger.info("Valid Loss: %2.5f" % eval_losses.avg)
    logger.info("Valid Accuracy: %2.5f" % accuracy)

    return accuracy


def train(args, model):
    # output dir
    os.makedirs(args.output_dir, exist_ok=True)

    train_loader, test_loader = get_data(data_path=args.data_path, name=args.dataset, batch_size=args.train_batch_size)

    epoch_steps = len(train_loader)
    total_steps = epoch_steps * args.num_epochs
    warm_steps = epoch_steps * args.warmup_epochs
    evl_steps = args.eval_every * epoch_steps


    optimizer = torch.optim.AdamW(model.parameters(),
                                lr=args.learning_rate,
                                weight_decay=args.weight_decay)

    if args.decay_type == "cosine":
        schedule = WarmupCosineSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)
    else:
        schedule = WarmupLinearSchedule(optimizer, warmup_steps=warm_steps, t_total=total_steps)

    logger.info("***** Running training *****")
    logger.info("  Total optimization steps = %d", total_steps)
    logger.info("  Instantaneous batch size per GPU = %d", args.train_batch_size)
    logger.info("  Total train batch size = %d", args.train_batch_size * args.gradient_accumulation_steps )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)

    model.zero_grad()
    set_seed(args)
    losses = AverageMeter()
    global_step, best_acc = 0, 0

    while True:
        model.train()
        epoch_iterator = tqdm(train_loader,
                                desc="Training (X / X steps) (loss = X.X)",
                                bar_format='{l_bar}-{r_bar}',
                                dynamic_ncols=True)
        for step, batch in enumerate(epoch_iterator):
            batch = tuple(t.to(args.device) for t in batch)
            x, label = batch

            loss = model(x, label)

            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps

            loss.backward()

            if (step + 1) % args.gradient_accumulation_steps == 0:
                losses.update(loss.item())

                torch.nn.utils.clip_grad_norm(model.parameters(), args.max_grad_norm)
                optimizer.step()
                schedule.step()
                optimizer.zero_grad()
                global_step += 1

                epoch_iterator.set_description(
                    "Training (%d / %d Steps) (loss=%2.5f)" % (global_step, total_steps, losses.val))

                if (global_step + 1) % evl_steps == 0:
                    accuracy = valid(args, model, test_loader)
                    if best_acc < accuracy:
                        save_model(args, model)
                        best_acc = accuracy
                    model.train()

            if (global_step + 1) % total_steps == 0:
                break

        losses.reset()
        if (global_step + 1) % total_steps == 0:
            break

    logger.info("Best Accuracy: \t%f" % best_acc)
    logger.info("End Training!")
    return best_acc

def main():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("--data_path",  default= "/root/autodl-tmp/VTAB_Data", help="where is the data store, should be changed according to local file")

    parser.add_argument("--name", default="ViT-B_16_dtd",
                        help="Name of this run. Used for monitoring.")
    parser.add_argument("--dataset", choices=["caltech101", "cifar", "smallnorb_ele", "sun397", "svhn", "resisc45", "dmlab"
                                              "oxford_iiit_pet","oxford_flowers102", 'dtd', "kitti", "eurosat","clevr_count"
                                              ,"dsprites_ori"],
                        default= "kitti",  help="Which downstream task.")


    parser.add_argument("--tuning_mode", choices=["ARC_att", "ARC"],
                        default= "ARC",  help="tuning mode,can be ARC_att or ARC")

    parser.add_argument("--model_type", choices=["ViT-B_16", "ViT-L_16", "ViT-H_14"],
                        default="ViT-B_16",
                        help="Which variant to use.")
    parser.add_argument("--pretrained_dir", type=str, default="ViT-B_16.npz",
                        help="Where to search for pretrained ViT models.")
    parser.add_argument("--output_dir", default="output", type=str,
                        help="The output directory where checkpoints will be written.")

    parser.add_argument("--img_size", default=224, type=int,
                        help="Resolution size")
    parser.add_argument("--train_batch_size", default=32, type=int,
                        help="Total batch size for training.")
    parser.add_argument("--eval_batch_size", default=256, type=int,
                        help="Total batch size for eval.")
    parser.add_argument("--eval_every", default=1, type=int,
                        help="Run prediction on validation set every so many steps."
                             "Will always run one evaluation at the end of training.")

    parser.add_argument("--learning_rate", default=0.01, type=float,
                        help="The initial learning rate for SGD.")
    parser.add_argument("--weight_decay", default=0, type=float,
                        help="Weight decay if we apply some.")
    parser.add_argument("--num_epochs", default=100, type=int,
                        help="Total number of training epochs to perform.")
    parser.add_argument("--decay_type", choices=["cosine", "linear"], default="cosine",
                        help="How to decay the learning rate.")
    parser.add_argument("--warmup_epochs", default=10, type=int,
                        help="Step of training to perform learning rate warmup for.")
    parser.add_argument("--max_grad_norm", default=1.0, type=float,
                        help="Max gradient norm.")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")

    parser.add_argument("--vit_drop", default=0, type=float,
                        help="dropout rate of ViT.")
    parser.add_argument("--adapt_drop", default=0.1, type=float,
                        help="dropout rate of adapter.")

    args = parser.parse_args()

    # Setup CUDA, GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = torch.cuda.device_count()
    args.device = device

    # Setup logging
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                        datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO)
    logger.warning("Process device: %s, n_gpu: %s" % (args.device, args.n_gpu))

    # Set seed
    set_seed(args)

    # Model & Tokenizer Setup
    model, num_params = setup(args)

    # Training
    best_acc = train(args, model)

    return best_acc, num_params

if __name__ == "__main__":
    main()
vtab_ARC_train.py
DavidYanAnDe/ARC