Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,2 @@
+**__pycache__**
+__pycache__\

app.py

@@ -0,0 +1,243 @@
+import gradio as gr
+from PIL import Image
+import numpy as np
+import os
+import torch
+import torchvision.transforms as transforms
+from models import get_model
+
+# os.environ["http_proxt"] = ""
+# 常量定义
+MEAN = {
+    "imagenet": [0.485, 0.456, 0.406],
+    "clip": [0.48145466, 0.4578275, 0.40821073]
+}
+
+STD = {
+    "imagenet": [0.229, 0.224, 0.225],
+    "clip": [0.26862954, 0.26130258, 0.27577711]
+}
+
+
+
+def collect_examples(base_dir="examples"):
+    """
+    递归扫描 base_dir 下所有图片，返回形如
+    [
+        ["examples/real/00001.jpg"],
+        ["examples/real/00002.jpg"],
+        ...
+        ["examples/fake/00001.jpg"],
+        ...
+    ]
+    的二维列表（每个 inner list 对应一次调用的所有输入，这里只有一张图）
+    """
+    exts = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
+    examples = []
+
+    # 让 real 在前、fake 在后，内层文件名按数字排序
+    for cls in ["real", "fake"]:
+        subdir = os.path.join(base_dir, cls)
+        if not os.path.isdir(subdir):
+            continue
+        for fname in sorted(os.listdir(subdir)):
+            if fname.lower().endswith(exts):
+                examples.append([os.path.join(subdir, fname)])
+    return examples
+
+class ForgeryDetector:
+    def __init__(self, arch='CLIP:ViT-L/14', ckpt_path='checkpoints/model_epoch_best.pth', device='cuda'):
+        """
+        初始化伪造检测器
+        
+        Args:
+            arch: 模型架构 (如 'res50', 'CLIP:ViT-B/32' 等)
+            ckpt_path: 预训练模型权重路径
+            device: 运行设备 ('cuda' 或 'cpu')
+        """
+        self.arch = arch
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+        
+        # 根据架构选择统计信息
+        stat_from = "imagenet" if arch.lower().startswith("imagenet") else "clip"
+        
+        # 数据预处理
+        self.transform = transforms.Compose([
+            transforms.Resize(256),  # 与验证代码保持一致
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=MEAN[stat_from], std=STD[stat_from]),
+        ])
+        
+        # 创建模型选项对象 (模拟TrainOptions)
+        self.opt = self._create_opt(arch)
+        
+        # 加载模型
+        self.model = self._load_model(ckpt_path)
+        
+    def _create_opt(self, arch):
+        """创建模型配置选项"""
+        class Options:
+            pass
+        
+        opt = Options()
+        opt.arch = arch
+        opt.head_type = "attention" if arch.startswith("CLIP:") else "fc"
+        opt.shuffle = True 
+        opt.shuffle_times = 1
+        opt.original_times = 1
+        opt.patch_size = [14]
+        opt.penultimate_feature = False
+        opt.patch_base = False
+        
+        return opt
+        
+    def _load_model(self, ckpt_path):
+        """加载预训练模型"""
+        print(f"正在加载模型架构: {self.arch}")
+        model = get_model(self.opt)
+        
+        # 加载权重
+        print(f"正在加载权重: {ckpt_path}")
+        state_dict = torch.load(ckpt_path, map_location=self.device)
+        
+        # 根据模型类型加载不同的权重
+        if self.opt.head_type == "fc":
+            model.fc.load_state_dict(state_dict)
+        elif self.opt.head_type == "attention":
+            model.attention_head.load_state_dict(state_dict)
+        else:
+            model.load_state_dict(state_dict)
+            
+        model.eval()
+        model.to(self.device)
+        print("模型加载完成!")
+        
+        return model
+    
+    def preprocess_image(self, image:Image.Image):
+        """
+        预处理单张图片
+        
+        Args:
+            image_path: 图片路径
+            
+        Returns:
+            处理后的tensor
+        """
+
+            
+        try:
+            image_tensor = self.transform(image).unsqueeze(0)  # 添加batch维度
+            return image_tensor.to(self.device)
+        except Exception as e:
+            raise ValueError(f"图片预处理失败: {e}")
+    
+    def predict(self, image: Image.Image, return_probabilities=True):
+        """
+        对单张图片进行伪造检测
+        
+        Args:
+            image: Image.Image
+            return_probabilities: 是否返回概率值，否则返回类别
+            
+        Returns:
+            如果return_probabilities=True: 返回 (real_prob, fake_prob)
+            如果return_probabilities=False: 返回 "real" 或 "fake"
+        """
+        # 预处理图片
+        image_tensor = self.preprocess_image(image)
+        
+        # 推理
+        with torch.no_grad():
+            output = self.model(image_tensor)
+            # 处理不同的输出格式
+            if output.shape[-1] == 2:
+                # 二分类输出
+                probs = torch.softmax(output, dim=1)[0]
+                
+                real_prob = probs[0].item()
+                fake_prob = probs[1].item()
+            else:
+                # 单值输出 (通常用sigmoid)
+                fake_prob = torch.sigmoid(output).item()
+                real_prob = 1 - fake_prob
+        
+        if return_probabilities:
+            return real_prob, fake_prob
+        else:
+            return "real" if real_prob > fake_prob else "fake"
+    
+    def batch_predict(self, image_paths, return_probabilities=True):
+        """
+        批量预测多张图片
+        
+        Args:
+            image_paths: 图片路径列表
+            return_probabilities: 是否返回概率值
+            
+        Returns:
+            预测结果列表
+        """
+        results = []
+        for image_path in image_paths:
+            try:
+                result = self.predict(image_path, return_probabilities)
+                results.append({
+                    'image_path': image_path,
+                    'result': result,
+                    'status': 'success'
+                })
+            except Exception as e:
+                results.append({
+                    'image_path': image_path,
+                    'result': None,
+                    'status': f'error: {e}'
+                })
+        return results
+
+
+model = ForgeryDetector()
+
+# ------- 1) （示例）假模型 -------
+def predict(img: Image.Image):
+    result = model.predict(img)
+
+    probs = np.array(result)
+    return {"Real Image": float(probs[0]), "Fake Image": float(probs[1])}
+
+EXAMPLES = collect_examples()      
+
+
+# ------- 2) UI -------
+with gr.Blocks() as demo:
+    gr.Markdown("## Real-vs-Fake Image Detector")
+
+    with gr.Row():
+        # —— 左列：图片 + 按钮 ——
+        with gr.Column(scale=1):
+            img_input = gr.Image(label="img", type="pil")
+
+            # ✨ 示例图片组件
+            gr.Examples(
+                examples=EXAMPLES,      # 刚才收集到的二维列表
+                inputs=img_input,       # 点击示例 -> 填到这个组件
+                cache_examples=False,   # 鉴于模型可能较大，关闭缓存以免占显存
+                label="Examples",
+            )
+
+            # 把按钮放到 **同一列** 下方
+            with gr.Row():
+                submit_btn = gr.Button("Submit", variant="primary")
+                clear_btn = gr.ClearButton(value="Clear", components=[img_input])
+
+        # —— 右列：结果条形图 ——
+        with gr.Column(scale=1):
+            label_output = gr.Label(label="output", num_top_classes=2)
+
+    # 交互
+    submit_btn.click(predict, inputs=img_input, outputs=label_output)
+    clear_btn.add([label_output])  # 同时清空右侧结果
+
+if __name__ == "__main__":
+    demo.launch()

checkpoints/model_epoch_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35273bf0539d3fc235ec24e5ce9bec3f17ff0572a48c0522b33fb89404213ff0
+size 12606764

models/__init__.py

@@ -0,0 +1,73 @@
+from .clip_models import CLIPModel, CLIPModelShuffleAttentionPenultimateLayer
+from .imagenet_models import ImagenetModel
+from .moco_models import MOCOModel_v3, MOCOModel_v3_ShuffleAttention, MOCOModel_v1, MOCOModel_v1_PatchesAttention, MOCOModel_v1_ShuffleAttention
+
+
+VALID_NAMES = [
+    'Imagenet:resnet18',
+    'Imagenet:resnet34',
+    'Imagenet:resnet50',
+    'Imagenet:resnet101',
+    'Imagenet:resnet152',
+    'Imagenet:vgg11',
+    'Imagenet:vgg19',
+    'Imagenet:swin-b',
+    'Imagenet:swin-s',
+    'Imagenet:swin-t',
+    'Imagenet:vit_b_16',
+    'Imagenet:vit_b_32',
+    'Imagenet:vit_l_16',
+    'Imagenet:vit_l_32',
+
+    'CLIP:RN50', 
+    'CLIP:RN101', 
+    'CLIP:RN50x4', 
+    'CLIP:RN50x16', 
+    'CLIP:RN50x64', 
+    'CLIP:ViT-B/32', 
+    'CLIP:ViT-B/16', 
+    'CLIP:ViT-L/14', 
+    'CLIP:ViT-L/14@336px',
+
+    'MOCO_v3:ViT',
+    'MOCO_v3:ViT_shuffle_attention',
+    'MOCO_v1:RN50',
+]
+
+
+
+
+
+def get_model(opt):
+    name = opt.arch
+    assert name in VALID_NAMES
+    if name.startswith("Imagenet:"):
+        return ImagenetModel(name[9:]) 
+    elif name.startswith("CLIP:"):
+        if opt.head_type == "fc":
+            if opt.penultimate_feature:
+                return CLIPModelPenultimateLayer(name[5:])
+            else:
+                return CLIPModel(name[5:])
+        elif opt.head_type == "attention":
+            # penultimate feature
+            if opt.shuffle:                
+                return CLIPModelShuffleAttentionPenultimateLayer(name[5:], shuffle_times=opt.shuffle_times, original_times=opt.original_times, patch_size=opt.patch_size)
+
+    elif name.startswith("MOCO_v1:"):
+        if opt.head_type == "fc":
+            return MOCOModel_v1()
+        elif opt.head_type == "attention":
+            if opt.patch_base:
+                return MOCOModel_v1_PatchesAttention()
+            else:
+                return MOCOModel_v1_ShuffleAttention()
+
+
+    elif name.startswith("MOCO_v3:"):
+        if opt.head_type == "attention":
+            return MOCOModel_v3_ShuffleAttention(opt)
+        else:
+            return MOCOModel_v3(opt)
+    else:
+        assert False 

models/clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *

models/clip/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

models/clip/clip.py

@@ -0,0 +1,237 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

models/clip/model.py

@@ -0,0 +1,452 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x[:1], key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x.squeeze(0)
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        out = {}
+        for idx, layer in enumerate(self.resblocks.children()):
+            x = layer(x)
+            out['layer'+str(idx)] = x[0] # shape:LND. choose cls token feature   
+        return out, x 
+
+        # return self.resblocks(x)  # This is the original code 
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        out, x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x[:, 0, :])
+
+
+        out['before_projection'] = x  
+
+        if self.proj is not None:
+            x = x @ self.proj
+        out['after_projection'] = x 
+
+        # Return both intermediate features and final clip feature 
+        # return out
+        
+        # This only returns CLIP features 
+        return x 
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith("transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

models/clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

models/clip_models.py

@@ -0,0 +1,115 @@
+from .clip import clip 
+from PIL import Image
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from models.transformer_attention import TransformerAttention
+import torchvision.transforms as transforms
+from .clip.model import VisionTransformer
+from .mlp import MLP
+
+CHANNELS = {
+    "RN50" : 1024,
+    "ViT-L/14" : 768,
+    "ViT-L/14-penultimate" : 1024
+}
+
+
+MEAN = {
+    "imagenet":[0.485, 0.456, 0.406],
+    "clip":[0.48145466, 0.4578275, 0.40821073]
+}
+
+STD = {
+    "imagenet":[0.229, 0.224, 0.225],
+    "clip":[0.26862954, 0.26130258, 0.27577711]
+}
+
+class CLIPModel(nn.Module):
+    """UFD"""
+    def __init__(self, name, num_classes=1):
+        super(CLIPModel, self).__init__()
+
+        self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class 
+        self.fc = nn.Linear( CHANNELS[name], num_classes )
+
+
+    def forward(self, x, return_feature=False):
+        features = self.model.encode_image(x) 
+        if return_feature:
+            return features
+        return self.fc(features)
+
+
+class CLIPModelPenultimateLayer(nn.Module):
+    def __init__(self, name, num_classes=1):
+        super(CLIPModelPenultimateLayer, self).__init__()
+
+        self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class 
+        self.register_hook()
+        self.fc = nn.Linear(CHANNELS[name+"-penultimate"], num_classes)
+
+    def register_hook(self):
+        
+        def hook(module, input, output):
+            self.features = torch.clone(output)
+        for name, module in self.model.visual.named_children():
+            if name == "ln_post":
+                module.register_forward_hook(hook)
+        return 
+
+    def forward(self, x):
+        self.model.encode_image(x) 
+        return self.fc(self.features)
+
+
+
+class CLIPModelShuffleAttentionPenultimateLayer(nn.Module):
+    def __init__(self, name, num_classes=1,shuffle_times=1, patch_size=32, original_times=1):
+        super(CLIPModelShuffleAttentionPenultimateLayer, self).__init__()
+        self.name = name
+        self.num_classes = num_classes
+        self.shuffle_times = shuffle_times
+        self.original_times = original_times
+        self.patch_size = patch_size
+        self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class
+        self.register_hook()
+        self.attention_head = TransformerAttention(CHANNELS[name+"-penultimate"], shuffle_times + original_times, last_dim=num_classes)
+        for name, param in self.model.named_parameters():
+            param.requires_grad = False
+
+    def register_hook(self):
+        
+        def hook(module, input, output):
+            self.features = torch.clone(output)
+        for name, module in self.model.visual.named_children():
+            if name == "ln_post":
+                module.register_forward_hook(hook)
+        return 
+    
+    def shuffle_patches(self, x, patch_size):
+        B, C, H, W = x.size()
+        # Unfold the input tensor to extract non-overlapping patches
+        patches = F.unfold(x, kernel_size=patch_size, stride=patch_size, dilation=1)
+        # Reshape the patches to (B, C, patch_H, patch_W, num_patches)
+        shuffled_patches = patches[:, :, torch.randperm(patches.size(-1))]
+        # Fold the shuffled patches back into images
+        shuffled_images = F.fold(shuffled_patches, output_size=(H, W), kernel_size=patch_size, stride=patch_size)
+        return shuffled_images
+
+
+    def forward(self, x, return_feature=False):
+        features = []
+        with torch.no_grad():
+            for i in range(self.shuffle_times):
+                self.model.encode_image(self.shuffle_patches(x, patch_size=self.patch_size[0]))
+                features.append(self.features)
+
+            self.model.encode_image(x)
+            for i in range(self.original_times):
+                features.append(self.features.clone())
+        features = self.attention_head(torch.stack(features, dim=-2))
+
+        return features
+
+

models/delete/vgg.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+from typing import Union, List, Dict, Any, cast
+import torchvision
+import torch.nn.functional as F
+
+
+
+
+
+class VGG(torch.nn.Module):
+    def __init__(self, arch_type, pretrained, progress):
+        super().__init__()
+        
+        self.layer1 = torch.nn.Sequential()
+        self.layer2 = torch.nn.Sequential()
+        self.layer3 = torch.nn.Sequential()
+        self.layer4 = torch.nn.Sequential()
+        self.layer5 = torch.nn.Sequential()
+
+        if arch_type == 'vgg11':
+            official_vgg = torchvision.models.vgg11(pretrained=pretrained, progress=progress)
+            blocks = [  [0,2], [2,5], [5,10], [10,15], [15,20] ]
+            last_idx = 20
+        elif arch_type == 'vgg19':
+            official_vgg = torchvision.models.vgg19(pretrained=pretrained, progress=progress)
+            blocks = [  [0,4], [4,9], [9,18], [18,27], [27,36] ]
+            last_idx = 36
+        else:
+            raise NotImplementedError
+        
+        
+        for x in range( *blocks[0] ):
+            self.layer1.add_module(str(x), official_vgg.features[x])
+        for x in range( *blocks[1] ):
+            self.layer2.add_module(str(x), official_vgg.features[x])
+        for x in range( *blocks[2] ):
+            self.layer3.add_module(str(x), official_vgg.features[x])
+        for x in range( *blocks[3] ):
+            self.layer4.add_module(str(x), official_vgg.features[x])
+        for x in range( *blocks[4] ):
+            self.layer5.add_module(str(x), official_vgg.features[x])
+            
+        self.max_pool = official_vgg.features[last_idx]
+        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+        
+        self.fc1 = official_vgg.classifier[0]
+        self.fc2 = official_vgg.classifier[3]
+        self.fc3 = official_vgg.classifier[6]
+        self.dropout = nn.Dropout()
+        
+        
+    def forward(self, x):
+        out = {}
+        
+        x = self.layer1(x)
+        out['f0'] = x
+        
+        x = self.layer2(x)
+        out['f1'] = x
+        
+        x = self.layer3(x)
+        out['f2'] = x
+        
+        x = self.layer4(x)
+        out['f3'] = x
+        
+        x = self.layer5(x)
+        out['f4'] = x
+        
+        x = self.max_pool(x)
+        x = self.avgpool(x)
+        x = x.view(-1,512*7*7) 
+        
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = self.dropout(x) 
+        x = self.fc2(x)
+        x = F.relu(x)
+        out['penultimate'] = x 
+        x = self.dropout(x) 
+        x = self.fc3(x)
+        out['logits'] = x 
+
+        return out
+
+
+
+
+
+
+
+
+
+
+def vgg11(pretrained=False, progress=True):
+    r"""VGG 11-layer model (configuration "A") from
+    `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return VGG('vgg11', pretrained, progress)
+
+
+
+def vgg19(pretrained=False, progress=True):
+    r"""VGG 19-layer model (configuration "E")
+    `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return VGG('vgg19', pretrained, progress)
+
+
+
+

models/imagenet_models.py

@@ -0,0 +1,40 @@
+from .resnet import resnet18, resnet34, resnet50, resnet101, resnet152
+from .vision_transformer import vit_b_16, vit_b_32, vit_l_16, vit_l_32
+
+from torchvision import transforms
+from PIL import Image
+import torch 
+import torch.nn as nn 
+
+
+model_dict = {
+    'resnet18': resnet18,
+    'resnet34': resnet34,
+    'resnet50': resnet50,
+    'resnet101': resnet101,
+    'resnet152': resnet152,
+    'vit_b_16': vit_b_16,
+    'vit_b_32': vit_b_32,
+    'vit_l_16': vit_l_16,
+    'vit_l_32': vit_l_32
+}
+
+
+CHANNELS = {
+    "resnet50" : 2048,
+    "vit_b_16" : 768,
+}
+
+
+
+class ImagenetModel(nn.Module):
+    def __init__(self, name, num_classes=1):
+        super(ImagenetModel, self).__init__()
+
+        self.model = model_dict[name](pretrained=True)
+        self.fc = nn.Linear(CHANNELS[name], num_classes) #manually define a fc layer here
+        
+
+    def forward(self, x):
+        feature = self.model(x)["penultimate"]
+        return self.fc(feature)

models/mlp.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MLP(nn.Module):
+    def __init__(self, input_size, h_sizes, output_size):
+        super(MLP, self).__init__()
+        self.input = nn.Linear(input_size, h_sizes[0])
+        self.hidden = nn.ModuleList([nn.Linear(h_sizes[i], h_sizes[i+1]) for i in range(len(h_sizes) - 1)])
+        self.output = nn.Linear(h_sizes[-1], output_size)
+
+    def forward(self, x):
+        x = F.relu(self.input(x))
+        for layer in self.hidden:
+            x = F.relu(layer(x))
+        x = self.output(x)
+        return x
+

models/moco_models.py

@@ -0,0 +1,364 @@
+import torch.nn as nn
+import torch
+import torchvision.transforms as transforms
+import os
+from models import vits
+from models.transformer_attention import TransformerAttention
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.optim
+import torch.utils.data
+import torch.utils.data.distributed
+import torchvision.transforms as transforms
+import torchvision.models as torchvision_models
+import torch.nn.functional as F
+import models.moco_v1.builder as builder
+import torchvision.models as models
+
+
+CHANNELS = {
+    "MOCO" : 128,
+    "ViT-L/14" : 768
+}
+
+MEAN = {
+    "imagenet":[0.485, 0.456, 0.406],
+    "clip":[0.48145466, 0.4578275, 0.40821073]
+}
+
+STD = {
+    "imagenet":[0.229, 0.224, 0.225],
+    "clip":[0.26862954, 0.26130258, 0.27577711]
+}
+
+
+
+class MOCOModel_v3(nn.Module):
+    def __init__(self, opt, num_classes=1):
+        super(MOCOModel_v3, self).__init__()
+        self.opt = opt
+        model_root = "path to your moco pretrained weight"
+        self.model = self.create_model(model_root) # self.preprecess will not be used during training, which is handled in Dataset class 
+        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.preprocess = transforms.Compose([
+                transforms.Resize([224, 224]),
+                transforms.ToTensor(),
+                normalize,
+            ])
+ 
+    def create_model(self, model_root):
+        opt = self.opt
+        opt.pretrained = model_root
+        dim_in = None
+        # create model
+        print("=> creating model '{}'".format(opt.arch))
+        if opt.arch.lower().find('vit') != -1:
+            model = vits.__dict__["vit_base"]()
+            linear_keyword = 'head'
+        else:
+            model = torchvision_models.__dict__[opt.arch]()
+            linear_keyword = 'fc'
+
+        # freeze all layers and remove the last fc
+        remove_name = None
+        for name, param in model.named_parameters():
+            if name not in ['%s.weight' % linear_keyword, '%s.bias' % linear_keyword]:
+                param.requires_grad = False
+        model.head = nn.Linear(model.head.weight.shape[-1], 1)
+        self.fc = model.head
+
+        # init the fc layer
+        getattr(model, linear_keyword).weight.data.normal_(mean=0.0, std=0.01)
+        getattr(model, linear_keyword).bias.data.zero_()
+
+        # load from pre-trained, before DistributedDataParallel constructor
+        if opt.pretrained:
+            if os.path.isfile(opt.pretrained):
+                print("=> loading checkpoint '{}'".format(opt.pretrained))
+                checkpoint = torch.load(opt.pretrained, map_location="cpu")
+
+                # rename moco pre-trained keys
+                state_dict = checkpoint['state_dict']
+                for k in list(state_dict.keys()):
+                    # retain only base_encoder up to before the embedding layer
+                    if k.startswith('module.base_encoder') and not k.startswith('module.base_encoder.%s' % linear_keyword):
+                        # remove prefix
+                        state_dict[k[len("module.base_encoder."):]] = state_dict[k]
+                    # delete renamed or unused k
+                    del state_dict[k]
+
+                opt.start_epoch = 0
+                msg = model.load_state_dict(state_dict, strict=False)
+                assert set(msg.missing_keys) == {"%s.weight" % linear_keyword, "%s.bias" % linear_keyword}
+
+                print("=> loaded pre-trained model '{}'".format(opt.pretrained))
+            else:
+                print("=> no checkpoint found at '{}'".format(opt.pretrained))
+
+
+
+        return model
+
+
+    def forward(self, x, return_feature=False):
+        features = self.model(x) 
+        return features
+
+
+
+class MOCOModel_v3_ShuffleAttention(nn.Module):
+    def __init__(self, opt, num_classes=1, shuffle_times=3):
+        self.shuffle_times = shuffle_times
+        super(MOCOModel_v3_ShuffleAttention, self).__init__()
+        self.opt = opt
+        self.attention_head = TransformerAttention(768, shuffle_times, 1)
+        model_root = "path to your moco pretrained weight"
+        self.model = self.create_model(model_root) # self.preprecess will not be used during training, which is handled in Dataset class 
+        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        self.preprocess = transforms.Compose([
+                transforms.Resize([224, 224]),
+                transforms.ToTensor(),
+                normalize,
+            ])
+ 
+    def create_model(self, model_root):
+        opt = self.opt
+        opt.pretrained = model_root
+        dim_in = None
+        # create model
+        print("=> creating model '{}'".format(opt.arch))
+        if opt.arch.lower().find('vit') != -1:
+            model = vits.__dict__["vit_base"]()
+            linear_keyword = 'head'
+        else:
+            model = torchvision_models.__dict__[opt.arch]()
+            linear_keyword = 'fc'
+
+        # init the fc layer
+        # getattr(model, linear_keyword).weight.data.one_()
+        # getattr(model, linear_keyword).bias.data.zero_()
+        for name, param in model.named_parameters():
+            # if name not in ['%s.weight' % linear_keyword, '%s.bias' % linear_keyword]:
+            param.requires_grad = False
+
+
+        # load from pre-trained, before DistributedDataParallel constructor
+        if opt.pretrained:
+            if os.path.isfile(opt.pretrained):
+                print("=> loading checkpoint '{}'".format(opt.pretrained))
+                checkpoint = torch.load(opt.pretrained, map_location="cpu")
+
+                # rename moco pre-trained keys
+                state_dict = checkpoint['state_dict']
+                for k in list(state_dict.keys()):
+                    # retain only base_encoder up to before the embedding layer
+                    if k.startswith('module.base_encoder') and not k.startswith('module.base_encoder.%s' % linear_keyword):
+                        # remove prefix
+                        state_dict[k[len("module.base_encoder."):]] = state_dict[k]
+                    # delete renamed or unused k
+                    del state_dict[k]
+
+                opt.start_epoch = 0
+                msg = model.load_state_dict(state_dict, strict=False)
+                assert set(msg.missing_keys) == {"%s.weight" % linear_keyword, "%s.bias" % linear_keyword}
+
+                print("=> loaded pre-trained model '{}'".format(opt.pretrained))
+            else:
+                print("=> no checkpoint found at '{}'".format(opt.pretrained))
+
+        # remove last layer
+        model.head = nn.Identity()
+
+
+        return model
+
+    def shuffle_patches(self, x, patch_size):
+        B, C, H, W = x.size()
+        # Unfold the input tensor to extract non-overlapping patches
+        patches = F.unfold(x, kernel_size=patch_size, stride=patch_size, dilation=1)
+        # Reshape the patches to (B, C, patch_H, patch_W, num_patches)
+        shuffled_patches = patches[:, :, torch.randperm(patches.size(-1))]
+        # Fold the shuffled patches back into images
+        shuffled_images = F.fold(shuffled_patches, output_size=(H, W), kernel_size=patch_size, stride=patch_size)
+        return shuffled_images
+
+
+    def forward(self, x, return_feature=False):
+        features = []
+        for i in range(self.shuffle_times):
+            features.append(self.model(self.shuffle_patches(x, patch_size=32)))
+        features = self.attention_head(torch.stack(features, dim=-2))
+        return features
+
+
+
+
+
+class MOCOModel_v1(nn.Module):
+    def __init__(self, num_classes=1):
+        super(MOCOModel_v1, self).__init__()
+        model_root = "path to your moco pretrained weight"
+        self.model = self.create_model(model_root) # self.preprecess will not be used during training, which is handled in Dataset class 
+        self.fc = nn.Linear( CHANNELS["MOCO"], num_classes)
+
+
+    def create_model(self, model_root):
+        print("=> creating model '{}'".format("resnet50"))
+        model = builder.MoCo(
+            models.__dict__["resnet50"],
+            128,
+            65536,
+            0.999,
+            0.07,
+            False,
+        )
+
+        model = model.encoder_q
+        pretrained_dict = torch.load(model_root, map_location="cpu")["state_dict"]
+        model_dict = model.state_dict()
+
+        # Create a new dictionary that maps the keys in the pretrained model to the keys in the built model
+        new_dict = {}
+        for k, v in pretrained_dict.items():
+            if 'module.encoder_q' in k:
+                k = k.replace('module.encoder_q.', '')
+            new_dict[k] = v
+
+        # Load the pretrained weights into the built model
+        model_dict.update(new_dict)
+        model.load_state_dict(model_dict)
+
+
+        return model
+
+
+    def forward(self, x, return_feature=False):
+        features = self.model(x) 
+        if return_feature:
+            return features
+        return self.fc(features)
+
+
+
+
+
+class MOCOModel_v1_ShuffleAttention(nn.Module):
+    def __init__(self, num_classes=1,shuffle_times=3):
+        super(MOCOModel_v1_ShuffleAttention, self).__init__()
+        model_root = "path to your moco pretrained weight"
+        self.model = self.create_model(model_root) # self.preprecess will not be used during training, which is handled in Dataset class 
+        self.shuffle_times = shuffle_times
+        self.attention_head = TransformerAttention(CHANNELS["MOCO"], shuffle_times + 1, 1)
+        for name, param in self.model.named_parameters():
+            param.requires_grad = False
+
+    def create_model(self, model_root):
+        print("=> creating model '{}'".format("resnet50"))
+        model = builder.MoCo(
+            models.__dict__["resnet50"],
+            128,
+            65536,
+            0.999,
+            0.07,
+            False,
+        )
+
+        model = model.encoder_q
+        pretrained_dict = torch.load(model_root, map_location="cpu")["state_dict"]
+        model_dict = model.state_dict()
+
+        # Create a new dictionary that maps the keys in the pretrained model to the keys in the built model
+        new_dict = {}
+        for k, v in pretrained_dict.items():
+            if 'module.encoder_q' in k:
+                k = k.replace('module.encoder_q.', '')
+            new_dict[k] = v
+
+        # Load the pretrained weights into the built model
+        model_dict.update(new_dict)
+        model.load_state_dict(model_dict)
+
+
+        return model
+
+
+    def shuffle_patches(self, x, patch_size):
+        B, C, H, W = x.size()
+        # Unfold the input tensor to extract non-overlapping patches
+        patches = F.unfold(x, kernel_size=patch_size, stride=patch_size, dilation=1)
+        # Reshape the patches to (B, C, patch_H, patch_W, num_patches)
+        shuffled_patches = patches[:, :, torch.randperm(patches.size(-1))]
+        # Fold the shuffled patches back into images
+        shuffled_images = F.fold(shuffled_patches, output_size=(H, W), kernel_size=patch_size, stride=patch_size)
+        return shuffled_images
+
+
+    def forward(self, x, return_feature=False):
+        features = []
+        for i in range(self.shuffle_times):
+            features.append(self.model(self.shuffle_patches(x, patch_size=32)))
+        features.append(self.model(x))
+        features = self.attention_head(torch.stack(features, dim=-2))
+        return features
+
+
+
+class MOCOModel_v1_PatchesAttention(nn.Module):
+    def __init__(self, num_classes=1, divide_nums=3):
+        super(MOCOModel_v1_PatchesAttention, self).__init__()
+        self.divide_nums = divide_nums
+        model_root = "path to your moco pretrained weight"
+        self.model = self.create_model(model_root) # self.preprecess will not be used during training, which is handled in Dataset class 
+        self.normal = transforms.Normalize( mean=MEAN["imagenet"], std=STD["imagenet"] )
+        self.attention_head = TransformerAttention(CHANNELS["MOCO"], divide_nums**2 + 1, 1)
+        for name, param in self.model.named_parameters():
+            param.requires_grad = False
+
+    def create_model(self, model_root):
+        print("=> creating model '{}'".format("resnet50"))
+        model = builder.MoCo(
+            models.__dict__["resnet50"],
+            128,
+            65536,
+            0.999,
+            0.07,
+            False,
+        )
+
+        model = model.encoder_q
+        pretrained_dict = torch.load(model_root, map_location="cpu")["state_dict"]
+        model_dict = model.state_dict()
+
+        # Create a new dictionary that maps the keys in the pretrained model to the keys in the built model
+        new_dict = {}
+        for k, v in pretrained_dict.items():
+            if 'module.encoder_q' in k:
+                k = k.replace('module.encoder_q.', '')
+            new_dict[k] = v
+
+        # Load the pretrained weights into the built model
+        model_dict.update(new_dict)
+        model.load_state_dict(model_dict)
+
+
+        return model
+
+    def divide(self, x):
+        patch_size = 224 // self.divide_nums
+        patches = []
+        for i in range(self.divide_nums):
+            for j in range(self.divide_nums):
+                patches.append(transforms.Resize([224, 224], antialias=True)(x[:, :, i*patch_size:(i+1)*patch_size, j*patch_size:(j+1)*patch_size]))
+        return patches
+
+
+    def forward(self, x, return_feature=False):
+        features = []
+        patches = self.divide(x)
+        for patch in patches:
+            features.append(self.model(self.normal(patch)))
+        features.append(self.model(self.normal(x)))
+        features = self.attention_head(torch.stack(features, dim=-2))
+        return features
+    

models/moco_v1/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.

models/moco_v1/builder.py

@@ -0,0 +1,380 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+from copy import deepcopy
+
+
+class MoCo(nn.Module):
+    """
+    Build a MoCo model with: a query encoder, a key encoder, and a queue
+    https://arxiv.org/abs/1911.05722
+    """
+
+    def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False):
+        """
+        dim: feature dimension (default: 128)
+        K: queue size; number of negative keys (default: 65536)
+        m: moco momentum of updating key encoder (default: 0.999)
+        T: softmax temperature (default: 0.07)
+        """
+        super(MoCo, self).__init__()
+
+        self.K = K
+        self.m = m
+        self.T = T
+
+        # create the encoders
+        # num_classes is the output fc dimension
+        self.encoder_q = base_encoder(num_classes=dim)
+        self.encoder_k = base_encoder(num_classes=dim)
+
+        if mlp:  # hack: brute-force replacement
+            dim_mlp = self.encoder_q.fc.weight.shape[1]
+            self.encoder_q.fc = nn.Sequential(
+                nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc
+            )
+            self.encoder_k.fc = nn.Sequential(
+                nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc
+            )
+
+        for param_q, param_k in zip(
+            self.encoder_q.parameters(), self.encoder_k.parameters()
+        ):
+            param_k.data.copy_(param_q.data)  # initialize
+            param_k.requires_grad = False  # not update by gradient
+
+        # create the queue
+        self.register_buffer("queue", torch.randn(dim, K))
+        self.queue = nn.functional.normalize(self.queue, dim=0)
+
+        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
+
+    @torch.no_grad()
+    def _momentum_update_key_encoder(self):
+        """
+        Momentum update of the key encoder
+        """
+        for param_q, param_k in zip(
+            self.encoder_q.parameters(), self.encoder_k.parameters()
+        ):
+            param_k.data = param_k.data * self.m + param_q.data * (1.0 - self.m)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, keys):
+        # gather keys before updating queue
+        keys = concat_all_gather(keys)
+
+        batch_size = keys.shape[0]
+
+        ptr = int(self.queue_ptr)
+        assert self.K % batch_size == 0  # for simplicity
+
+        # replace the keys at ptr (dequeue and enqueue)
+        self.queue[:, ptr : ptr + batch_size] = keys.T
+        ptr = (ptr + batch_size) % self.K  # move pointer
+
+        self.queue_ptr[0] = ptr
+
+    @torch.no_grad()
+    def _batch_shuffle_ddp(self, x):
+        """
+        Batch shuffle, for making use of BatchNorm.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # random shuffle index
+        idx_shuffle = torch.randperm(batch_size_all).cuda()
+
+        # broadcast to all gpus
+        torch.distributed.broadcast(idx_shuffle, src=0)
+
+        # index for restoring
+        idx_unshuffle = torch.argsort(idx_shuffle)
+
+        # shuffled index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this], idx_unshuffle
+
+    @torch.no_grad()
+    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
+        """
+        Undo batch shuffle.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # restored index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this]
+
+    def forward(self, im_q, im_k):
+        """
+        Input:
+            im_q: a batch of query images
+            im_k: a batch of key images
+        Output:
+            logits, targets
+        """
+
+        # compute query features
+        q = self.encoder_q(im_q)  # queries: NxC
+        q = nn.functional.normalize(q, dim=1)
+
+        # compute key features
+        with torch.no_grad():  # no gradient to keys
+            self._momentum_update_key_encoder()  # update the key encoder
+
+            # shuffle for making use of BN
+            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
+
+            k = self.encoder_k(im_k)  # keys: NxC
+            k = nn.functional.normalize(k, dim=1)
+
+            # undo shuffle
+            k = self._batch_unshuffle_ddp(k, idx_unshuffle)
+
+        # compute logits
+        # Einstein sum is more intuitive
+        # positive logits: Nx1
+        l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1)
+        # negative logits: NxK
+        l_neg = torch.einsum("nc,ck->nk", [q, self.queue.clone().detach()])
+
+        # logits: Nx(1+K)
+        logits = torch.cat([l_pos, l_neg], dim=1)
+
+        # apply temperature
+        logits /= self.T
+
+        # labels: positive key indicators
+        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
+
+        # dequeue and enqueue
+        self._dequeue_and_enqueue(k)
+
+        return logits, labels
+
+
+class MoCoMAE(nn.Module):
+
+    def __init__(self, base_encoder, dim=1024, K=65536, m=0.999, T=0.07, norm=False):
+        super().__init__()
+
+        self.K = K
+        self.m = m
+        self.T = T
+        self.norm = norm
+
+        self.patch_embed = deepcopy(base_encoder.patch_embed)
+        self.pos_embed = deepcopy(base_encoder.pos_embed)
+        self.encoder_q = deepcopy(base_encoder)
+        self.encoder_k = deepcopy(base_encoder)
+
+        for param_k in self.encoder_k.parameters():
+            param_k.requires_grad = False  # not update by gradient
+        for param_q in self.encoder_q.parameters():
+            param_q.requires_grad = False
+        for param_q in self.encoder_q.blocks.parameters():
+            param_q.requires_grad = True
+        self.encoder_q.cls_token.requires_grad = True
+
+
+        # create the queue
+        self.register_buffer("queue", torch.randn(dim, K))
+        self.queue = nn.functional.normalize(self.queue, dim=0)
+
+        self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
+
+    @torch.no_grad()
+    def _momentum_update_key_encoder(self):
+        """
+        Momentum update of the key encoder
+        """
+        for param_q, param_k in zip(
+            self.encoder_q.parameters(), self.encoder_k.parameters()
+        ):
+            param_k.data = param_k.data * self.m + param_q.data * (1.0 - self.m)
+
+    @torch.no_grad()
+    def _dequeue_and_enqueue(self, keys):
+        # gather keys before updating queue
+        keys = concat_all_gather(keys)
+
+        batch_size = keys.shape[0]
+
+        ptr = int(self.queue_ptr)
+        assert self.K % batch_size == 0  # for simplicity
+
+        # replace the keys at ptr (dequeue and enqueue)
+        self.queue[:, ptr : ptr + batch_size] = keys.T
+        ptr = (ptr + batch_size) % self.K  # move pointer
+
+        self.queue_ptr[0] = ptr
+
+    @torch.no_grad()
+    def _batch_shuffle_ddp(self, x):
+        """
+        Batch shuffle, for making use of BatchNorm.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # random shuffle index
+        idx_shuffle = torch.randperm(batch_size_all).cuda()
+
+        # broadcast to all gpus
+        torch.distributed.broadcast(idx_shuffle, src=0)
+
+        # index for restoring
+        idx_unshuffle = torch.argsort(idx_shuffle)
+
+        # shuffled index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this], idx_unshuffle
+
+    @torch.no_grad()
+    def _batch_unshuffle_ddp(self, x, idx_unshuffle):
+        """
+        Undo batch shuffle.
+        *** Only support DistributedDataParallel (DDP) model. ***
+        """
+        # gather from all gpus
+        batch_size_this = x.shape[0]
+        x_gather = concat_all_gather(x)
+        batch_size_all = x_gather.shape[0]
+
+        num_gpus = batch_size_all // batch_size_this
+
+        # restored index for this gpu
+        gpu_idx = torch.distributed.get_rank()
+        idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]
+
+        return x_gather[idx_this]
+
+    def augment(self, x, mask_ratio=0.75):
+        # embed patches
+        x = self.patch_embed(x)
+        # add pos embed w/o cls token
+        x = x + self.pos_embed[:, 1:, :]
+        # masking: length -> length * mask_ratio
+        x = random_masking(x, mask_ratio)
+        return x
+    
+    def forward(self, im_q, im_k):
+        """
+        Input:
+            im_q: a batch of query images
+            im_k: a batch of key images
+        Output:
+            logits, targets
+        """
+
+        # compute query features
+        q = self.encoder_q(im_q)  # queries: NxC
+        if self.norm:
+            q = nn.functional.normalize(q, dim=1)
+
+        # compute key features
+        with torch.no_grad():  # no gradient to keys
+            self._momentum_update_key_encoder()  # update the key encoder
+
+            # shuffle for making use of BN
+            im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)
+
+            k = self.encoder_k(im_k)  # keys: NxC
+            if self.norm:
+                k = nn.functional.normalize(k, dim=1)
+
+            # undo shuffle
+            k = self._batch_unshuffle_ddp(k, idx_unshuffle)
+
+        # compute logits
+        # Einstein sum is more intuitive
+        # positive logits: Nx1
+        l_pos = torch.einsum("nc,nc->n", [q, k]).unsqueeze(-1)
+        # negative logits: NxK
+        l_neg = torch.einsum("nc,ck->nk", [q, self.queue.clone().detach()])
+
+        # logits: Nx(1+K)
+        logits = torch.cat([l_pos, l_neg], dim=1)
+
+        # apply temperature
+        logits /= self.T
+
+        # labels: positive key indicators
+        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
+
+        # dequeue and enqueue
+        self._dequeue_and_enqueue(k)
+
+        return logits, labels
+    
+
+# utils
+@torch.no_grad()
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [
+        torch.ones_like(tensor) for _ in range(torch.distributed.get_world_size())
+    ]
+    torch.distributed.all_gather(tensors_gather, tensor.contiguous(), async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output
+
+
+def random_masking(x, mask_ratio):
+    """
+    Perform per-sample random masking by per-sample shuffling.
+    Per-sample shuffling is done by argsort random noise.
+    x: [N, L, D], sequence
+    """
+    N, L, D = x.shape  # batch, length, dim
+    len_keep = int(L * (1 - mask_ratio))
+    
+    noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+    
+    # sort noise for each sample
+    ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+    ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+    # keep the first subset
+    ids_keep = ids_shuffle[:, :len_keep]
+    x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+    # generate the binary mask: 0 is keep, 1 is remove
+    mask = torch.ones([N, L], device=x.device)
+    mask[:, :len_keep] = 0
+    # unshuffle to get the binary mask
+    mask = torch.gather(mask, dim=1, index=ids_restore)
+
+    return x_masked

models/moco_v1/loader.py

@@ -0,0 +1,32 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+
+from PIL import ImageFilter
+
+
+class TwoCropsTransform:
+    """Take two random crops of one image as the query and key."""
+
+    def __init__(self, base_transform):
+        self.base_transform = base_transform
+
+    def __call__(self, x):
+        q = self.base_transform(x)
+        k = self.base_transform(x)
+        return [q, k]
+
+
+class GaussianBlur:
+    """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""
+
+    def __init__(self, sigma=[0.1, 2.0]):
+        self.sigma = sigma
+
+    def __call__(self, x):
+        sigma = random.uniform(self.sigma[0], self.sigma[1])
+        x = x.filter(ImageFilter.GaussianBlur(radius=sigma))
+        return x

models/resnet.py

@@ -0,0 +1,337 @@
+import torch
+from torch import Tensor
+import torch.nn as nn
+from typing import Type, Any, Callable, Union, List, Optional
+
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+
+
+def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion: int = 4
+
+    def __init__(
+        self,
+        inplanes: int,
+        planes: int,
+        stride: int = 1,
+        downsample: Optional[nn.Module] = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(
+        self,
+        block: Type[Union[BasicBlock, Bottleneck]],
+        layers: List[int],
+        num_classes: int = 1000,
+        zero_init_residual: bool = False,
+        groups: int = 1,
+        width_per_group: int = 64,
+        replace_stride_with_dilation: Optional[List[bool]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
+
+    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,
+                    stride: int = 1, dilate: bool = False) -> nn.Sequential:
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        # The comment resolution is based on input size is 224*224 imagenet
+        out = {}
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        out['f0'] = x # N*64*56*56 
+
+        x = self.layer1(x)
+        out['f1'] = x # N*64*56*56
+
+        x = self.layer2(x)
+        out['f2'] = x # N*128*28*28
+
+        x = self.layer3(x)
+        out['f3'] = x # N*256*14*14
+        
+        x = self.layer4(x)
+        out['f4'] = x # N*512*7*7
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        out['penultimate'] = x # N*512
+
+        x = self.fc(x)
+        out['logits'] = x # N*1000
+
+        # return all features 
+        return out
+
+        # return final classification result 
+        # return x
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+
+def _resnet(
+    arch: str,
+    block: Type[Union[BasicBlock, Bottleneck]],
+    layers: List[int],
+    pretrained: bool,
+    progress: bool,
+    **kwargs: Any
+) -> ResNet:
+    model = ResNet(block, layers, **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, **kwargs)
+
+
+def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, **kwargs)
+
+
+def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, **kwargs)
+
+
+def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, **kwargs)
+

models/transformer_attention.py

@@ -0,0 +1,217 @@
+import torch
+import torch.nn as nn
+
+
+
+
+
+class TransformerAttention(nn.Module):
+    def __init__(self, input_dim, output_dim, last_dim=1):
+        super(TransformerAttention, self).__init__()
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim*output_dim, last_dim)
+        # self.fc = nn.Linear(input_dim, output_dim)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        # x = x.unsqueeze(dim=1)
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1)))
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+        output = output.view([output.shape[0], -1])
+        output = self.fc(output)
+        return output
+    
+
+
+
+
+
+class TransformerMultiHeadAttention(nn.Module):
+    def __init__(self, input_dim, output_dim, head_num, predict_dim=1):
+        super(TransformerMultiHeadAttention, self).__init__()
+        self.head_num = head_num
+        self.query = nn.Linear(input_dim, input_dim * head_num)
+        self.key = nn.Linear(input_dim, input_dim * head_num)
+        self.value = nn.Linear(input_dim, input_dim * head_num)
+        self.fc = nn.Linear(input_dim * head_num * output_dim, predict_dim)
+        self.softmax = nn.Softmax(dim=1)
+
+    def split_heads(self, tensor):
+        # Split the last dimension into (head_num, new_last_dim)
+        tensor = tensor.view(tensor.size(0), tensor.size(1), self.head_num, tensor.size(-1) // self.head_num)
+        # Transpose the result to (batch, head_num, new_last_dim, -1)
+        return tensor.transpose(1, 2)
+    
+
+    def scaled_dot_product_attention(self, q, k, v):
+        d_k = q.size(-1)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
+        attention = self.softmax(scores)
+        return torch.matmul(attention, v)
+
+    def combine_heads(self, tensor):
+        # Transpose and reshape the tensor to (batch, -1, head_num * new_last_dim)
+        return tensor.transpose(1, 2).contiguous().view(tensor.size(0), -1)
+    
+    def forward(self, x):
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+        q, k, v = [self.split_heads(tensor) for tensor in (q, k, v)]
+        attention = self.scaled_dot_product_attention(q, k, v)
+        output = self.combine_heads(attention)
+        output = self.fc(output)
+        return output
+
+
+
+class TransformerAttentionwithClassifierToken(nn.Module):
+    def __init__(self, input_dim, middle_dim, output_dim):
+        super(TransformerAttentionwithClassifierToken, self).__init__()
+        self.input_dim = input_dim
+        self.classifier_token = nn.Linear(input_dim, 1)
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim, output_dim)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        classifier_token = torch.ones((x.shape[0], 1, self.input_dim)).to(x.device)*self.classifier_token.weight.data.squeeze(dim=-1)
+        x = torch.cat([classifier_token.to(x.device), x], dim=-2)
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1)))
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+        # using only the classifier token
+        output = self.fc(output[:, 0, :])
+        return output
+
+
+class TransformerCrossAttention(nn.Module):
+    def __init__(self, input_dim, middle_dim, output_dim):
+        super(TransformerAttention, self).__init__()
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim*middle_dim, output_dim)
+        # self.fc = nn.Linear(input_dim, output_dim)
+        self.softmax = nn.Softmax(dim=1)
+
+    def forward(self, x):
+        # x = x.unsqueeze(dim=1)
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1)))
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+        output = output.view([output.shape[0], -1])
+        output = self.fc(output)
+        return output
+
+
+
+
+
+class TransformerAttentionwithPisition(nn.Module):
+    def __init__(self, input_dim, output_dim, last_dim=1, token_num=2):
+        super(TransformerAttentionwithPisition, self).__init__()
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim * token_num, last_dim)
+        self.softmax = nn.Softmax(dim=1)
+        self.position_embeddings = nn.Parameter(torch.zeros([token_num, input_dim]))
+
+    def forward(self, x):
+        seq_len = x.size(1)
+        x_with_position = x + self.position_embeddings
+
+        q = self.query(x_with_position)
+        k = self.key(x_with_position)
+        v = self.value(x_with_position)
+
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1), dtype=torch.float32))
+
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+
+        output = output.view([output.shape[0], -1])
+        output = self.fc(output)
+        return output
+
+
+
+class TransformerCrossAttentionwithPisition(nn.Module):
+    def __init__(self, input_dim, output_dim, last_dim=1, token_num=2):
+        super(TransformerCrossAttentionwithPisition, self).__init__()
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim * output_dim // 2, last_dim)
+        self.softmax = nn.Softmax(dim=1)
+        self.position_embeddings = nn.Parameter(torch.zeros([token_num, input_dim]))
+
+    def forward(self, x):
+        # first half to be q and second half to be k, v
+        seq_len = x.size(1)
+        x_with_position = x + self.position_embeddings
+
+        q = self.query(x_with_position[:, :x_with_position.shape[1]//2, :])
+        k = self.key(x_with_position[:, x_with_position.shape[1]//2:, :])
+        v = self.value(x_with_position[:, x_with_position.shape[1]//2:, :])
+
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1), dtype=torch.float32))
+
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+
+        output = output.view([output.shape[0], -1])
+        output = self.fc(output)
+        return output
+
+
+class TransformerAttentionwithCatPe(nn.Module):
+    def __init__(self, input_dim, output_dim, last_dim=1, token_num=2):
+        super(TransformerAttentionwithCatPe, self).__init__()
+        input_dim = input_dim * 2
+        self.query = nn.Linear(input_dim, input_dim)
+        self.key = nn.Linear(input_dim, input_dim)
+        self.value = nn.Linear(input_dim, input_dim)
+        self.fc = nn.Linear(input_dim * output_dim, last_dim)
+        self.softmax = nn.Softmax(dim=1)
+        self.position_embeddings = nn.Parameter(torch.zeros([token_num, input_dim//2]))
+
+    def forward(self, x):
+        seq_len = x.size(1)
+        pe = self.position_embeddings.expand(x.shape[0], -1, -1)
+        x_with_position = torch.cat([x, pe], dim=-1)
+
+        q = self.query(x_with_position)
+        k = self.key(x_with_position)
+        v = self.value(x_with_position)
+
+        attention = torch.matmul(q, k.transpose(1, 2))
+        attention = attention / torch.sqrt(torch.tensor(k.size(-1), dtype=torch.float32))
+
+        attention = self.softmax(attention)
+        output = torch.matmul(attention, v)
+
+        output = output.view([output.shape[0], -1])
+        output = self.fc(output)
+        return output
+

models/vision_transformer.py

@@ -0,0 +1,481 @@
+import math
+from collections import OrderedDict
+from functools import partial
+from typing import Any, Callable, List, NamedTuple, Optional
+
+import torch
+import torch.nn as nn
+
+# from .._internally_replaced_utils import load_state_dict_from_url
+from .vision_transformer_misc import ConvNormActivation
+from .vision_transformer_utils import _log_api_usage_once
+
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# __all__ = [
+#     "VisionTransformer",
+#     "vit_b_16",
+#     "vit_b_32",
+#     "vit_l_16",
+#     "vit_l_32",
+# ]
+
+model_urls = {
+    "vit_b_16": "https://download.pytorch.org/models/vit_b_16-c867db91.pth",
+    "vit_b_32": "https://download.pytorch.org/models/vit_b_32-d86f8d99.pth",
+    "vit_l_16": "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth",
+    "vit_l_32": "https://download.pytorch.org/models/vit_l_32-c7638314.pth",
+}
+
+
+class ConvStemConfig(NamedTuple):
+    out_channels: int
+    kernel_size: int
+    stride: int
+    norm_layer: Callable[..., nn.Module] = nn.BatchNorm2d
+    activation_layer: Callable[..., nn.Module] = nn.ReLU
+
+
+class MLPBlock(nn.Sequential):
+    """Transformer MLP block."""
+
+    def __init__(self, in_dim: int, mlp_dim: int, dropout: float):
+        super().__init__()
+        self.linear_1 = nn.Linear(in_dim, mlp_dim)
+        self.act = nn.GELU()
+        self.dropout_1 = nn.Dropout(dropout)
+        self.linear_2 = nn.Linear(mlp_dim, in_dim)
+        self.dropout_2 = nn.Dropout(dropout)
+
+        nn.init.xavier_uniform_(self.linear_1.weight)
+        nn.init.xavier_uniform_(self.linear_2.weight)
+        nn.init.normal_(self.linear_1.bias, std=1e-6)
+        nn.init.normal_(self.linear_2.bias, std=1e-6)
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+
+        # Attention block
+        self.ln_1 = norm_layer(hidden_dim)
+        self.self_attention = nn.MultiheadAttention(hidden_dim, num_heads, dropout=attention_dropout, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+
+        # MLP block
+        self.ln_2 = norm_layer(hidden_dim)
+        self.mlp = MLPBlock(hidden_dim, mlp_dim, dropout)
+
+    def forward(self, input: torch.Tensor):
+        torch._assert(input.dim() == 3, f"Expected (seq_length, batch_size, hidden_dim) got {input.shape}")
+        x = self.ln_1(input)
+        x, _ = self.self_attention(query=x, key=x, value=x, need_weights=False)
+        x = self.dropout(x)
+        x = x + input
+
+        y = self.ln_2(x)
+        y = self.mlp(y)
+        return x + y
+
+
+class Encoder(nn.Module):
+    """Transformer Model Encoder for sequence to sequence translation."""
+
+    def __init__(
+        self,
+        seq_length: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float,
+        attention_dropout: float,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+    ):
+        super().__init__()
+        # Note that batch_size is on the first dim because
+        # we have batch_first=True in nn.MultiAttention() by default
+        self.pos_embedding = nn.Parameter(torch.empty(1, seq_length, hidden_dim).normal_(std=0.02))  # from BERT
+        self.dropout = nn.Dropout(dropout)
+        layers: OrderedDict[str, nn.Module] = OrderedDict()
+        for i in range(num_layers):
+            layers[f"encoder_layer_{i}"] = EncoderBlock(
+                num_heads,
+                hidden_dim,
+                mlp_dim,
+                dropout,
+                attention_dropout,
+                norm_layer,
+            )
+        self.layers = nn.Sequential(layers)
+        self.ln = norm_layer(hidden_dim)
+
+    def forward(self, input: torch.Tensor):
+        torch._assert(input.dim() == 3, f"Expected (batch_size, seq_length, hidden_dim) got {input.shape}")
+        input = input + self.pos_embedding
+        return self.ln(self.layers(self.dropout(input)))
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer as per https://arxiv.org/abs/2010.11929."""
+
+    def __init__(
+        self,
+        image_size: int,
+        patch_size: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_dim: int,
+        mlp_dim: int,
+        dropout: float = 0.0,
+        attention_dropout: float = 0.0,
+        num_classes: int = 1000,
+        representation_size: Optional[int] = None,
+        norm_layer: Callable[..., torch.nn.Module] = partial(nn.LayerNorm, eps=1e-6),
+        conv_stem_configs: Optional[List[ConvStemConfig]] = None,
+    ):
+        super().__init__()
+        _log_api_usage_once(self)
+        torch._assert(image_size % patch_size == 0, "Input shape indivisible by patch size!")
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_dim = hidden_dim
+        self.mlp_dim = mlp_dim
+        self.attention_dropout = attention_dropout
+        self.dropout = dropout
+        self.num_classes = num_classes
+        self.representation_size = representation_size
+        self.norm_layer = norm_layer
+
+        if conv_stem_configs is not None:
+            # As per https://arxiv.org/abs/2106.14881
+            seq_proj = nn.Sequential()
+            prev_channels = 3
+            for i, conv_stem_layer_config in enumerate(conv_stem_configs):
+                seq_proj.add_module(
+                    f"conv_bn_relu_{i}",
+                    ConvNormActivation(
+                        in_channels=prev_channels,
+                        out_channels=conv_stem_layer_config.out_channels,
+                        kernel_size=conv_stem_layer_config.kernel_size,
+                        stride=conv_stem_layer_config.stride,
+                        norm_layer=conv_stem_layer_config.norm_layer,
+                        activation_layer=conv_stem_layer_config.activation_layer,
+                    ),
+                )
+                prev_channels = conv_stem_layer_config.out_channels
+            seq_proj.add_module(
+                "conv_last", nn.Conv2d(in_channels=prev_channels, out_channels=hidden_dim, kernel_size=1)
+            )
+            self.conv_proj: nn.Module = seq_proj
+        else:
+            self.conv_proj = nn.Conv2d(
+                in_channels=3, out_channels=hidden_dim, kernel_size=patch_size, stride=patch_size
+            )
+
+        seq_length = (image_size // patch_size) ** 2
+
+        # Add a class token
+        self.class_token = nn.Parameter(torch.zeros(1, 1, hidden_dim))
+        seq_length += 1
+
+        self.encoder = Encoder(
+            seq_length,
+            num_layers,
+            num_heads,
+            hidden_dim,
+            mlp_dim,
+            dropout,
+            attention_dropout,
+            norm_layer,
+        )
+        self.seq_length = seq_length
+
+        heads_layers: OrderedDict[str, nn.Module] = OrderedDict()
+        if representation_size is None:
+            heads_layers["head"] = nn.Linear(hidden_dim, num_classes)
+        else:
+            heads_layers["pre_logits"] = nn.Linear(hidden_dim, representation_size)
+            heads_layers["act"] = nn.Tanh()
+            heads_layers["head"] = nn.Linear(representation_size, num_classes)
+
+        self.heads = nn.Sequential(heads_layers)
+
+        if isinstance(self.conv_proj, nn.Conv2d):
+            # Init the patchify stem
+            fan_in = self.conv_proj.in_channels * self.conv_proj.kernel_size[0] * self.conv_proj.kernel_size[1]
+            nn.init.trunc_normal_(self.conv_proj.weight, std=math.sqrt(1 / fan_in))
+            if self.conv_proj.bias is not None:
+                nn.init.zeros_(self.conv_proj.bias)
+        elif self.conv_proj.conv_last is not None and isinstance(self.conv_proj.conv_last, nn.Conv2d):
+            # Init the last 1x1 conv of the conv stem
+            nn.init.normal_(
+                self.conv_proj.conv_last.weight, mean=0.0, std=math.sqrt(2.0 / self.conv_proj.conv_last.out_channels)
+            )
+            if self.conv_proj.conv_last.bias is not None:
+                nn.init.zeros_(self.conv_proj.conv_last.bias)
+
+        if hasattr(self.heads, "pre_logits") and isinstance(self.heads.pre_logits, nn.Linear):
+            fan_in = self.heads.pre_logits.in_features
+            nn.init.trunc_normal_(self.heads.pre_logits.weight, std=math.sqrt(1 / fan_in))
+            nn.init.zeros_(self.heads.pre_logits.bias)
+
+        if isinstance(self.heads.head, nn.Linear):
+            nn.init.zeros_(self.heads.head.weight)
+            nn.init.zeros_(self.heads.head.bias)
+
+    def _process_input(self, x: torch.Tensor) -> torch.Tensor:
+        n, c, h, w = x.shape
+        p = self.patch_size
+        torch._assert(h == self.image_size, "Wrong image height!")
+        torch._assert(w == self.image_size, "Wrong image width!")
+        n_h = h // p
+        n_w = w // p
+
+        # (n, c, h, w) -> (n, hidden_dim, n_h, n_w)
+        x = self.conv_proj(x)
+        # (n, hidden_dim, n_h, n_w) -> (n, hidden_dim, (n_h * n_w))
+        x = x.reshape(n, self.hidden_dim, n_h * n_w)
+
+        # (n, hidden_dim, (n_h * n_w)) -> (n, (n_h * n_w), hidden_dim)
+        # The self attention layer expects inputs in the format (N, S, E)
+        # where S is the source sequence length, N is the batch size, E is the
+        # embedding dimension
+        x = x.permute(0, 2, 1)
+
+        return x
+
+    def forward(self, x: torch.Tensor):
+        out = {}
+
+        # Reshape and permute the input tensor
+        x = self._process_input(x)
+        n = x.shape[0]
+
+        # Expand the class token to the full batch
+        batch_class_token = self.class_token.expand(n, -1, -1)
+        x = torch.cat([batch_class_token, x], dim=1)
+
+        
+        x = self.encoder(x)
+        img_feature = x[:,1:]
+        H = W = int(self.image_size / self.patch_size)
+        out['f4'] = img_feature.view(n, H, W, self.hidden_dim).permute(0,3,1,2)
+
+        # Classifier "token" as used by standard language architectures
+        x = x[:, 0]
+        out['penultimate'] = x 
+
+        x = self.heads(x) # I checked that for all pretrained ViT, this is just a fc 
+        out['logits'] = x 
+
+        return out
+
+
+def _vision_transformer(
+    arch: str,
+    patch_size: int,
+    num_layers: int,
+    num_heads: int,
+    hidden_dim: int,
+    mlp_dim: int,
+    pretrained: bool,
+    progress: bool,
+    **kwargs: Any,
+) -> VisionTransformer:
+    image_size = kwargs.pop("image_size", 224)
+
+    model = VisionTransformer(
+        image_size=image_size,
+        patch_size=patch_size,
+        num_layers=num_layers,
+        num_heads=num_heads,
+        hidden_dim=hidden_dim,
+        mlp_dim=mlp_dim,
+        **kwargs,
+    )
+
+    if pretrained:
+        if arch not in model_urls:
+            raise ValueError(f"No checkpoint is available for model type '{arch}'!")
+        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        model.load_state_dict(state_dict)
+
+    return model
+
+
+def vit_b_16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+    """
+    Constructs a vit_b_16 architecture from
+    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" <https://arxiv.org/abs/2010.11929>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _vision_transformer(
+        arch="vit_b_16",
+        patch_size=16,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        pretrained=pretrained,
+        progress=progress,
+        **kwargs,
+    )
+
+
+def vit_b_32(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+    """
+    Constructs a vit_b_32 architecture from
+    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" <https://arxiv.org/abs/2010.11929>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _vision_transformer(
+        arch="vit_b_32",
+        patch_size=32,
+        num_layers=12,
+        num_heads=12,
+        hidden_dim=768,
+        mlp_dim=3072,
+        pretrained=pretrained,
+        progress=progress,
+        **kwargs,
+    )
+
+
+def vit_l_16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+    """
+    Constructs a vit_l_16 architecture from
+    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" <https://arxiv.org/abs/2010.11929>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _vision_transformer(
+        arch="vit_l_16",
+        patch_size=16,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        pretrained=pretrained,
+        progress=progress,
+        **kwargs,
+    )
+
+
+def vit_l_32(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VisionTransformer:
+    """
+    Constructs a vit_l_32 architecture from
+    `"An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale" <https://arxiv.org/abs/2010.11929>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _vision_transformer(
+        arch="vit_l_32",
+        patch_size=32,
+        num_layers=24,
+        num_heads=16,
+        hidden_dim=1024,
+        mlp_dim=4096,
+        pretrained=pretrained,
+        progress=progress,
+        **kwargs,
+    )
+
+
+def interpolate_embeddings(
+    image_size: int,
+    patch_size: int,
+    model_state: "OrderedDict[str, torch.Tensor]",
+    interpolation_mode: str = "bicubic",
+    reset_heads: bool = False,
+) -> "OrderedDict[str, torch.Tensor]":
+    """This function helps interpolating positional embeddings during checkpoint loading,
+    especially when you want to apply a pre-trained model on images with different resolution.
+
+    Args:
+        image_size (int): Image size of the new model.
+        patch_size (int): Patch size of the new model.
+        model_state (OrderedDict[str, torch.Tensor]): State dict of the pre-trained model.
+        interpolation_mode (str): The algorithm used for upsampling. Default: bicubic.
+        reset_heads (bool): If true, not copying the state of heads. Default: False.
+
+    Returns:
+        OrderedDict[str, torch.Tensor]: A state dict which can be loaded into the new model.
+    """
+    # Shape of pos_embedding is (1, seq_length, hidden_dim)
+    pos_embedding = model_state["encoder.pos_embedding"]
+    n, seq_length, hidden_dim = pos_embedding.shape
+    if n != 1:
+        raise ValueError(f"Unexpected position embedding shape: {pos_embedding.shape}")
+
+    new_seq_length = (image_size // patch_size) ** 2 + 1
+
+    # Need to interpolate the weights for the position embedding.
+    # We do this by reshaping the positions embeddings to a 2d grid, performing
+    # an interpolation in the (h, w) space and then reshaping back to a 1d grid.
+    if new_seq_length != seq_length:
+        # The class token embedding shouldn't be interpolated so we split it up.
+        seq_length -= 1
+        new_seq_length -= 1
+        pos_embedding_token = pos_embedding[:, :1, :]
+        pos_embedding_img = pos_embedding[:, 1:, :]
+
+        # (1, seq_length, hidden_dim) -> (1, hidden_dim, seq_length)
+        pos_embedding_img = pos_embedding_img.permute(0, 2, 1)
+        seq_length_1d = int(math.sqrt(seq_length))
+        torch._assert(seq_length_1d * seq_length_1d == seq_length, "seq_length is not a perfect square!")
+
+        # (1, hidden_dim, seq_length) -> (1, hidden_dim, seq_l_1d, seq_l_1d)
+        pos_embedding_img = pos_embedding_img.reshape(1, hidden_dim, seq_length_1d, seq_length_1d)
+        new_seq_length_1d = image_size // patch_size
+
+        # Perform interpolation.
+        # (1, hidden_dim, seq_l_1d, seq_l_1d) -> (1, hidden_dim, new_seq_l_1d, new_seq_l_1d)
+        new_pos_embedding_img = nn.functional.interpolate(
+            pos_embedding_img,
+            size=new_seq_length_1d,
+            mode=interpolation_mode,
+            align_corners=True,
+        )
+
+        # (1, hidden_dim, new_seq_l_1d, new_seq_l_1d) -> (1, hidden_dim, new_seq_length)
+        new_pos_embedding_img = new_pos_embedding_img.reshape(1, hidden_dim, new_seq_length)
+
+        # (1, hidden_dim, new_seq_length) -> (1, new_seq_length, hidden_dim)
+        new_pos_embedding_img = new_pos_embedding_img.permute(0, 2, 1)
+        new_pos_embedding = torch.cat([pos_embedding_token, new_pos_embedding_img], dim=1)
+
+        model_state["encoder.pos_embedding"] = new_pos_embedding
+
+        if reset_heads:
+            model_state_copy: "OrderedDict[str, torch.Tensor]" = OrderedDict()
+            for k, v in model_state.items():
+                if not k.startswith("heads"):
+                    model_state_copy[k] = v
+            model_state = model_state_copy
+
+    return model_state

models/vision_transformer_misc.py

@@ -0,0 +1,163 @@
+from typing import Callable, List, Optional
+
+import torch
+from torch import Tensor
+
+from .vision_transformer_utils import _log_api_usage_once
+
+
+interpolate = torch.nn.functional.interpolate
+
+
+# This is not in nn
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed
+
+    Args:
+        num_features (int): Number of features ``C`` from an expected input of size ``(N, C, H, W)``
+        eps (float): a value added to the denominator for numerical stability. Default: 1e-5
+    """
+
+    def __init__(
+        self,
+        num_features: int,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        _log_api_usage_once(self)
+        self.eps = eps
+        self.register_buffer("weight", torch.ones(num_features))
+        self.register_buffer("bias", torch.zeros(num_features))
+        self.register_buffer("running_mean", torch.zeros(num_features))
+        self.register_buffer("running_var", torch.ones(num_features))
+
+    def _load_from_state_dict(
+        self,
+        state_dict: dict,
+        prefix: str,
+        local_metadata: dict,
+        strict: bool,
+        missing_keys: List[str],
+        unexpected_keys: List[str],
+        error_msgs: List[str],
+    ):
+        num_batches_tracked_key = prefix + "num_batches_tracked"
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        # move reshapes to the beginning
+        # to make it fuser-friendly
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        scale = w * (rv + self.eps).rsqrt()
+        bias = b - rm * scale
+        return x * scale + bias
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.weight.shape[0]}, eps={self.eps})"
+
+
+class ConvNormActivation(torch.nn.Sequential):
+    """
+    Configurable block used for Convolution-Normalzation-Activation blocks.
+
+    Args:
+        in_channels (int): Number of channels in the input image
+        out_channels (int): Number of channels produced by the Convolution-Normalzation-Activation block
+        kernel_size: (int, optional): Size of the convolving kernel. Default: 3
+        stride (int, optional): Stride of the convolution. Default: 1
+        padding (int, tuple or str, optional): Padding added to all four sides of the input. Default: None, in wich case it will calculated as ``padding = (kernel_size - 1) // 2 * dilation``
+        groups (int, optional): Number of blocked connections from input channels to output channels. Default: 1
+        norm_layer (Callable[..., torch.nn.Module], optional): Norm layer that will be stacked on top of the convolutiuon layer. If ``None`` this layer wont be used. Default: ``torch.nn.BatchNorm2d``
+        activation_layer (Callable[..., torch.nn.Module], optinal): Activation function which will be stacked on top of the normalization layer (if not None), otherwise on top of the conv layer. If ``None`` this layer wont be used. Default: ``torch.nn.ReLU``
+        dilation (int): Spacing between kernel elements. Default: 1
+        inplace (bool): Parameter for the activation layer, which can optionally do the operation in-place. Default ``True``
+        bias (bool, optional): Whether to use bias in the convolution layer. By default, biases are included if ``norm_layer is None``.
+
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: Optional[int] = None,
+        groups: int = 1,
+        norm_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.BatchNorm2d,
+        activation_layer: Optional[Callable[..., torch.nn.Module]] = torch.nn.ReLU,
+        dilation: int = 1,
+        inplace: Optional[bool] = True,
+        bias: Optional[bool] = None,
+    ) -> None:
+        if padding is None:
+            padding = (kernel_size - 1) // 2 * dilation
+        if bias is None:
+            bias = norm_layer is None
+        layers = [
+            torch.nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                dilation=dilation,
+                groups=groups,
+                bias=bias,
+            )
+        ]
+        if norm_layer is not None:
+            layers.append(norm_layer(out_channels))
+        if activation_layer is not None:
+            params = {} if inplace is None else {"inplace": inplace}
+            layers.append(activation_layer(**params))
+        super().__init__(*layers)
+        _log_api_usage_once(self)
+        self.out_channels = out_channels
+
+
+class SqueezeExcitation(torch.nn.Module):
+    """
+    This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1).
+    Parameters ``activation``, and ``scale_activation`` correspond to ``delta`` and ``sigma`` in in eq. 3.
+
+    Args:
+        input_channels (int): Number of channels in the input image
+        squeeze_channels (int): Number of squeeze channels
+        activation (Callable[..., torch.nn.Module], optional): ``delta`` activation. Default: ``torch.nn.ReLU``
+        scale_activation (Callable[..., torch.nn.Module]): ``sigma`` activation. Default: ``torch.nn.Sigmoid``
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        squeeze_channels: int,
+        activation: Callable[..., torch.nn.Module] = torch.nn.ReLU,
+        scale_activation: Callable[..., torch.nn.Module] = torch.nn.Sigmoid,
+    ) -> None:
+        super().__init__()
+        _log_api_usage_once(self)
+        self.avgpool = torch.nn.AdaptiveAvgPool2d(1)
+        self.fc1 = torch.nn.Conv2d(input_channels, squeeze_channels, 1)
+        self.fc2 = torch.nn.Conv2d(squeeze_channels, input_channels, 1)
+        self.activation = activation()
+        self.scale_activation = scale_activation()
+
+    def _scale(self, input: Tensor) -> Tensor:
+        scale = self.avgpool(input)
+        scale = self.fc1(scale)
+        scale = self.activation(scale)
+        scale = self.fc2(scale)
+        return self.scale_activation(scale)
+
+    def forward(self, input: Tensor) -> Tensor:
+        scale = self._scale(input)
+        return scale * input

models/vision_transformer_utils.py

@@ -0,0 +1,549 @@
+import math
+import pathlib
+import warnings
+from types import FunctionType
+from typing import Any, BinaryIO, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from PIL import Image, ImageColor, ImageDraw, ImageFont
+
+__all__ = [
+    "make_grid",
+    "save_image",
+    "draw_bounding_boxes",
+    "draw_segmentation_masks",
+    "draw_keypoints",
+    "flow_to_image",
+]
+
+
+@torch.no_grad()
+def make_grid(
+    tensor: Union[torch.Tensor, List[torch.Tensor]],
+    nrow: int = 8,
+    padding: int = 2,
+    normalize: bool = False,
+    value_range: Optional[Tuple[int, int]] = None,
+    scale_each: bool = False,
+    pad_value: float = 0.0,
+    **kwargs,
+) -> torch.Tensor:
+    """
+    Make a grid of images.
+
+    Args:
+        tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
+            or a list of images all of the same size.
+        nrow (int, optional): Number of images displayed in each row of the grid.
+            The final grid size is ``(B / nrow, nrow)``. Default: ``8``.
+        padding (int, optional): amount of padding. Default: ``2``.
+        normalize (bool, optional): If True, shift the image to the range (0, 1),
+            by the min and max values specified by ``value_range``. Default: ``False``.
+        value_range (tuple, optional): tuple (min, max) where min and max are numbers,
+            then these numbers are used to normalize the image. By default, min and max
+            are computed from the tensor.
+        range (tuple. optional):
+            .. warning::
+                This parameter was deprecated in ``0.12`` and will be removed in ``0.14``. Please use ``value_range``
+                instead.
+        scale_each (bool, optional): If ``True``, scale each image in the batch of
+            images separately rather than the (min, max) over all images. Default: ``False``.
+        pad_value (float, optional): Value for the padded pixels. Default: ``0``.
+
+    Returns:
+        grid (Tensor): the tensor containing grid of images.
+    """
+    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+        _log_api_usage_once(make_grid)
+    if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
+        raise TypeError(f"tensor or list of tensors expected, got {type(tensor)}")
+
+    if "range" in kwargs.keys():
+        warnings.warn(
+            "The parameter 'range' is deprecated since 0.12 and will be removed in 0.14. "
+            "Please use 'value_range' instead."
+        )
+        value_range = kwargs["range"]
+
+    # if list of tensors, convert to a 4D mini-batch Tensor
+    if isinstance(tensor, list):
+        tensor = torch.stack(tensor, dim=0)
+
+    if tensor.dim() == 2:  # single image H x W
+        tensor = tensor.unsqueeze(0)
+    if tensor.dim() == 3:  # single image
+        if tensor.size(0) == 1:  # if single-channel, convert to 3-channel
+            tensor = torch.cat((tensor, tensor, tensor), 0)
+        tensor = tensor.unsqueeze(0)
+
+    if tensor.dim() == 4 and tensor.size(1) == 1:  # single-channel images
+        tensor = torch.cat((tensor, tensor, tensor), 1)
+
+    if normalize is True:
+        tensor = tensor.clone()  # avoid modifying tensor in-place
+        if value_range is not None:
+            assert isinstance(
+                value_range, tuple
+            ), "value_range has to be a tuple (min, max) if specified. min and max are numbers"
+
+        def norm_ip(img, low, high):
+            img.clamp_(min=low, max=high)
+            img.sub_(low).div_(max(high - low, 1e-5))
+
+        def norm_range(t, value_range):
+            if value_range is not None:
+                norm_ip(t, value_range[0], value_range[1])
+            else:
+                norm_ip(t, float(t.min()), float(t.max()))
+
+        if scale_each is True:
+            for t in tensor:  # loop over mini-batch dimension
+                norm_range(t, value_range)
+        else:
+            norm_range(tensor, value_range)
+
+    assert isinstance(tensor, torch.Tensor)
+    if tensor.size(0) == 1:
+        return tensor.squeeze(0)
+
+    # make the mini-batch of images into a grid
+    nmaps = tensor.size(0)
+    xmaps = min(nrow, nmaps)
+    ymaps = int(math.ceil(float(nmaps) / xmaps))
+    height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
+    num_channels = tensor.size(1)
+    grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
+    k = 0
+    for y in range(ymaps):
+        for x in range(xmaps):
+            if k >= nmaps:
+                break
+            # Tensor.copy_() is a valid method but seems to be missing from the stubs
+            # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_
+            grid.narrow(1, y * height + padding, height - padding).narrow(  # type: ignore[attr-defined]
+                2, x * width + padding, width - padding
+            ).copy_(tensor[k])
+            k = k + 1
+    return grid
+
+
+@torch.no_grad()
+def save_image(
+    tensor: Union[torch.Tensor, List[torch.Tensor]],
+    fp: Union[str, pathlib.Path, BinaryIO],
+    format: Optional[str] = None,
+    **kwargs,
+) -> None:
+    """
+    Save a given Tensor into an image file.
+
+    Args:
+        tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
+            saves the tensor as a grid of images by calling ``make_grid``.
+        fp (string or file object): A filename or a file object
+        format(Optional):  If omitted, the format to use is determined from the filename extension.
+            If a file object was used instead of a filename, this parameter should always be used.
+        **kwargs: Other arguments are documented in ``make_grid``.
+    """
+
+    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+        _log_api_usage_once(save_image)
+    grid = make_grid(tensor, **kwargs)
+    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
+    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    im = Image.fromarray(ndarr)
+    im.save(fp, format=format)
+
+
+@torch.no_grad()
+def draw_bounding_boxes(
+    image: torch.Tensor,
+    boxes: torch.Tensor,
+    labels: Optional[List[str]] = None,
+    colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
+    fill: Optional[bool] = False,
+    width: int = 1,
+    font: Optional[str] = None,
+    font_size: int = 10,
+) -> torch.Tensor:
+
+    """
+    Draws bounding boxes on given image.
+    The values of the input image should be uint8 between 0 and 255.
+    If fill is True, Resulting Tensor should be saved as PNG image.
+
+    Args:
+        image (Tensor): Tensor of shape (C x H x W) and dtype uint8.
+        boxes (Tensor): Tensor of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format. Note that
+            the boxes are absolute coordinates with respect to the image. In other words: `0 <= xmin < xmax < W` and
+            `0 <= ymin < ymax < H`.
+        labels (List[str]): List containing the labels of bounding boxes.
+        colors (color or list of colors, optional): List containing the colors
+            of the boxes or single color for all boxes. The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for boxes.
+        fill (bool): If `True` fills the bounding box with specified color.
+        width (int): Width of bounding box.
+        font (str): A filename containing a TrueType font. If the file is not found in this filename, the loader may
+            also search in other directories, such as the `fonts/` directory on Windows or `/Library/Fonts/`,
+            `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
+        font_size (int): The requested font size in points.
+
+    Returns:
+        img (Tensor[C, H, W]): Image Tensor of dtype uint8 with bounding boxes plotted.
+    """
+
+    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+        _log_api_usage_once(draw_bounding_boxes)
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"Tensor expected, got {type(image)}")
+    elif image.dtype != torch.uint8:
+        raise ValueError(f"Tensor uint8 expected, got {image.dtype}")
+    elif image.dim() != 3:
+        raise ValueError("Pass individual images, not batches")
+    elif image.size(0) not in {1, 3}:
+        raise ValueError("Only grayscale and RGB images are supported")
+
+    num_boxes = boxes.shape[0]
+
+    if labels is None:
+        labels: Union[List[str], List[None]] = [None] * num_boxes  # type: ignore[no-redef]
+    elif len(labels) != num_boxes:
+        raise ValueError(
+            f"Number of boxes ({num_boxes}) and labels ({len(labels)}) mismatch. Please specify labels for each box."
+        )
+
+    if colors is None:
+        colors = _generate_color_palette(num_boxes)
+    elif isinstance(colors, list):
+        if len(colors) < num_boxes:
+            raise ValueError(f"Number of colors ({len(colors)}) is less than number of boxes ({num_boxes}). ")
+    else:  # colors specifies a single color for all boxes
+        colors = [colors] * num_boxes
+
+    colors = [(ImageColor.getrgb(color) if isinstance(color, str) else color) for color in colors]
+
+    # Handle Grayscale images
+    if image.size(0) == 1:
+        image = torch.tile(image, (3, 1, 1))
+
+    ndarr = image.permute(1, 2, 0).cpu().numpy()
+    img_to_draw = Image.fromarray(ndarr)
+    img_boxes = boxes.to(torch.int64).tolist()
+
+    if fill:
+        draw = ImageDraw.Draw(img_to_draw, "RGBA")
+    else:
+        draw = ImageDraw.Draw(img_to_draw)
+
+    txt_font = ImageFont.load_default() if font is None else ImageFont.truetype(font=font, size=font_size)
+
+    for bbox, color, label in zip(img_boxes, colors, labels):  # type: ignore[arg-type]
+        if fill:
+            fill_color = color + (100,)
+            draw.rectangle(bbox, width=width, outline=color, fill=fill_color)
+        else:
+            draw.rectangle(bbox, width=width, outline=color)
+
+        if label is not None:
+            margin = width + 1
+            draw.text((bbox[0] + margin, bbox[1] + margin), label, fill=color, font=txt_font)
+
+    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+
+
+@torch.no_grad()
+def draw_segmentation_masks(
+    image: torch.Tensor,
+    masks: torch.Tensor,
+    alpha: float = 0.8,
+    colors: Optional[Union[List[Union[str, Tuple[int, int, int]]], str, Tuple[int, int, int]]] = None,
+) -> torch.Tensor:
+
+    """
+    Draws segmentation masks on given RGB image.
+    The values of the input image should be uint8 between 0 and 255.
+
+    Args:
+        image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+        masks (Tensor): Tensor of shape (num_masks, H, W) or (H, W) and dtype bool.
+        alpha (float): Float number between 0 and 1 denoting the transparency of the masks.
+            0 means full transparency, 1 means no transparency.
+        colors (color or list of colors, optional): List containing the colors
+            of the masks or single color for all masks. The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+            By default, random colors are generated for each mask.
+
+    Returns:
+        img (Tensor[C, H, W]): Image Tensor, with segmentation masks drawn on top.
+    """
+
+    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+        _log_api_usage_once(draw_segmentation_masks)
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"The image must be a tensor, got {type(image)}")
+    elif image.dtype != torch.uint8:
+        raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+    elif image.dim() != 3:
+        raise ValueError("Pass individual images, not batches")
+    elif image.size()[0] != 3:
+        raise ValueError("Pass an RGB image. Other Image formats are not supported")
+    if masks.ndim == 2:
+        masks = masks[None, :, :]
+    if masks.ndim != 3:
+        raise ValueError("masks must be of shape (H, W) or (batch_size, H, W)")
+    if masks.dtype != torch.bool:
+        raise ValueError(f"The masks must be of dtype bool. Got {masks.dtype}")
+    if masks.shape[-2:] != image.shape[-2:]:
+        raise ValueError("The image and the masks must have the same height and width")
+
+    num_masks = masks.size()[0]
+    if colors is not None and num_masks > len(colors):
+        raise ValueError(f"There are more masks ({num_masks}) than colors ({len(colors)})")
+
+    if colors is None:
+        colors = _generate_color_palette(num_masks)
+
+    if not isinstance(colors, list):
+        colors = [colors]
+    if not isinstance(colors[0], (tuple, str)):
+        raise ValueError("colors must be a tuple or a string, or a list thereof")
+    if isinstance(colors[0], tuple) and len(colors[0]) != 3:
+        raise ValueError("It seems that you passed a tuple of colors instead of a list of colors")
+
+    out_dtype = torch.uint8
+
+    colors_ = []
+    for color in colors:
+        if isinstance(color, str):
+            color = ImageColor.getrgb(color)
+        colors_.append(torch.tensor(color, dtype=out_dtype))
+
+    img_to_draw = image.detach().clone()
+    # TODO: There might be a way to vectorize this
+    for mask, color in zip(masks, colors_):
+        img_to_draw[:, mask] = color[:, None]
+
+    out = image * (1 - alpha) + img_to_draw * alpha
+    return out.to(out_dtype)
+
+
+@torch.no_grad()
+def draw_keypoints(
+    image: torch.Tensor,
+    keypoints: torch.Tensor,
+    connectivity: Optional[List[Tuple[int, int]]] = None,
+    colors: Optional[Union[str, Tuple[int, int, int]]] = None,
+    radius: int = 2,
+    width: int = 3,
+) -> torch.Tensor:
+
+    """
+    Draws Keypoints on given RGB image.
+    The values of the input image should be uint8 between 0 and 255.
+
+    Args:
+        image (Tensor): Tensor of shape (3, H, W) and dtype uint8.
+        keypoints (Tensor): Tensor of shape (num_instances, K, 2) the K keypoints location for each of the N instances,
+            in the format [x, y].
+        connectivity (List[Tuple[int, int]]]): A List of tuple where,
+            each tuple contains pair of keypoints to be connected.
+        colors (str, Tuple): The color can be represented as
+            PIL strings e.g. "red" or "#FF00FF", or as RGB tuples e.g. ``(240, 10, 157)``.
+        radius (int): Integer denoting radius of keypoint.
+        width (int): Integer denoting width of line connecting keypoints.
+
+    Returns:
+        img (Tensor[C, H, W]): Image Tensor of dtype uint8 with keypoints drawn.
+    """
+
+    if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+        _log_api_usage_once(draw_keypoints)
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"The image must be a tensor, got {type(image)}")
+    elif image.dtype != torch.uint8:
+        raise ValueError(f"The image dtype must be uint8, got {image.dtype}")
+    elif image.dim() != 3:
+        raise ValueError("Pass individual images, not batches")
+    elif image.size()[0] != 3:
+        raise ValueError("Pass an RGB image. Other Image formats are not supported")
+
+    if keypoints.ndim != 3:
+        raise ValueError("keypoints must be of shape (num_instances, K, 2)")
+
+    ndarr = image.permute(1, 2, 0).cpu().numpy()
+    img_to_draw = Image.fromarray(ndarr)
+    draw = ImageDraw.Draw(img_to_draw)
+    img_kpts = keypoints.to(torch.int64).tolist()
+
+    for kpt_id, kpt_inst in enumerate(img_kpts):
+        for inst_id, kpt in enumerate(kpt_inst):
+            x1 = kpt[0] - radius
+            x2 = kpt[0] + radius
+            y1 = kpt[1] - radius
+            y2 = kpt[1] + radius
+            draw.ellipse([x1, y1, x2, y2], fill=colors, outline=None, width=0)
+
+        if connectivity:
+            for connection in connectivity:
+                start_pt_x = kpt_inst[connection[0]][0]
+                start_pt_y = kpt_inst[connection[0]][1]
+
+                end_pt_x = kpt_inst[connection[1]][0]
+                end_pt_y = kpt_inst[connection[1]][1]
+
+                draw.line(
+                    ((start_pt_x, start_pt_y), (end_pt_x, end_pt_y)),
+                    width=width,
+                )
+
+    return torch.from_numpy(np.array(img_to_draw)).permute(2, 0, 1).to(dtype=torch.uint8)
+
+
+# Flow visualization code adapted from https://github.com/tomrunia/OpticalFlow_Visualization
+@torch.no_grad()
+def flow_to_image(flow: torch.Tensor) -> torch.Tensor:
+
+    """
+    Converts a flow to an RGB image.
+
+    Args:
+        flow (Tensor): Flow of shape (N, 2, H, W) or (2, H, W) and dtype torch.float.
+
+    Returns:
+        img (Tensor): Image Tensor of dtype uint8 where each color corresponds
+            to a given flow direction. Shape is (N, 3, H, W) or (3, H, W) depending on the input.
+    """
+
+    if flow.dtype != torch.float:
+        raise ValueError(f"Flow should be of dtype torch.float, got {flow.dtype}.")
+
+    orig_shape = flow.shape
+    if flow.ndim == 3:
+        flow = flow[None]  # Add batch dim
+
+    if flow.ndim != 4 or flow.shape[1] != 2:
+        raise ValueError(f"Input flow should have shape (2, H, W) or (N, 2, H, W), got {orig_shape}.")
+
+    max_norm = torch.sum(flow ** 2, dim=1).sqrt().max()
+    epsilon = torch.finfo((flow).dtype).eps
+    normalized_flow = flow / (max_norm + epsilon)
+    img = _normalized_flow_to_image(normalized_flow)
+
+    if len(orig_shape) == 3:
+        img = img[0]  # Remove batch dim
+    return img
+
+
+@torch.no_grad()
+def _normalized_flow_to_image(normalized_flow: torch.Tensor) -> torch.Tensor:
+
+    """
+    Converts a batch of normalized flow to an RGB image.
+
+    Args:
+        normalized_flow (torch.Tensor): Normalized flow tensor of shape (N, 2, H, W)
+    Returns:
+       img (Tensor(N, 3, H, W)): Flow visualization image of dtype uint8.
+    """
+
+    N, _, H, W = normalized_flow.shape
+    device = normalized_flow.device
+    flow_image = torch.zeros((N, 3, H, W), dtype=torch.uint8, device=device)
+    colorwheel = _make_colorwheel().to(device)  # shape [55x3]
+    num_cols = colorwheel.shape[0]
+    norm = torch.sum(normalized_flow ** 2, dim=1).sqrt()
+    a = torch.atan2(-normalized_flow[:, 1, :, :], -normalized_flow[:, 0, :, :]) / torch.pi
+    fk = (a + 1) / 2 * (num_cols - 1)
+    k0 = torch.floor(fk).to(torch.long)
+    k1 = k0 + 1
+    k1[k1 == num_cols] = 0
+    f = fk - k0
+
+    for c in range(colorwheel.shape[1]):
+        tmp = colorwheel[:, c]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1 - f) * col0 + f * col1
+        col = 1 - norm * (1 - col)
+        flow_image[:, c, :, :] = torch.floor(255 * col)
+    return flow_image
+
+
+def _make_colorwheel() -> torch.Tensor:
+    """
+    Generates a color wheel for optical flow visualization as presented in:
+    Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+    URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf.
+
+    Returns:
+        colorwheel (Tensor[55, 3]): Colorwheel Tensor.
+    """
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = torch.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = torch.floor(255 * torch.arange(0, RY) / RY)
+    col = col + RY
+    # YG
+    colorwheel[col : col + YG, 0] = 255 - torch.floor(255 * torch.arange(0, YG) / YG)
+    colorwheel[col : col + YG, 1] = 255
+    col = col + YG
+    # GC
+    colorwheel[col : col + GC, 1] = 255
+    colorwheel[col : col + GC, 2] = torch.floor(255 * torch.arange(0, GC) / GC)
+    col = col + GC
+    # CB
+    colorwheel[col : col + CB, 1] = 255 - torch.floor(255 * torch.arange(CB) / CB)
+    colorwheel[col : col + CB, 2] = 255
+    col = col + CB
+    # BM
+    colorwheel[col : col + BM, 2] = 255
+    colorwheel[col : col + BM, 0] = torch.floor(255 * torch.arange(0, BM) / BM)
+    col = col + BM
+    # MR
+    colorwheel[col : col + MR, 2] = 255 - torch.floor(255 * torch.arange(MR) / MR)
+    colorwheel[col : col + MR, 0] = 255
+    return colorwheel
+
+
+def _generate_color_palette(num_objects: int):
+    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
+    return [tuple((i * palette) % 255) for i in range(num_objects)]
+
+
+def _log_api_usage_once(obj: Any) -> None:
+
+    """
+    Logs API usage(module and name) within an organization.
+    In a large ecosystem, it's often useful to track the PyTorch and
+    TorchVision APIs usage. This API provides the similar functionality to the
+    logging module in the Python stdlib. It can be used for debugging purpose
+    to log which methods are used and by default it is inactive, unless the user
+    manually subscribes a logger via the `SetAPIUsageLogger method <https://github.com/pytorch/pytorch/blob/eb3b9fe719b21fae13c7a7cf3253f970290a573e/c10/util/Logging.cpp#L114>`_.
+    Please note it is triggered only once for the same API call within a process.
+    It does not collect any data from open-source users since it is no-op by default.
+    For more information, please refer to
+    * PyTorch note: https://pytorch.org/docs/stable/notes/large_scale_deployments.html#api-usage-logging;
+    * Logging policy: https://github.com/pytorch/vision/issues/5052;
+
+    Args:
+        obj (class instance or method): an object to extract info from.
+    """
+    if not obj.__module__.startswith("torchvision"):
+        return
+    name = obj.__class__.__name__
+    if isinstance(obj, FunctionType):
+        name = obj.__name__
+    torch._C._log_api_usage_once(f"{obj.__module__}.{name}")

models/vits.py

@@ -0,0 +1,143 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import torch
+import torch.nn as nn
+from functools import partial, reduce
+from operator import mul
+
+from timm.models.vision_transformer import VisionTransformer, _cfg
+from timm.models.layers.helpers import to_2tuple
+from timm.models.layers import PatchEmbed
+
+__all__ = [
+    'vit_small', 
+    'vit_base',
+    'vit_conv_small',
+    'vit_conv_base',
+]
+
+
+class VisionTransformerMoCo(VisionTransformer):
+    def __init__(self, stop_grad_conv1=False, **kwargs):
+        super().__init__(**kwargs)
+        # Use fixed 2D sin-cos position embedding
+        self.build_2d_sincos_position_embedding()
+
+        # weight initialization
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear):
+                if 'qkv' in name:
+                    # treat the weights of Q, K, V separately
+                    val = math.sqrt(6. / float(m.weight.shape[0] // 3 + m.weight.shape[1]))
+                    nn.init.uniform_(m.weight, -val, val)
+                else:
+                    nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+        nn.init.normal_(self.cls_token, std=1e-6)
+
+        if isinstance(self.patch_embed, PatchEmbed):
+            # xavier_uniform initialization
+            val = math.sqrt(6. / float(3 * reduce(mul, self.patch_embed.patch_size, 1) + self.embed_dim))
+            nn.init.uniform_(self.patch_embed.proj.weight, -val, val)
+            nn.init.zeros_(self.patch_embed.proj.bias)
+
+            if stop_grad_conv1:
+                self.patch_embed.proj.weight.requires_grad = False
+                self.patch_embed.proj.bias.requires_grad = False
+
+    def build_2d_sincos_position_embedding(self, temperature=10000.):
+        h, w = self.patch_embed.grid_size
+        grid_w = torch.arange(w, dtype=torch.float32)
+        grid_h = torch.arange(h, dtype=torch.float32)
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
+        assert self.embed_dim % 4 == 0, 'Embed dimension must be divisible by 4 for 2D sin-cos position embedding'
+        pos_dim = self.embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+        omega = 1. / (temperature**omega)
+        out_w = torch.einsum('m,d->md', [grid_w.flatten(), omega])
+        out_h = torch.einsum('m,d->md', [grid_h.flatten(), omega])
+        pos_emb = torch.cat([torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)], dim=1)[None, :, :]
+
+        assert self.num_tokens == 1, 'Assuming one and only one token, [cls]'
+        pe_token = torch.zeros([1, 1, self.embed_dim], dtype=torch.float32)
+        self.pos_embed = nn.Parameter(torch.cat([pe_token, pos_emb], dim=1))
+        self.pos_embed.requires_grad = False
+
+
+class ConvStem(nn.Module):
+    """ 
+    ConvStem, from Early Convolutions Help Transformers See Better, Tete et al. https://arxiv.org/abs/2106.14881
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+
+        assert patch_size == 16, 'ConvStem only supports patch size of 16'
+        assert embed_dim % 8 == 0, 'Embed dimension must be divisible by 8 for ConvStem'
+
+        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
+
+        # build stem, similar to the design in https://arxiv.org/abs/2106.14881
+        stem = []
+        input_dim, output_dim = 3, embed_dim // 8
+        for l in range(4):
+            stem.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=2, padding=1, bias=False))
+            stem.append(nn.BatchNorm2d(output_dim))
+            stem.append(nn.ReLU(inplace=True))
+            input_dim = output_dim
+            output_dim *= 2
+        stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1))
+        self.proj = nn.Sequential(*stem)
+
+        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] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+def vit_small(**kwargs):
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=384, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_base(**kwargs):
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_conv_small(**kwargs):
+    # minus one ViT block
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=384, depth=11, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), embed_layer=ConvStem, **kwargs)
+    model.default_cfg = _cfg()
+    return model
+
+def vit_conv_base(**kwargs):
+    # minus one ViT block
+    model = VisionTransformerMoCo(
+        patch_size=16, embed_dim=768, depth=11, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), embed_layer=ConvStem, **kwargs)
+    model.default_cfg = _cfg()
+    return model

networks/base_model.py

@@ -0,0 +1,78 @@
+import os
+import torch
+import torch.nn as nn
+from torch.nn import init
+from torch.optim import lr_scheduler
+
+
+class BaseModel(nn.Module):
+    def __init__(self, opt):
+        super(BaseModel, self).__init__()
+        self.opt = opt
+        self.total_steps = 0
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        self.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+
+    def save_networks(self, save_filename):
+        save_path = os.path.join(self.save_dir, save_filename)
+
+        # serialize model and optimizer to dict
+        # state_dict = {
+        #     'model': self.model.state_dict(),
+        #     'optimizer' : self.optimizer.state_dict(),
+        #     'total_steps' : self.total_steps,
+        # }
+        if self.opt.fix_backbone:
+            if self.opt.head_type == "attention" or self.opt.head_type == "crossattention":
+                state_dict = self.model.attention_head.state_dict()
+            elif self.opt.head_type == "mlp":
+                # state_dict = {
+                # 'mlp': self.model.mlp.state_dict(),
+                # 'fc' : self.model.fc.state_dict(),
+                # }
+                state_dict = self.model.mlp.state_dict()
+            elif self.opt.head_type == "transformer":
+                state_dict = {
+                                'transformer': self.model.transformer_block.state_dict(),
+                                'fc' : self.model.fc.state_dict(),
+                                }
+            else:
+                state_dict = {
+                "weight": self.model.fc.state_dict()["weight"],
+                "bias": self.model.fc.state_dict()["bias"],}
+        else:
+            state_dict = self.model.state_dict()
+
+        torch.save(state_dict, save_path)
+
+
+    def eval(self):
+        self.model.eval()
+
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+
+
+def init_weights(net, init_type='normal', gain=0.02):
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print('initialize network with %s' % init_type)
+    net.apply(init_func)

networks/lpf.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2019, Adobe Inc. All rights reserved.
+#
+# This work is licensed under the Creative Commons Attribution-NonCommercial-ShareAlike
+# 4.0 International Public License. To view a copy of this license, visit
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.
+
+import torch
+import torch.nn.parallel
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from IPython import embed
+
+class Downsample(nn.Module):
+    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+        super(Downsample, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)), int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2))]
+        self.pad_sizes = [pad_size+pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride-1)/2.)
+        self.channels = channels
+
+        # print('Filter size [%i]'%filt_size)
+        if(self.filt_size==1):
+            a = np.array([1.,])
+        elif(self.filt_size==2):
+            a = np.array([1., 1.])
+        elif(self.filt_size==3):
+            a = np.array([1., 2., 1.])
+        elif(self.filt_size==4):    
+            a = np.array([1., 3., 3., 1.])
+        elif(self.filt_size==5):    
+            a = np.array([1., 4., 6., 4., 1.])
+        elif(self.filt_size==6):    
+            a = np.array([1., 5., 10., 10., 5., 1.])
+        elif(self.filt_size==7):    
+            a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+        filt = torch.Tensor(a[:,None]*a[None,:])
+        filt = filt/torch.sum(filt)
+        self.register_buffer('filt', filt[None,None,:,:].repeat((self.channels,1,1,1)))
+
+        self.pad = get_pad_layer(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if(self.filt_size==1):
+            if(self.pad_off==0):
+                return inp[:,:,::self.stride,::self.stride]    
+            else:
+                return self.pad(inp)[:,:,::self.stride,::self.stride]
+        else:
+            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+def get_pad_layer(pad_type):
+    if(pad_type in ['refl','reflect']):
+        PadLayer = nn.ReflectionPad2d
+    elif(pad_type in ['repl','replicate']):
+        PadLayer = nn.ReplicationPad2d
+    elif(pad_type=='zero'):
+        PadLayer = nn.ZeroPad2d
+    else:
+        print('Pad type [%s] not recognized'%pad_type)
+    return PadLayer
+
+
+class Downsample1D(nn.Module):
+    def __init__(self, pad_type='reflect', filt_size=3, stride=2, channels=None, pad_off=0):
+        super(Downsample1D, self).__init__()
+        self.filt_size = filt_size
+        self.pad_off = pad_off
+        self.pad_sizes = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))]
+        self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]
+        self.stride = stride
+        self.off = int((self.stride - 1) / 2.)
+        self.channels = channels
+
+        # print('Filter size [%i]' % filt_size)
+        if(self.filt_size == 1):
+            a = np.array([1., ])
+        elif(self.filt_size == 2):
+            a = np.array([1., 1.])
+        elif(self.filt_size == 3):
+            a = np.array([1., 2., 1.])
+        elif(self.filt_size == 4):
+            a = np.array([1., 3., 3., 1.])
+        elif(self.filt_size == 5):
+            a = np.array([1., 4., 6., 4., 1.])
+        elif(self.filt_size == 6):
+            a = np.array([1., 5., 10., 10., 5., 1.])
+        elif(self.filt_size == 7):
+            a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+        filt = torch.Tensor(a)
+        filt = filt / torch.sum(filt)
+        self.register_buffer('filt', filt[None, None, :].repeat((self.channels, 1, 1)))
+
+        self.pad = get_pad_layer_1d(pad_type)(self.pad_sizes)
+
+    def forward(self, inp):
+        if(self.filt_size == 1):
+            if(self.pad_off == 0):
+                return inp[:, :, ::self.stride]
+            else:
+                return self.pad(inp)[:, :, ::self.stride]
+        else:
+            return F.conv1d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])
+
+
+def get_pad_layer_1d(pad_type):
+    if(pad_type in ['refl', 'reflect']):
+        PadLayer = nn.ReflectionPad1d
+    elif(pad_type in ['repl', 'replicate']):
+        PadLayer = nn.ReplicationPad1d
+    elif(pad_type == 'zero'):
+        PadLayer = nn.ZeroPad1d
+    else:
+        print('Pad type [%s] not recognized' % pad_type)
+    return PadLayer

networks/resnet_lpf.py

@@ -0,0 +1,313 @@
+# This code is built from the PyTorch examples repository: https://github.com/pytorch/vision/tree/master/torchvision/models.
+# Copyright (c) 2017 Torch Contributors.
+# The Pytorch examples are available under the BSD 3-Clause License.
+#
+# ==========================================================================================
+#
+# Adobe’s modifications are Copyright 2019 Adobe. All rights reserved.
+# Adobe’s modifications are licensed under the Creative Commons Attribution-NonCommercial-ShareAlike
+# 4.0 International Public License (CC-NC-SA-4.0). To view a copy of the license, visit
+# https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode.
+#
+# ==========================================================================================
+#
+# BSD-3 License
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from .lpf import *
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d']
+
+
+# model_urls = {
+#     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+#     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+#     'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+#     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+#     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+# }
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                 padding=1, groups=groups, bias=False)
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, norm_layer=None, filter_size=1):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1:
+            raise ValueError('BasicBlock only supports groups=1')
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        if(stride==1):
+            self.conv2 = conv3x3(planes,planes)
+        else:
+            self.conv2 = nn.Sequential(Downsample(filt_size=filter_size, stride=stride, channels=planes),
+                conv3x3(planes, planes),)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, norm_layer=None, filter_size=1):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = conv3x3(planes, planes, groups) # stride moved
+        self.bn2 = norm_layer(planes)
+        if(stride==1):
+            self.conv3 = conv1x1(planes, planes * self.expansion)
+        else:
+            self.conv3 = nn.Sequential(Downsample(filt_size=filter_size, stride=stride, channels=planes),
+                conv1x1(planes, planes * self.expansion))
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, norm_layer=None, filter_size=1, pool_only=True):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        planes = [int(width_per_group * groups * 2 ** i) for i in range(4)]
+        self.inplanes = planes[0]
+
+        if(pool_only):
+            self.conv1 = nn.Conv2d(3, planes[0], kernel_size=7, stride=2, padding=3, bias=False)
+        else:
+            self.conv1 = nn.Conv2d(3, planes[0], kernel_size=7, stride=1, padding=3, bias=False)
+        self.bn1 = norm_layer(planes[0])
+        self.relu = nn.ReLU(inplace=True)
+
+        if(pool_only):
+            self.maxpool = nn.Sequential(*[nn.MaxPool2d(kernel_size=2, stride=1), 
+                Downsample(filt_size=filter_size, stride=2, channels=planes[0])])
+        else:
+            self.maxpool = nn.Sequential(*[Downsample(filt_size=filter_size, stride=2, channels=planes[0]), 
+                nn.MaxPool2d(kernel_size=2, stride=1), 
+                Downsample(filt_size=filter_size, stride=2, channels=planes[0])])
+
+        self.layer1 = self._make_layer(block, planes[0], layers[0], groups=groups, norm_layer=norm_layer)
+        self.layer2 = self._make_layer(block, planes[1], layers[1], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+        self.layer3 = self._make_layer(block, planes[2], layers[2], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+        self.layer4 = self._make_layer(block, planes[3], layers[3], stride=2, groups=groups, norm_layer=norm_layer, filter_size=filter_size)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(planes[3] * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                if(m.in_channels!=m.out_channels or m.out_channels!=m.groups or m.bias is not None):
+                    # don't want to reinitialize downsample layers, code assuming normal conv layers will not have these characteristics
+                    nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                else:
+                    print('Not initializing')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, groups=1, norm_layer=None, filter_size=1):
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            # downsample = nn.Sequential(
+            #     conv1x1(self.inplanes, planes * block.expansion, stride, filter_size=filter_size),
+            #     norm_layer(planes * block.expansion),
+            # )
+
+            downsample = [Downsample(filt_size=filter_size, stride=stride, channels=self.inplanes),] if(stride !=1) else []
+            downsample += [conv1x1(self.inplanes, planes * block.expansion, 1),
+                norm_layer(planes * block.expansion)]
+            # print(downsample)
+            downsample = nn.Sequential(*downsample)
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, groups, norm_layer, filter_size=filter_size))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=groups, norm_layer=norm_layer, filter_size=filter_size))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+def resnet18(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], filter_size=filter_size, pool_only=pool_only, **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+    return model
+
+
+def resnet34(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    """Constructs a ResNet-34 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+    return model
+
+
+def resnet50(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def resnet101(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
+
+
+def resnet152(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    """Constructs a ResNet-152 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], filter_size=filter_size, pool_only=pool_only, **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model
+
+
+def resnext50_32x4d(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    model = ResNet(Bottleneck, [3, 4, 6, 3], groups=4, width_per_group=32, filter_size=filter_size, pool_only=pool_only, **kwargs)
+    # if pretrained:
+    #     model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def resnext101_32x8d(pretrained=False, filter_size=1, pool_only=True, **kwargs):
+    model = ResNet(Bottleneck, [3, 4, 23, 3], groups=8, width_per_group=32, filter_size=filter_size, pool_only=pool_only, **kwargs)
+    # if pretrained:
+    #     model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model

networks/trainer.py

@@ -0,0 +1,107 @@
+import functools
+import torch
+import torch.nn as nn
+from networks.base_model import BaseModel, init_weights
+import sys
+from models import get_model
+
+class Trainer(BaseModel):
+    def name(self):
+        return 'Trainer'
+
+    def __init__(self, opt):
+        super(Trainer, self).__init__(opt)
+        self.opt = opt  
+        self.model = get_model(opt)
+        print(f"using {self.model.__class__.__name__}")
+
+        if opt.head_type == "fc":
+            torch.nn.init.normal_(self.model.fc.weight.data, 0.0, opt.init_gain)
+        elif opt.head_type == "attention":
+            for name, params in self.model.attention_head.named_parameters():
+                torch.nn.init.normal_(params, 0.0, opt.init_gain)
+        
+        if opt.resume_path is not None:
+            state_dict = torch.load(opt.resume_path)
+            if self.opt.fix_backbone:
+                if self.opt.head_type == "attention" or opt.head_type == "crossattention":
+                    self.model.attention_head.load_state_dict(state_dict)
+                else:
+                    self.model.fc.load_state_dict(state_dict)
+            else:
+                self.model.load_state_dict(state_dict)
+
+        if opt.fix_backbone:
+            params = []
+            if opt.head_type == "fc":
+                for name, p in self.model.named_parameters():
+                    if  name=="fc.weight" or name=="fc.bias": 
+                        params.append(p) 
+                    else:
+                        p.requires_grad = False
+            elif opt.head_type == "mlp":
+                for p in self.model.mlp.parameters():
+                    params.append(p)
+            elif opt.head_type == "attention" or opt.head_type == "crossattention":
+                for p in self.model.attention_head.parameters():
+                    params.append(p)
+
+            elif opt.head_type == "transformer":
+                params = [{'params': self.model.transformer_block.parameters()},
+                {'params': self.model.fc.parameters()}]
+                # params = self.model.transformer.parameters()
+                # params["fc"] = self.model.fc.parameters()
+
+
+        else:
+            print("Your backbone is not fixed. Are you sure you want to proceed? If this is a mistake, enable the --fix_backbone command during training and rerun")
+            import time 
+            time.sleep(3)
+            params = self.model.parameters()
+
+        if opt.optim == 'adam':
+            self.optimizer = torch.optim.AdamW(params, lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=opt.weight_decay)
+        elif opt.optim == 'sgd':
+            self.optimizer = torch.optim.SGD(params, lr=opt.lr, momentum=0.0, weight_decay=opt.weight_decay)
+        else:
+            raise ValueError("optim should be [adam, sgd]")
+
+        self.loss_fn = nn.BCEWithLogitsLoss()
+
+        # self.model = nn.parallel.DistributedDataParallel(self.model)
+        self.model.to(opt.gpu_ids[0])
+
+
+
+    def adjust_learning_rate(self, min_lr=1e-6):
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] /= 10.
+            if param_group['lr'] < min_lr:
+                return False
+        return True
+
+
+    def set_input(self, input):
+        self.input = input[0].to(self.device)
+        self.label = input[1].to(self.device).float()
+
+
+    def forward(self):
+        self.output = self.model(self.input)
+        # self.output = self.output.view(-1).unsqueeze(1)
+        self.output = self.output
+
+
+
+
+    def get_loss(self):
+        return self.loss_fn(self.output.squeeze(1), self.label)
+
+    def optimize_parameters(self):
+        self.forward()
+        self.loss = self.loss_fn(self.output.squeeze(1), self.label) 
+        self.optimizer.zero_grad()
+        self.loss.backward()
+        self.optimizer.step()
+
+

options/base_options.py

@@ -0,0 +1,137 @@
+import argparse
+import os
+import torch
+
+
+class BaseOptions():
+    def __init__(self):
+        self.initialized = False
+
+    def initialize(self, parser):
+
+        parser.add_argument('--head_type', default='attention')
+        parser.add_argument('--no_augment', action='store_true', default=False)
+        parser.add_argument('--patch_base', action='store_true', default=False)
+        parser.add_argument('--shuffle', action='store_true', default=False)
+        parser.add_argument('--shuffle_times', type=int, default=1)
+        parser.add_argument('--original_times', type=int, default=1)
+        parser.add_argument('--patch_size', nargs='+', type=int, default=[14])
+        parser.add_argument('--resume_path', type=str, default=None)
+        parser.add_argument('--load_whole_model', action='store_true', default=False)
+        parser.add_argument('--scale', type=float, default=1.0)
+        
+
+
+        
+
+
+
+
+
+        parser.add_argument('--mode', default='binary')
+        parser.add_argument('--arch', type=str, default='res50', help='see my_models/__init__.py')
+        parser.add_argument('--fix_backbone', action='store_true')  
+
+        # data augmentation
+        parser.add_argument('--rz_interp', default='bilinear')
+        parser.add_argument('--blur_prob', type=float, default=0.5)
+        parser.add_argument('--blur_sig', default='0.0,3.0')
+        parser.add_argument('--jpg_prob', type=float, default=0.5)
+        parser.add_argument('--jpg_method', default='cv2,pil')
+        parser.add_argument('--jpg_qual', default='30,100')
+        
+         
+        parser.add_argument('--real_list_path', default=None, help='only used if data_mode==ours: path for the list of real images, which should contain train.pickle and val.pickle')
+        parser.add_argument('--fake_list_path', default=None, help='only used if data_mode==ours: path for the list of fake images, which should contain train.pickle and val.pickle')
+        parser.add_argument('--wang2020_data_path', default=None, help='only used if data_mode==wang2020 it should contain train and test folders')
+        parser.add_argument('--data_mode',  default='ours', help='wang2020 or ours')
+        parser.add_argument('--data_label', default='train', help='label to decide whether train or validation dataset')
+        parser.add_argument('--weight_decay', type=float, default=0.0, help='loss weight for l2 reg')
+        
+        parser.add_argument('--class_bal', action='store_true') # what is this ?
+        parser.add_argument('--batch_size', type=int, default=256, help='input batch size')
+        parser.add_argument('--loadSize', type=int, default=256, help='scale images to this size')
+        parser.add_argument('--cropSize', type=int, default=224, help='then crop to this size')
+        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
+        parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
+        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
+        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+        parser.add_argument('--serial_batches', action='store_true', help='if true, takes images in order to make batches, otherwise takes them randomly')
+        parser.add_argument('--resize_or_crop', type=str, default='scale_and_crop', help='scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop|none]')
+        parser.add_argument('--no_flip', action='store_true', help='if specified, do not flip the images for data augmentation')
+        parser.add_argument('--init_type', type=str, default='normal', help='network initialization [normal|xavier|kaiming|orthogonal]')
+        parser.add_argument('--init_gain', type=float, default=0.02, help='scaling factor for normal, xavier and orthogonal.')
+        parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{loadSize}')
+        self.initialized = True
+        return parser
+
+    def gather_options(self):
+        # initialize parser with basic options
+        if not self.initialized:
+            parser = argparse.ArgumentParser(
+                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+            parser = self.initialize(parser)
+
+        # get the basic options
+        opt, _ = parser.parse_known_args()
+        self.parser = parser
+
+        return parser.parse_args()
+
+    def print_options(self, opt):
+        message = ''
+        message += '----------------- Options ---------------\n'
+        for k, v in sorted(vars(opt).items()):
+            comment = ''
+            default = self.parser.get_default(k)
+            if v != default:
+                comment = '\t[default: %s]' % str(default)
+            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
+        message += '----------------- End -------------------'
+        print(message)
+
+        # save to the disk
+        expr_dir = os.path.join(opt.checkpoints_dir, opt.name)
+        if not os.path.exists(expr_dir):
+            os.makedirs(expr_dir)
+        file_name = os.path.join(expr_dir, 'opt.txt')
+        with open(file_name, 'wt') as opt_file:
+            opt_file.write(message)
+            opt_file.write('\n')
+
+    def parse(self, print_options=True):
+
+        opt = self.gather_options()
+        opt.isTrain = self.isTrain   # train or test
+
+        # process opt.suffix
+        if opt.suffix:
+            suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
+            opt.name = opt.name + suffix
+
+        if print_options:
+            self.print_options(opt)
+
+        # set gpu ids
+        str_ids = opt.gpu_ids.split(',')
+        opt.gpu_ids = []
+        for str_id in str_ids:
+            id = int(str_id)
+            if id >= 0:
+                opt.gpu_ids.append(id)
+        if len(opt.gpu_ids) > 0:
+            torch.cuda.set_device(opt.gpu_ids[0])
+
+        # additional
+        #opt.classes = opt.classes.split(',')
+        opt.rz_interp = opt.rz_interp.split(',')
+        opt.blur_sig = [float(s) for s in opt.blur_sig.split(',')]
+        opt.jpg_method = opt.jpg_method.split(',')
+        opt.jpg_qual = [int(s) for s in opt.jpg_qual.split(',')]
+        if len(opt.jpg_qual) == 2:
+            opt.jpg_qual = list(range(opt.jpg_qual[0], opt.jpg_qual[1] + 1))
+        elif len(opt.jpg_qual) > 2:
+            raise ValueError("Shouldn't have more than 2 values for --jpg_qual.")
+
+        self.opt = opt
+        return self.opt

options/test_options.py

@@ -0,0 +1,13 @@
+from .base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)
+        parser.add_argument('--model_path')
+        parser.add_argument('--no_resize', action='store_true')
+        parser.add_argument('--no_crop', action='store_true')
+        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
+
+        self.isTrain = False
+        return parser

options/train_options.py

@@ -0,0 +1,22 @@
+from .base_options import BaseOptions
+
+
+class TrainOptions(BaseOptions):
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)
+        parser.add_argument('--earlystop_epoch', type=int, default=5)
+        parser.add_argument('--data_aug', action='store_true', help='if specified, perform additional data augmentation (photometric, blurring, jpegging)')
+        parser.add_argument('--optim', type=str, default='adam', help='optim to use [sgd, adam]')
+        parser.add_argument('--new_optim', action='store_true', help='new optimizer instead of loading the optim state')
+        parser.add_argument('--loss_freq', type=int, default=400, help='frequency of showing loss on tensorboard')
+        parser.add_argument('--save_epoch_freq', type=int, default=1, help='frequency of saving checkpoints at the end of epochs')
+        parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count, we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>, ...')
+        parser.add_argument('--last_epoch', type=int, default=-1, help='starting epoch count for scheduler intialization')
+        parser.add_argument('--train_split', type=str, default='train', help='train, val, test, etc')
+        parser.add_argument('--val_split', type=str, default='val', help='train, val, test, etc')
+        parser.add_argument('--niter', type=int, default=100, help='total epoches')
+        parser.add_argument('--beta1', type=float, default=0.9, help='momentum term of adam')
+        parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate for adam')
+
+        self.isTrain = True
+        return parser

requirements.txt

@@ -0,0 +1,84 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+anyio==4.5.2
+blessed==1.20.0
+brotlipy==0.7.0
+captum==0.7.0
+click==8.1.8
+contourpy==1.1.1
+cycler==0.12.1
+et-xmlfile==1.1.0
+exceptiongroup==1.3.0
+fastapi==0.115.12
+ffmpy==0.5.0
+fonttools==4.43.1
+fsspec==2025.3.0
+ftfy==6.1.1
+gpustat==1.1.1
+gradio==4.44.1
+gradio_client==1.3.0
+h11==0.16.0
+hf-xet==1.1.3
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.32.4
+imageio==2.31.5
+importlib-resources==6.1.0
+joblib==1.3.2
+kiwisolver==1.4.5
+lazy_loader==0.3
+markdown-it-py==3.0.0
+matplotlib==3.7.3
+mdurl==0.1.2
+mkl-fft==1.3.1
+mkl-service==2.4.0
+nvidia-ml-py==12.535.108
+opencv-python==4.8.1.78
+openpyxl==3.0.10
+orjson==3.10.15
+packaging==23.2
+pandas
+Pillow 
+protobuf==4.24.4
+psutil==5.9.5
+pydantic==2.10.6
+pydantic_core==2.27.2
+pydub==0.25.1
+Pygments==2.19.1
+pyparsing==3.1.1
+python-multipart==0.0.20
+pytorchtools==0.0.2
+pytz
+PyWavelets==1.4.1
+PyYAML
+regex==2023.10.3
+requests
+rich==14.0.0
+ruff==0.11.13
+scikit-image==0.21.0
+scikit-learn==1.3.1
+scipy==1.10.1
+semantic-version==2.10.0
+shellingham==1.5.4
+six
+sniffio==1.3.1
+starlette==0.44.0
+sympy 
+tensorboardX==2.6.2.2
+threadpoolctl==3.2.0
+tifffile==2023.7.10
+timm
+tomlkit==0.12.0
+torch==2.1.0
+torchaudio==2.1.0
+torchvision==0.16.0
+tqdm==4.66.1
+triton==2.1.0
+typer==0.16.0
+typing_extensions==4.13.2
+urllib3==2.2.3
+utils==1.0.1
+uvicorn==0.33.0
+wcwidth==0.2.8
+websockets==12.0
+zipp==3.17.0

upload_space.py

@@ -0,0 +1,9 @@
+from huggingface_hub import upload_folder
+
+upload_folder(
+    repo_id="XavierJiezou/face-forgery-detection",  # 你的 Space 名称
+    folder_path="./",  # 当前目录
+    repo_type="space",
+    path_in_repo=".",  # 上传到 Space 的根目录
+    token=True  # 自动从 CLI 登录中读取 token
+)