Upload 10 files

Get_feature.ipynb

@@ -0,0 +1,395 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1a160d17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from PIL import Image\n",
+    "import torchvision.transforms as T\n",
+    "from glob import glob\n",
+    "import tqdm\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "from scipy.stats import binned_statistic_2d\n",
+    "import umap\n",
+    "from matplotlib.offsetbox import OffsetImage, AnnotationBbox\n",
+    "\n",
+    "# --- 1. MAE Model Definition (Assumes models_mae is available) ---\n",
+    "# NOTE: In a public release, you must ensure 'models_mae.py' or equivalent \n",
+    "# model definitions (e.g., from the official MAE repository) are accessible.\n",
+    "try:\n",
+    "    # This line assumes you have a models_mae.py file defining the MAE architecture\n",
+    "    import models_mae \n",
+    "except ImportError:\n",
+    "    print(\"Warning: 'models_mae.py' not found. Please ensure it's in the path or replace this section with your model definition.\")\n",
+    "    # Define a dummy module for demonstration if models_mae is missing\n",
+    "    class DummyModel:\n",
+    "        def __init__(self):\n",
+    "            pass\n",
+    "        def to(self, device):\n",
+    "            return self\n",
+    "        def load_state_dict(self, state_dict, strict):\n",
+    "            pass\n",
+    "        def eval(self):\n",
+    "            pass\n",
+    "        def forward_encoder(self, x, mask_ratio):\n",
+    "            # Simulate a feature tensor of shape (1, 197, 768) for base model\n",
+    "            B, C, H, W = x.shape\n",
+    "            dummy_features = torch.randn(B, 197, 768) \n",
+    "            return dummy_features, None, None # Features, mask, ids_restore\n",
+    "    \n",
+    "    class DummyModelsMae:\n",
+    "        def mae_vit_base_patch16(self):\n",
+    "            return DummyModel()\n",
+    "\n",
+    "    models_mae = DummyModelsMae()\n",
+    "\n",
+    "\n",
+    "def prepare_model(chkpt_dir, arch='mae_vit_base_patch16'):\n",
+    "    \"\"\"\n",
+    "    Loads the MAE model and its pre-trained weights.\n",
+    "    \"\"\"\n",
+    "    # 1. Instantiate the model architecture\n",
+    "    model = getattr(models_mae, arch)()\n",
+    "    \n",
+    "    # 2. Load the checkpoint\n",
+    "    # Note: Using map_location='cpu' ensures it loads even if a GPU is not available initially.\n",
+    "    checkpoint = torch.load(chkpt_dir, map_location='cpu')\n",
+    "    \n",
+    "    # 3. Clean up the state dictionary keys (e.g., removing 'module.' prefix from DataParallel)\n",
+    "    state_dict = {k.replace('module.', ''): v for k, v in checkpoint['model'].items()}\n",
+    "\n",
+    "    # 4. Move model to GPU (if available) and load weights\n",
+    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "    print(f\"Using device: {device}\")\n",
+    "    model = model.to(device)\n",
+    "    model.load_state_dict(state_dict, strict=False)\n",
+    "    \n",
+    "    # 5. Set model to evaluation mode\n",
+    "    model.eval()\n",
+    "    return model, device\n",
+    "\n",
+    "# --- 2. Image Loading and Preprocessing ---\n",
+    "\n",
+    "def load_image(image_path, scale_size=256, target_size=224):\n",
+    "    \"\"\"\n",
+    "    Loads and preprocesses an image for MAE input.\n",
+    "    \n",
+    "    Note: The normalization values here [0.1689, 0.1536, 0.1516] and [0.1284, 0.0963, 0.1051]\n",
+    "    appear to be custom statistics, likely calculated from a specific dataset\n",
+    "    (like an astronomical image dataset, given the chkpt_dir path).\n",
+    "    Standard ImageNet stats are typically [0.485, 0.456, 0.406] and [0.229, 0.224, 0.225].\n",
+    "    \"\"\"\n",
+    "    transform = T.Compose([\n",
+    "        T.Resize(scale_size, interpolation=T.InterpolationMode.BICUBIC), # Using BICUBIC as specified by '3'\n",
+    "        T.CenterCrop(target_size),\n",
+    "        T.ToTensor(),\n",
+    "        # Custom normalization based on the dataset the MAE was trained on\n",
+    "        T.Normalize(mean=[0.1689, 0.1536, 0.1516], std=[0.1284, 0.0963, 0.1051])\n",
+    "    ])\n",
+    "    image = Image.open(image_path).convert('RGB')\n",
+    "    return transform(image)\n",
+    "\n",
+    "# --- 3. Feature Extraction ---\n",
+    "\n",
+    "def extract_features(model, image_path, device):\n",
+    "    \"\"\"\n",
+    "    Passes an image through the MAE encoder to get the feature representation.\n",
+    "    \n",
+    "    Returns the Class Token (CLS) feature vector.\n",
+    "    \"\"\"\n",
+    "    img = load_image(image_path).unsqueeze(0) # Add batch dimension\n",
+    "    img = img.to(device)\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        # model.forward_encoder returns (features, mask, ids_restore)\n",
+    "        # We set mask_ratio=0 to ensure we encode the full image\n",
+    "        x, _, _ = model.forward_encoder(img.float(), mask_ratio=0) \n",
+    "        \n",
+    "        # x is typically of shape [B, L, D] where B=Batch, L=Num_Tokens (e.g., 197), D=Feature_Dim (e.g., 768)\n",
+    "        # The CLS token is the first token (index 0).\n",
+    "        # We extract the CLS token feature vector x[0, 0, :]\n",
+    "        return x[0, 0, :].cpu().numpy()\n",
+    "\n",
+    "# --- 4. Dimensionality Reduction (UMAP) ---\n",
+    "\n",
+    "def perform_umap(features, n_components=2, random_state=42):\n",
+    "    \"\"\"\n",
+    "    Scales features and applies UMAP to reduce dimensions to 2D.\n",
+    "    \n",
+    "    \"\"\"\n",
+    "    print(\"Scaling features...\")\n",
+    "    scaler = StandardScaler()\n",
+    "    features_scaled = scaler.fit_transform(features)\n",
+    "    \n",
+    "    print(\"Applying UMAP...\")\n",
+    "    reducer = umap.UMAP(n_components=n_components, random_state=random_state)\n",
+    "    return reducer.fit_transform(features_scaled)\n",
+    "\n",
+    "# --- 5. Main Workflow ---\n",
+    "\n",
+    "def embedding_feature(chkpt_dir, image_dir):\n",
+    "    \"\"\"\n",
+    "    Main function to load model, extract features, and perform UMAP.\n",
+    "    \"\"\"\n",
+    "    # 1. Find all image files\n",
+    "    image_path_list = glob(os.path.join(image_dir, \"*.jpg\"))\n",
+    "    print(f\"Number of images to process: {len(image_path_list)}\")\n",
+    "\n",
+    "    # 2. Prepare the model and device\n",
+    "    model, device = prepare_model(chkpt_dir)\n",
+    "\n",
+    "    # 3. Extract features\n",
+    "    features, image_paths = [], []\n",
+    "    print(\"Extracting features...\")\n",
+    "    for image_path in tqdm.tqdm(image_path_list):\n",
+    "        try:\n",
+    "            features.append(extract_features(model, image_path, device))\n",
+    "            image_paths.append(image_path)\n",
+    "        except Exception as e:\n",
+    "            # Handle potential file read errors or other issues\n",
+    "            print(f\"Skipping image {image_path} due to error: {e}\")\n",
+    "            continue\n",
+    "\n",
+    "    features = np.array(features)\n",
+    "    \n",
+    "    # 4. Perform UMAP\n",
+    "    embedding = perform_umap(features)\n",
+    "\n",
+    "    return embedding, image_paths\n",
+    "\n",
+    "# --- 6. Visualization Functions (Grid Plotting) ---\n",
+    "\n",
+    "def plt_image(components, image_paths, save_path='umap_test.png', nx=100, ny=100):\n",
+    "    \"\"\"\n",
+    "    Creates a detailed visualization where each bin in the UMAP space \n",
+    "    is represented by a single, randomly selected image from that bin.\n",
+    "    \n",
+    "    This is often called a 'datamap' or 'tile plot'.\n",
+    "    \"\"\"\n",
+    "    print(f\"Creating a {nx}x{ny} tile plot visualization...\")\n",
+    "    z_emb = components\n",
+    "    iseed = 13579 # Fixed seed for reproducible random selection\n",
+    "\n",
+    "    # 1. Define bounds for the UMAP space\n",
+    "    xmin, xmax = z_emb[:, 0].min(), z_emb[:, 0].max()\n",
+    "    ymin, ymax = z_emb[:, 1].min(), z_emb[:, 1].max()\n",
+    "\n",
+    "    # 2. Define bins\n",
+    "    binx = np.linspace(xmin, xmax, nx + 1)\n",
+    "    biny = np.linspace(ymin, ymax, ny + 1)\n",
+    "\n",
+    "    # 3. Bin the UMAP coordinates (x and y)\n",
+    "    ret = binned_statistic_2d(z_emb[:, 0], z_emb[:, 1], z_emb[:, 1], 'count', \n",
+    "                              bins=[binx, biny], expand_binnumbers=True)\n",
+    "    z_emb_bins = ret.binnumber.T # Transposed bin numbers (ix, iy) for each point\n",
+    "\n",
+    "    inds_lin = np.arange(z_emb.shape[0])\n",
+    "    inds_used = [] # Indices of images selected for plotting\n",
+    "    plotq = []     # Subplot positions for the selected images\n",
+    "\n",
+    "    # 4. Select one image per populated bin\n",
+    "    for ix in tqdm.tqdm(range(nx), desc=\"Selecting images for grid\"):\n",
+    "        for iy in range(ny):\n",
+    "            # Find all image indices (inds) that fall into the current bin (ix, iy)\n",
+    "            # Bin numbers are 1-based, so we check for ix+1 and iy+1\n",
+    "            dm = (z_emb_bins[:, 0] == ix + 1) & (z_emb_bins[:, 1] == iy + 1)\n",
+    "            inds = inds_lin[dm]\n",
+    "\n",
+    "            if len(inds) > 0:\n",
+    "                # Use a fixed seed based on bin location for reproducible random choice\n",
+    "                np.random.seed(ix * ny + iy + iseed) \n",
+    "                ind_plt = np.random.choice(inds)\n",
+    "                \n",
+    "                inds_used.append(ind_plt)\n",
+    "                # Calculate the 1D index for the subplot: (row * num_cols) + col + 1\n",
+    "                plotq.append(iy * nx + ix) # Adjusted from original: ix + iy * nx\n",
+    "\n",
+    "    # 5. Create the plot\n",
+    "    print(f\"Plotting {len(inds_used)} images...\")\n",
+    "    fig = plt.figure(figsize=(20, 20)) \n",
+    "    \n",
+    "    for index, i in enumerate(inds_used):\n",
+    "        image_path = image_paths[i]\n",
+    "        if os.path.exists(image_path): \n",
+    "            try:\n",
+    "                img_jpg = plt.imread(image_path)\n",
+    "                # Subplot index is 1-based. Use plotq[index] + 1\n",
+    "                plt.subplot(nx, ny, plotq[index] + 1) \n",
+    "                plt.xticks([])\n",
+    "                plt.yticks([])\n",
+    "                plt.imshow(img_jpg)\n",
+    "            except Exception as e:\n",
+    "                print(f\"Error loading or plotting image {image_path}: {e}\")\n",
+    "        else:\n",
+    "            print(f\"File does not exist: {image_path}\")\n",
+    "\n",
+    "    plt.suptitle(f\"MAE Features Embedded with UMAP ({len(inds_used)} images in {nx}x{ny} grid)\", fontsize=24)\n",
+    "    plt.tight_layout(rect=[0, 0, 1, 0.98]) # Adjust layout for suptitle\n",
+    "    plt.savefig(save_path)\n",
+    "    print(f\"Visualization saved to {save_path}\")\n",
+    "    plt.close(fig) # Use close(fig) instead of just close() if you are running this in a loop or non-interactive environment\n",
+    "    # Note: plt.show() is commented out for standard command-line use.\n",
+    "\n",
+    "# --- 7. Execution Block ---\n",
+    "\n",
+    "if __name__ == '__main__':\n",
+    "    # Configuration - REPLACE THESE WITH YOUR ACTUAL PATHS\n",
+    "    # The checkpoint directory of the pre-trained MAE model\n",
+    "    chkpt_dir = './ckpt/norm_ckpt/base/weights/best/epoch_777_loss_0.7236/ckpt.pth' \n",
+    "    # The directory containing the images (*.jpg) for feature extraction\n",
+    "    image_dir = './dataset'  \n",
+    "\n",
+    "    if not os.path.exists(chkpt_dir):\n",
+    "        print(f\"Error: Checkpoint file not found at {chkpt_dir}. Please update the chkpt_dir variable.\")\n",
+    "    elif not os.path.exists(image_dir):\n",
+    "        print(f\"Error: Image directory not found at {image_dir}. Please update the image_dir variable.\")\n",
+    "    else:\n",
+    "        # Perform feature extraction and UMAP reduction\n",
+    "        embedding, image_paths = embedding_feature(chkpt_dir, image_dir)\n",
+    "    \n",
+    "        # Create and save the UMAP visualization\n",
+    "        plt_image(embedding, image_paths, save_path='mae_umap_visualization.png')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e5928015",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "import torchvision.transforms as T\n",
+    "# 假设 models_mae 模块包含您的 MAE 模型定义\n",
+    "import models_mae \n",
+    "\n",
+    "# 定义自定义的归一化值（与您的训练保持一致）\n",
+    "CUSTOM_MEAN = [0.1689, 0.1536, 0.1516]\n",
+    "CUSTOM_STD = [0.1284, 0.0963, 0.1051]\n",
+    "TARGET_SIZE = 224 # ViT-Base 默认的输入尺寸\n",
+    "\n",
+    "# --- 1. 模型准备 (与原代码相同) ---\n",
+    "def prepare_model(chkpt_dir, arch='mae_vit_base_patch16'):\n",
+    "    # ... (与原代码 prepare_model 函数完全相同) ...\n",
+    "    try:\n",
+    "        model = getattr(models_mae, arch)()\n",
+    "    except AttributeError:\n",
+    "        print(f\"错误: models_mae 模块中找不到架构 '{arch}'。请确保 models_mae.py 文件存在。\")\n",
+    "        return None, None\n",
+    "        \n",
+    "    checkpoint = torch.load(chkpt_dir, map_location='cpu')\n",
+    "    state_dict = {k.replace('module.', ''): v for k, v in checkpoint['model'].items()}\n",
+    "\n",
+    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "    print(f\"使用的设备: {device}\")\n",
+    "    model = model.to(device)\n",
+    "    \n",
+    "    model.load_state_dict(state_dict, strict=False)\n",
+    "    model.eval()\n",
+    "    return model, device\n",
+    "\n",
+    "# --- 2. 图像加载与预处理 (关键修改部分) ---\n",
+    "def load_and_preprocess_image(image_path, target_size=TARGET_SIZE):\n",
+    "    \"\"\"\n",
+    "    修改后的图像加载函数：\n",
+    "    - 不进行缩放（Resize to scale_size）\n",
+    "    - 不进行中心裁剪（CenterCrop）\n",
+    "    - 而是直接强制缩放到模型的输入尺寸 (target_size x target_size)。\n",
+    "    \"\"\"\n",
+    "    transform = T.Compose([\n",
+    "        # 强制将图像尺寸变为 (target_size, target_size)，可能导致拉伸\n",
+    "        T.Resize((target_size, target_size), interpolation=T.InterpolationMode.BICUBIC),\n",
+    "        T.ToTensor(),\n",
+    "        T.Normalize(mean=CUSTOM_MEAN, std=CUSTOM_STD)\n",
+    "    ])\n",
+    "    \n",
+    "    if not os.path.exists(image_path):\n",
+    "        raise FileNotFoundError(f\"图像文件未找到: {image_path}\")\n",
+    "        \n",
+    "    image = Image.open(image_path).convert('RGB')\n",
+    "    return transform(image)\n",
+    "\n",
+    "# --- 3. 特征提取核心函数 (与原代码相同) ---\n",
+    "def extract_single_feature(model, image_path, device):\n",
+    "    \"\"\"\n",
+    "    提取单个图像的特征向量（CLS Token）。\n",
+    "    \"\"\"\n",
+    "    img_tensor = load_and_preprocess_image(image_path).unsqueeze(0) # [1, C, H, W]\n",
+    "    img_tensor = img_tensor.to(device)\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        x, _, _ = model.forward_encoder(img_tensor.float(), mask_ratio=0) \n",
+    "        feature_vector = x[0, 0, :].cpu().numpy()\n",
+    "        \n",
+    "    return feature_vector\n",
+    "\n",
+    "# --- 4. 示例使用 ---\n",
+    "if __name__ == '__main__':\n",
+    "    # 请替换为您的模型和图像路径\n",
+    "    CHKPT_DIR = './ckpt/norm_ckpt/base/weights/best/epoch_777_loss_0.7236/ckpt.pth' \n",
+    "    SINGLE_IMAGE_PATH = './dataset/example_image_001.jpg' \n",
+    "    \n",
+    "    print(\"--- MAE 单张图像特征提取开始 (无中心裁剪/预缩放) ---\")\n",
+    "\n",
+    "    try:\n",
+    "        # 1. 准备模型\n",
+    "        model, device = prepare_model(CHKPT_DIR)\n",
+    "        if model is None:\n",
+    "            exit()\n",
+    "\n",
+    "        # 2. 提取特征\n",
+    "        feature = extract_single_feature(model, SINGLE_IMAGE_PATH, device)\n",
+    "\n",
+    "        # 3. 结果展示\n",
+    "        print(\"\\n--- 特征提取结果 ---\")\n",
+    "        print(f\"图像路径: {SINGLE_IMAGE_PATH}\")\n",
+    "        print(f\"ViT 输入尺寸: {TARGET_SIZE}x{TARGET_SIZE} (通过强制缩放获得)\")\n",
+    "        print(f\"特征向量形状: {feature.shape}\")\n",
+    "        print(f\"特征向量（前5个值）: {feature[:5]}\")\n",
+    "        print(\"--- 特征提取完成 ---\")\n",
+    "\n",
+    "    except FileNotFoundError as e:\n",
+    "        print(f\"致命错误: {e}\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"发生其他错误: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f641d493",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

models_mae.py

@@ -0,0 +1,242 @@
+import torch
+import torch.nn as nn
+from functools import partial
+from timm.models.vision_transformer import PatchEmbed, Block
+from util.pos_embed import get_2d_sincos_pos_embed
+
+
+class MaskedAutoEncoderViT(nn.Module):
+    """ Masked Autoencoder with VisionTransformer backbone
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, 
+                 embed_dim=1024, depth=24, num_heads=16,
+                 decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+                 mlp_ratio=4.0, norm_layer=nn.LayerNorm, norm_pix_loss=False):
+        super().__init__()
+        
+
+        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)   # fixed sin-cos embedding
+
+        self.blocks = nn.ModuleList([
+            Block(embed_dim, num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+
+        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False)   # fixed sin-cos embedding
+
+        self.decoder_blocks = nn.ModuleList([
+            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias=True, norm_layer=norm_layer)
+            for i in range(decoder_depth)
+        ])
+
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias=True)  # decoder to patch
+
+        self.norm_pix_loss = norm_pix_loss
+
+        self.initialize_weights()
+    
+    def initialize_weights(self):
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
+        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))
+
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        torch.nn.init.normal_(self.mask_token, std=.02)
+
+        self.apply(self._init_weights)
+        
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+    
+    def random_masking(self, 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]
+
+        ids_shuffle = torch.argsort(noise, dim=1)   # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+    
+    def patchify(self, imgs):
+        """
+        imgs: (N, 3, H, W)
+        x: (N, L, patch_size**2 *3)
+        """
+        p = self.patch_embed.patch_size[0]
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
+        x = torch.einsum('nchpwq->nhwpqc', x)
+        x = x.reshape(shape=(imgs.shape[0], h*w, p**2*3))
+        return x
+    
+    def unpatchify(self, x):
+        """
+        x: (N, L, patch_size**2 *3)
+        imgs: (N, 3, H, W)
+        """
+        p = self.patch_embed.patch_size[0]
+        h = w = int(x.shape[1]**0.5)
+        assert h *w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
+        return imgs
+    
+    def forward_encoder(self, x, mask_ratio):
+        x = self.patch_embed(x)
+
+        x = x + self.pos_embed[:, 1:, :]
+
+        x, mask, ids_restore = self.random_masking(x, mask_ratio)
+
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x, mask, ids_restore
+    
+    def forward_decoder(self, x, ids_restore):
+        x = self.decoder_embed(x)
+
+        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
+        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)   # no cls token
+        x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))   # unshuffle
+        x = torch.cat([x[:, :1, :], x_], dim=1)   # append cls token
+
+        x = x + self.decoder_pos_embed
+
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        x = self.decoder_pred(x)
+
+        x = x[:, 1:, :]
+
+        return x
+    
+    def forward_loss(self, imgs, pred, mask):
+        """
+        imgs: [N, 3, H, W]
+        pred: [N, L, p*p*3]
+        mask: [N, L], 0 is keep, 1 is move.
+        """
+        target = self.patchify(imgs)
+        if self.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6)**0.5
+        
+        loss = (pred - target) ** 2
+        loss = loss.mean(dim=-1)   # [N, L], mean loss per patch
+
+        loss = (loss * mask).sum() / mask.sum()   # mean loss on removed patches
+        return loss
+
+    def forward(self, imgs, mask_ratio=0.75):
+        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)
+        pred = self.forward_decoder(latent, ids_restore)   # [N, L, p*p*3]
+        loss = self.forward_loss(imgs, pred, mask)
+        return loss, pred, mask
+    
+
+    def forward_encoder_with_given_mask(self, x, given_patch_mask):
+
+        x = self.patch_embed(x)  # (N, L, D)
+
+        x = x + self.pos_embed[:, 1:, :]  # (N, L, D)
+
+        N, L, D = x.shape
+        noise = torch.rand(N, L, device=x.device)
+
+        mask_float = given_patch_mask.float()
+        ids_shuffle = torch.argsort(mask_float * (noise.max() + 1) + noise, dim=1)  # (N, L)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        len_keep = L - given_patch_mask.sum(dim=1).max().int().item()
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x_masked), dim=1)
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x, given_patch_mask, ids_restore
+    
+    def forward_with_given_mask(self, imgs, given_patch_mask):
+
+        latent, mask, ids_restore = self.forward_encoder_with_given_mask(imgs, given_patch_mask)
+        pred = self.forward_decoder(latent, ids_restore)
+        loss = self.forward_loss(imgs, pred, mask)
+        return loss, pred, mask
+    
+    
+
+
+def mae_vit_base_patch16(**kwargs):
+    model = MaskedAutoEncoderViT(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_large_patch16(**kwargs):
+    model = MaskedAutoEncoderViT(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def mae_vit_huge_patch14(**kwargs):
+    model = MaskedAutoEncoderViT(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16,
+        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,
+        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model

util/lr_sched.py

@@ -0,0 +1,98 @@
+import torch, math
+import torch.nn as nn
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = args.lr * epoch / args.warmup_epochs 
+    else:
+        lr = args.min_lr + (args.lr - args.min_lr) * 0.5 * \
+            (1. + math.cos(math.pi * (epoch - args.warmup_epochs) / (args.epochs - args.warmup_epochs)))
+    for param_group in optimizer.param_groups:
+        if "lr_scale" in param_group:
+            param_group["lr"] = lr * param_group["lr_scale"]
+        else:
+            param_group["lr"] = lr
+    return lr
+
+
+def param_groups_weight_decay(model: nn.Module, weight_decay=1e-5, no_weight_decay_list=()):
+    no_weight_decay_list = set(no_weight_decay_list)
+    decay = []
+    no_decay = []
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if param.ndim <= 1 or name.endswith(".bias") or name in no_weight_decay_list:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+
+    return [
+        {'params': no_decay, 'weight_decay': 0.},
+        {'params': decay, 'weight_decay': weight_decay}]
+
+
+def param_groups_lrd(model, weight_decay=0.05, no_weight_decay_list=[], layer_decay=.75):
+    """
+    Parameter groups for layer-wise lr decay
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
+    """
+    param_group_names = {}
+    param_groups = {}
+
+    num_layers = len(model.blocks) + 1
+
+    layer_scales = list(layer_decay ** (num_layers - i) for i in range(num_layers + 1))
+
+    for n, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+
+        # no decay: all 1D parameters and model specific ones
+        if p.ndim == 1 or n in no_weight_decay_list:
+            g_decay = "no_decay"
+            this_decay = 0.
+        else:
+            g_decay = "decay"
+            this_decay = weight_decay
+            
+        layer_id = get_layer_id_for_vit(n, num_layers)
+        group_name = "layer_%d_%s" % (layer_id, g_decay)
+
+        if group_name not in param_group_names:
+            this_scale = layer_scales[layer_id]
+
+            param_group_names[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+            param_groups[group_name] = {
+                "lr_scale": this_scale,
+                "weight_decay": this_decay,
+                "params": [],
+            }
+
+        param_group_names[group_name]["params"].append(n)
+        param_groups[group_name]["params"].append(p)
+
+    # print("parameter groups: \n%s" % json.dumps(param_group_names, indent=2))
+
+    return list(param_groups.values())
+
+
+def get_layer_id_for_vit(name, num_layers):
+    """
+    Assign a parameter with its layer id
+    Following BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L33
+    """
+    if name in ['cls_token', 'pos_embed']:
+        return 0
+    elif name.startswith('patch_embed'):
+        return 0
+    elif name.startswith('blocks'):
+        return int(name.split('.')[1]) + 1
+    else:
+        return num_layers

util/misc.py

@@ -0,0 +1,40 @@
+import os, shutil
+import torch, math
+
+def colorstr(*input):
+    # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e.  colorstr('blue', 'hello world')
+    *args, string = input if len(input) > 1 else ('blue', 'bold', input[0])  # color arguments, string
+    colors = {'black': '\033[30m',  # basic colors
+              'red': '\033[31m',
+              'green': '\033[32m',
+              'yellow': '\033[33m',
+              'blue': '\033[34m',
+              'magenta': '\033[35m',
+              'cyan': '\033[36m',
+              'white': '\033[37m',
+              'bright_black': '\033[90m',  # bright colors
+              'bright_red': '\033[91m',
+              'bright_green': '\033[92m',
+              'bright_yellow': '\033[93m',
+              'bright_blue': '\033[94m',
+              'bright_magenta': '\033[95m',
+              'bright_cyan': '\033[96m',
+              'bright_white': '\033[97m',
+              'end': '\033[0m',  # misc
+              'bold': '\033[1m',
+              'underline': '\033[4m'}
+    return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
+
+
+def SaveCheckpoint(state, last, last_path, best, best_path, is_best):
+    if os.path.exists(last):
+        shutil.rmtree(last)
+    last_path.mkdir(parents=True, exist_ok=True)
+    torch.save(state, os.path.join(last_path, 'ckpt.pth'))
+
+    if is_best:
+        if os.path.exists(best):
+            shutil.rmtree(best)
+        best_path.mkdir(parents=True, exist_ok=True)
+        torch.save(state, os.path.join(best_path, 'ckpt.pth'))
+        

util/pos_embed.py

@@ -0,0 +1,85 @@
+import numpy as np
+import torch
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed