MAE Food101 finetune scratch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "44a895cf",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ae0c778",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "561a337a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from transformers import AutoFeatureExtractor\n",
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "import timm.models.vision_transformer\n",
    "import os\n",
    "\n",
    "from tqdm import tqdm\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms\n",
    "from timm.models.vision_transformer import PatchEmbed, Block"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e97b3c",
   "metadata": {},
   "source": [
    "## Food101 Dataset "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb0acb3b",
   "metadata": {},
   "source": [
    "original train set: 75,750\n",
    "\n",
    "original validation set: 25,250\n",
    "\n",
    "total: 101,000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4d2c7525",
   "metadata": {},
   "outputs": [],
   "source": [
    "food101 = load_dataset('food101')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6d876e36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=384x512>,\n",
       " 'label': 6}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "food101['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dd9a1138",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 75750/75750 [05:04<00:00, 248.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean RGB: 0.5449871888617703, 0.4434935563380693, 0.34361316599832514\n",
      "Std RGB: 0.27093838406970966, 0.2734508551865403, 0.2780531622290323\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "total_r = 0\n",
    "total_g = 0\n",
    "total_b = 0\n",
    "total_pixels = 0\n",
    "\n",
    "# Variables for standard deviation\n",
    "sum_squared_r = 0\n",
    "sum_squared_g = 0\n",
    "sum_squared_b = 0\n",
    "\n",
    "indices_to_drop = []\n",
    "\n",
    "for i in tqdm(range(len(food101['train']))):\n",
    "    img = np.array(food101['train'][i]['image'].resize((224, 224))) / 255.\n",
    "    \n",
    "    if img.shape == (224, 224, 3):\n",
    "        total_r += img[:, :, 0].sum()\n",
    "        total_g += img[:, :, 1].sum()\n",
    "        total_b += img[:, :, 2].sum()\n",
    "        \n",
    "        # For standard deviation\n",
    "        sum_squared_r += np.sum(np.square(img[:, :, 0]))\n",
    "        sum_squared_g += np.sum(np.square(img[:, :, 1]))\n",
    "        sum_squared_b += np.sum(np.square(img[:, :, 2]))\n",
    "\n",
    "        total_pixels += 224 * 224\n",
    "    else:\n",
    "        indices_to_drop.append(i)\n",
    "\n",
    "mean_r = total_r / total_pixels\n",
    "mean_g = total_g / total_pixels\n",
    "mean_b = total_b / total_pixels\n",
    "\n",
    "# Compute std for each channel\n",
    "std_r = np.sqrt((sum_squared_r / total_pixels) - (mean_r ** 2))\n",
    "std_g = np.sqrt((sum_squared_g / total_pixels) - (mean_g ** 2))\n",
    "std_b = np.sqrt((sum_squared_b / total_pixels) - (mean_b ** 2))\n",
    "\n",
    "print(f\"Mean RGB: {mean_r}, {mean_g}, {mean_b}\")\n",
    "print(f\"Std RGB: {std_r}, {std_g}, {std_b}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5ce42ea6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(indices_to_drop)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7ebe6c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "food101_mean = np.array([mean_r, mean_g, mean_b])\n",
    "food101_std = np.array([std_r, std_g, std_b])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "177d2824",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform(example_batch):\n",
    "    \"\"\"\n",
    "    reshape the images into 224 * 224\n",
    "    \"\"\"\n",
    "    inputs = {}\n",
    "    \n",
    "    pixel_values = []\n",
    "    labels = []\n",
    "    for i in range(len(example_batch['image'])):\n",
    "        x = example_batch['image'][i]\n",
    "        y = example_batch['label'][i]\n",
    "        if np.array(x.resize((224, 224))).shape == (224, 224, 3):\n",
    "            pixel_values.append(torch.tensor(((np.array(x.resize((224, 224))) / 255. - food101_mean) / food101_std), dtype = torch.float).permute(2, 0, 1))\n",
    "            labels.append(y)\n",
    "\n",
    "    inputs['pixel_values'] = pixel_values\n",
    "    inputs['label'] = labels\n",
    "    return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "91632b41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(224, 224, 3)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(food101['train'][0]['image'].resize((224, 224))).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7888bf3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_food101 = food101.with_transform(transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3ffbf30a",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = processed_food101['train']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9952eaf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices_to_choose = list(set(range(len(train_dataset))) - set(indices_to_drop))\n",
    "filtered_train_dataset = train_dataset.select(indices_to_choose)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e2e5e384",
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_dataset = processed_food101['validation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "65c9c819",
   "metadata": {},
   "outputs": [],
   "source": [
    "indices_to_choose = list(set(range(len(validation_dataset))) - set(indices_to_drop))\n",
    "filtered_valid_dataset = validation_dataset.select(indices_to_choose)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0acaa800",
   "metadata": {},
   "source": [
    "## Masked AutoEnocder "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f77060f",
   "metadata": {},
   "source": [
    "Positional Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0b478a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):\n",
    "    \"\"\"\n",
    "    embed_dim: output dimension for each position\n",
    "    pos: a list of positions to be encoded: size (M,)\n",
    "    out: (M, D)\n",
    "    \"\"\"\n",
    "    \n",
    "    assert embed_dim % 2 == 0\n",
    "    omega = np.arange(embed_dim // 2, dtype = float)\n",
    "    omega /= embed_dim /2.\n",
    "    omega = 1. / 10000 ** omega # (D/2,)\n",
    "    \n",
    "    pos = pos.reshape(-1) # (M,)\n",
    "    out = np.einsum('m, d->md', pos, omega) # (M, D/2), outer product\n",
    "    \n",
    "    emb_sin = np.sin(out) # (M, D/2)\n",
    "    emb_cos = np.cos(out) # (M, D/2)\n",
    "    emb = np.concatenate([emb_sin, emb_cos], axis = 1) # (M, D)\n",
    "    return emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b7c1b0fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):\n",
    "    assert embed_dim % 2 == 0\n",
    "    \n",
    "    # use half of dimensions to encode grid_h\n",
    "    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)\n",
    "    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)\n",
    "    \n",
    "    emb = np.concatenate([emb_h, emb_w], axis = 1)\n",
    "    return emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "20fdfd07",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token = False):\n",
    "    \"\"\"\n",
    "    grid_size: int of the grid height and width\n",
    "    return:\n",
    "    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)\n",
    "    \"\"\"\n",
    "    grid_h = np.arange(grid_size, dtype = np.float32)\n",
    "    grid_w = np.arange(grid_size, dtype = np.float32)\n",
    "    grid = np.meshgrid(grid_h, grid_w)\n",
    "    grid = np.stack(grid, axis = 0)\n",
    "    \n",
    "    grid = grid.reshape([2, 1, grid_size, grid_size])\n",
    "    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)\n",
    "    if cls_token:\n",
    "        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis = 0)\n",
    "    return pos_embed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b04896fd",
   "metadata": {},
   "source": [
    "Model Structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d6e23373",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskedAutoencoderViT(nn.Module):\n",
    "    \"\"\"\n",
    "    Masked Autoencoder with VisionTransformer backbone\n",
    "    \"\"\"\n",
    "    def __init__(self, img_size = 224, patch_size = 16, in_chans = 3,\\\n",
    "                 embed_dim = 1024, depth = 24, num_heads = 16,\\\n",
    "                decoder_embed_dim = 512, decoder_depth = 8, decoder_num_heads = 16,\\\n",
    "                mlp_ratio = 4., norm_layer = nn.LayerNorm, norm_pix_loss = False):\n",
    "        super().__init__()\n",
    "        \n",
    "        # ---------------\n",
    "        # MAE encoder specifics\n",
    "        self.in_chans = in_chans\n",
    "        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)\n",
    "        num_patches = self.patch_embed.num_patches\n",
    "        \n",
    "        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))\n",
    "        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad = False) # fixed sin-cos embedding\n",
    "        \n",
    "        self.blocks = nn.ModuleList([\n",
    "            Block(embed_dim, num_heads, mlp_ratio, qkv_bias = True, norm_layer = norm_layer)\n",
    "            for i in range(depth)\n",
    "        ])\n",
    "        self.norm = norm_layer(embed_dim)\n",
    "        # ---------------\n",
    "        \n",
    "        # ---------------\n",
    "        # MAE decoder specifics\n",
    "        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias = True)\n",
    "        \n",
    "        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))\n",
    "        \n",
    "        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad = False)\n",
    "        \n",
    "        self.decoder_blocks = nn.ModuleList([\n",
    "            Block(decoder_embed_dim, decoder_num_heads, mlp_ratio, qkv_bias = True, norm_layer = norm_layer)\n",
    "            for i in range(decoder_depth)\n",
    "        ])\n",
    "        \n",
    "        self.decoder_norm = norm_layer(decoder_embed_dim)\n",
    "        self.decoder_pred = nn.Linear(decoder_embed_dim, patch_size**2 * in_chans, bias = True) # decoder to patch\n",
    "        # ---------------\n",
    "        \n",
    "        self.norm_pix_loss = norm_pix_loss\n",
    "        \n",
    "        self.initialize_weights()\n",
    "        \n",
    "    def initialize_weights(self):\n",
    "        # initilization\n",
    "        # initialize (and freeze) pos_embed by sin-cos embedding\n",
    "        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token = True)\n",
    "        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))\n",
    "        \n",
    "        decoder_pos_embed = get_2d_sincos_pos_embed(self.decoder_pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token = True)\n",
    "        self.decoder_pos_embed.data.copy_(torch.from_numpy(decoder_pos_embed).float().unsqueeze(0))\n",
    "        \n",
    "        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)\n",
    "        w = self.patch_embed.proj.weight.data\n",
    "        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))\n",
    "        \n",
    "        # timm's trunc_normal_(std=.02) is effectively normal_(std=.02) as cutoff is too big (2.)\n",
    "        torch.nn.init.normal_(self.cls_token, std=.02)\n",
    "        torch.nn.init.normal_(self.mask_token, std=.02)\n",
    "        \n",
    "        # initialize nn.Linear and nn.LayerNorm\n",
    "        self.apply(self._init_weights)\n",
    "        \n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            # we use xavier_uniform following official JAX ViT:\n",
    "            torch.nn.init.xavier_uniform_(m.weight)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "            \n",
    "    def patchify(self, imgs):\n",
    "        \"\"\"\n",
    "        imgs: (N, 3, H, W)\n",
    "        x: (N, L, patch_size**2 * 3)\n",
    "        \"\"\"\n",
    "        p = self.patch_embed.patch_size[0]\n",
    "        \n",
    "        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0\n",
    "        h = w = imgs.shape[2] // p\n",
    "        x = imgs.reshape(shape = (imgs.shape[0], self.in_chans, h, p, w, p))\n",
    "        x = torch.einsum('nchpwq -> nhpwqc', x)\n",
    "        x = x.reshape(shape = (imgs.shape[0], h * w, p**2 * self.in_chans))\n",
    "        return x\n",
    "    \n",
    "    def unpatchify(self, x):\n",
    "        \"\"\"\n",
    "        x: (N, L, patch_size**2 * 3)\n",
    "        imgs: (N, 3, H, W)\n",
    "        \"\"\"\n",
    "        p = self.patch_embed.patch_size[0]\n",
    "        h = w = int(x.shape[1]**.5)\n",
    "        assert h * w == x.shape[1]\n",
    "        \n",
    "        x = x.reshape(shape = (x.shape[0], h, w, p, p, self.in_chans))\n",
    "        x = torch.einsum('nhwpqc -> nchpwq', x)\n",
    "        imgs = x.reshape(shape = (x.shape[0], self.in_chans, h * p, h * p))\n",
    "        return imgs\n",
    "    \n",
    "    def random_masking(self, x, mask_ratio):\n",
    "        \"\"\"\n",
    "        Perform per-sample masking by per-sample shuffling.\n",
    "        Per-sample shuffling is done by argsort random noise.\n",
    "        x: [N, L, D], sequence\n",
    "        \"\"\"\n",
    "        \n",
    "        N, L, D = x.shape # batch, length, dim\n",
    "        len_keep = int(L * (1 - mask_ratio))\n",
    "        \n",
    "        noise = torch.rand(N, L, device = x.device) # noise in [0, 1]\n",
    "        \n",
    "        # sort noise for each sample\n",
    "        ids_shuffle = torch.argsort(noise, dim = 1) # ascend: small is kept, large is removed\n",
    "        ids_restore = torch.argsort(ids_shuffle, dim = 1)\n",
    "        \n",
    "        # keep the first subset\n",
    "        ids_keep = ids_shuffle[:, :len_keep]\n",
    "        x_masked = torch.gather(x, dim = 1, index = ids_keep.unsqueeze(-1).repeat(1, 1, D))\n",
    "        \n",
    "        # generate the binary mask: 0 is kept, 1 is removed\n",
    "        mask = torch.ones([N, L], device = x.device)\n",
    "        mask[:, :len_keep] = 0\n",
    "        # unshuffle to get the binary mask\n",
    "        mask = torch.gather(mask, dim = 1, index = ids_restore)\n",
    "        \n",
    "        return x_masked, mask, ids_restore\n",
    "    \n",
    "    def forward_encoder(self, x, mask_ratio):\n",
    "        # embed patches\n",
    "        x = self.patch_embed(x)\n",
    "        \n",
    "        # add pos embed w/o cls token\n",
    "        x = x + self.pos_embed[:, 1:, :]\n",
    "        \n",
    "        # masking: length -> length * mask_ratio\n",
    "        x, mask, ids_restore = self.random_masking(x, mask_ratio)\n",
    "        \n",
    "        # append cls token\n",
    "        cls_token = self.cls_token + self.pos_embed[:, :1, :]\n",
    "        cls_tokens = cls_token.expand(x.shape[0], -1, -1)\n",
    "        x = torch.cat((cls_tokens, x), dim = 1)\n",
    "        \n",
    "        # apply Transformer blocks\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x)\n",
    "            \n",
    "        x = self.norm(x)\n",
    "        \n",
    "        return x, mask, ids_restore\n",
    "        \n",
    "    def forward_decoder(self, x, ids_restore):\n",
    "        # embed tokens\n",
    "        x = self.decoder_embed(x)\n",
    "        \n",
    "        # append mask tokens to sequence\n",
    "        mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)\n",
    "        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim = 1) # no cls token\n",
    "        x_ = torch.gather(x_, dim = 1, index = ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])) # unshuffle\n",
    "        x = torch.cat([x[:, :1, :], x_], dim = 1) # append cls token\n",
    "        \n",
    "        # add pos embed\n",
    "        x = x + self.decoder_pos_embed\n",
    "        \n",
    "        # apply Transformer blocks\n",
    "        for blk in self.decoder_blocks:\n",
    "            x = blk(x)\n",
    "        x = self.decoder_norm(x)\n",
    "        \n",
    "        # predictor projection\n",
    "        x = self.decoder_pred(x)\n",
    "        \n",
    "        # remove cls token\n",
    "        x = x[:, 1:, :]\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    def forward_loss(self, imgs, pred, mask, batch_mean):\n",
    "        \"\"\"\n",
    "        imgs: [N, 3, H, W]\n",
    "        pred: [N, L, p*p*3]\n",
    "        mask: [N, L], 0 is kept, 1 is removed\n",
    "        \"\"\"\n",
    "        target = self.patchify(imgs)\n",
    "        if self.norm_pix_loss:\n",
    "            mean = target.mean(dim = -1, keepdim = True)\n",
    "            var = target.var(dim = -1, keepdim = True)\n",
    "            target = (target - mean) / (var + 1.e-6) ** .5\n",
    "        \n",
    "        loss = (pred - target) ** 2\n",
    "        loss = loss.mean(dim = -1) # [N, L], mean loss per patch\n",
    "        if batch_mean:\n",
    "            loss = (loss * mask).sum() / mask.sum()\n",
    "        else:\n",
    "            loss = (loss * mask).sum(dim = 1) / mask.sum(dim = 1)\n",
    "        return loss\n",
    "    \n",
    "    def forward_feature(self, imgs):\n",
    "        latent, _, _ = self.forward_encoder(imgs)\n",
    "        return latent[:, 0]\n",
    "    \n",
    "    def forward(self, imgs, mask_ratio = 0.75, batch_mean = True):\n",
    "        latent, mask, ids_restore = self.forward_encoder(imgs, mask_ratio)\n",
    "        pred = self.forward_decoder(latent, ids_restore) # [N, L, p*p*3]\n",
    "        loss = self.forward_loss(imgs, pred, mask, batch_mean)\n",
    "        return loss, self.unpatchify(pred), mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b32f084",
   "metadata": {},
   "source": [
    "Model Build-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "75c123ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mae_vit_base_patch16_dec512d8b(**kwargs):\n",
    "    model = MaskedAutoencoderViT(\n",
    "        patch_size=16, embed_dim=768, depth=12, num_heads=12,\n",
    "        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,\n",
    "        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def mae_vit_large_patch16_dec512d8b(**kwargs):\n",
    "    model = MaskedAutoencoderViT(\n",
    "        patch_size=16, embed_dim=1024, depth=24, num_heads=16,\n",
    "        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,\n",
    "        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)\n",
    "    return model\n",
    "\n",
    "\n",
    "def mae_vit_huge_patch14_dec512d8b(**kwargs):\n",
    "    model = MaskedAutoencoderViT(\n",
    "        patch_size=14, embed_dim=1280, depth=32, num_heads=16,\n",
    "        decoder_embed_dim=512, decoder_depth=8, decoder_num_heads=16,\n",
    "        mlp_ratio=4, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "471af4f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set recommended archs\n",
    "mae_vit_base_patch16 = mae_vit_base_patch16_dec512d8b  # decoder: 512 dim, 8 blocks\n",
    "mae_vit_large_patch16 = mae_vit_large_patch16_dec512d8b  # decoder: 512 dim, 8 blocks\n",
    "mae_vit_huge_patch14 = mae_vit_huge_patch14_dec512d8b  # decoder: 512 dim, 8 blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "162c493d",
   "metadata": {},
   "source": [
    "## Fine-tune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "017076cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.optim import Adam, AdamW\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import torchvision.transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from torchvision.transforms import ToTensor, Compose, Normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "26cdba31",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Finetune_cls(nn.Module):\n",
    "    \n",
    "    def __init__(self, model, d_hidden = 1024, label_classes = 101):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.model = model\n",
    "        self.proj_label_classes = nn.Linear(d_hidden, label_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "            \n",
    "        cls_state = self.model.forward_encoder(x, 0)[0][:, 0, :]\n",
    "        \n",
    "        return self.proj_label_classes(cls_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "20a76bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def finetune_epoch(clf: Finetune_cls,\n",
    "                   train_dataset,\n",
    "                   batch_size=128,\n",
    "                   lr=5e-5,\n",
    "                   device=\"cuda:0\"):\n",
    "    clf.train()\n",
    "    loader = DataLoader(train_dataset, batch_size, drop_last=True, shuffle = True)\n",
    "    \n",
    "    params_to_update = []\n",
    "    for name, param in clf.named_parameters():\n",
    "        if param.requires_grad == True:\n",
    "            params_to_update.append(param)\n",
    "\n",
    "    optimizer = torch.optim.AdamW(params_to_update, lr=lr)\n",
    "    criterion = nn.CrossEntropyLoss().to(device)\n",
    "\n",
    "    loss_list = []\n",
    "    avg_loss = 0\n",
    "    for batch in tqdm(loader, desc=\"Train\"):\n",
    "        # Train loop.\n",
    "        optimizer.zero_grad()\n",
    "        cls_logit = clf(batch['pixel_values'].to(device))\n",
    "        \n",
    "        loss_cls = criterion(cls_logit, batch['label'].to(device).squeeze())\n",
    "        \n",
    "        loss = loss_cls\n",
    "\n",
    "        loss.backward()\n",
    "        \n",
    "        # Have gradients at this point.\n",
    "        nn.utils.clip_grad_norm_(clf.parameters(), max_norm=5.0, norm_type=2)\n",
    "        optimizer.step()\n",
    "        \n",
    "        avg_loss += loss.item()\n",
    "        \n",
    "    return loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6444dbd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def evaluate(clf: Finetune_cls, eval_dataset, batch_size):\n",
    "    clf.eval()\n",
    "    loader = DataLoader(eval_dataset, batch_size=batch_size, drop_last=True)\n",
    "\n",
    "    n_right_classes = 0\n",
    "    n_total = 0\n",
    "\n",
    "    for batch in tqdm(loader, desc=\"Eval\"):\n",
    "        # Compute accuracy.\n",
    "        cls_logit = clf(batch['pixel_values'].to(device))\n",
    "        \n",
    "        pred = cls_logit.argmax(dim=1)\n",
    "        \n",
    "        n_right_classes_batch = sum(pred == batch['label'].to(device)).item()\n",
    "        \n",
    "        n_right_classes += n_right_classes_batch\n",
    "        \n",
    "        n_total += pred.numel()\n",
    "\n",
    "    print(\"  Acc_cls:\", n_right_classes / n_total)\n",
    "\n",
    "    return n_right_classes / n_total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "d0a27f5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def finetune(clf: Finetune_cls, train_dataset, test_dataset, n_epochs: int = 1, model_name=None, **args):\n",
    "    print(\"Using device:\", args[\"device\"])\n",
    "    train = train_dataset\n",
    "\n",
    "    valid = test_dataset\n",
    "    loss = []\n",
    "    acc = []\n",
    "    for epoch in range(n_epochs):\n",
    "        print(f\"Starting epoch {epoch+1}...\")\n",
    "        loss_list = finetune_epoch(clf, train, **args)\n",
    "        loss += loss_list\n",
    "\n",
    "        # Save the final checkpoints of the model\n",
    "        if model_name is not None:\n",
    "            torch.save(clf, model_path + model_name + 'epoch_' + str(epoch+1) + '.pt')\n",
    "\n",
    "        acc_i = evaluate(clf, valid, 32)\n",
    "        acc.append(acc_i)\n",
    "    \n",
    "    return loss, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8bde3c17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n",
      "Starting epoch 1...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [41:22<00:00,  2.10s/it]\n",
      "Eval: 100%|██████████| 789/789 [10:22<00:00,  1.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.19415399239543726\n",
      "Starting epoch 2...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [40:43<00:00,  2.07s/it]\n",
      "Eval: 100%|██████████| 789/789 [10:17<00:00,  1.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.2658032319391635\n",
      "Starting epoch 3...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [26:29<00:00,  1.34s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:37<00:00,  2.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.32180766793409377\n",
      "Starting epoch 4...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [24:07<00:00,  1.22s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:39<00:00,  2.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3465224968314322\n",
      "Starting epoch 5...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [24:07<00:00,  1.22s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:37<00:00,  2.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.36695975918884666\n",
      "Starting epoch 6...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [24:34<00:00,  1.25s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:52<00:00,  2.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.38165399239543724\n",
      "Starting epoch 7...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [29:21<00:00,  1.49s/it]\n",
      "Eval: 100%|██████████| 789/789 [05:34<00:00,  2.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.38620880861850443\n",
      "Starting epoch 8...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [24:18<00:00,  1.23s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:32<00:00,  2.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3827233840304182\n",
      "Starting epoch 9...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:33<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:31<00:00,  2.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3696134347275032\n",
      "Starting epoch 10...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:35<00:00,  1.20s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:36<00:00,  2.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3676330798479088\n",
      "Starting epoch 11...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:39<00:00,  1.20s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:35<00:00,  2.86it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3599096958174905\n",
      "Starting epoch 12...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:29<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:28<00:00,  2.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3493742078580482\n",
      "Starting epoch 13...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:28<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:31<00:00,  2.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3502059569074778\n",
      "Starting epoch 14...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:27<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:29<00:00,  2.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3523051330798479\n",
      "Starting epoch 15...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:25<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:31<00:00,  2.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3514733840304182\n",
      "Starting epoch 16...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:31<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:33<00:00,  2.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3493742078580482\n",
      "Starting epoch 17...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:28<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:29<00:00,  2.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.34897813688212925\n",
      "Starting epoch 18...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:25<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:32<00:00,  2.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3464432826362484\n",
      "Starting epoch 19...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:29<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:27<00:00,  2.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.34521546261089986\n",
      "Starting epoch 20...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train: 100%|██████████| 1183/1183 [23:31<00:00,  1.19s/it]\n",
      "Eval: 100%|██████████| 789/789 [04:32<00:00,  2.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Acc_cls: 0.3458491761723701\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "EPOCHS = 20\n",
    "model = MaskedAutoencoderViT(img_size = 224, in_chans = 3, patch_size = 16, depth = 12, decoder_depth = 4).to(device)\n",
    "classifier = Finetune_cls(model.to(device))\n",
    "classifier.to(device)\n",
    "optimizer = AdamW(model.parameters(), lr = 5e-5, betas=(0.9, 0.95), weight_decay=0.05)\n",
    "\n",
    "loss = finetune(classifier, filtered_train_dataset, filtered_valid_dataset, EPOCHS, batch_size=64, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "edcadd18",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.19415399239543726,\n",
       " 0.2658032319391635,\n",
       " 0.32180766793409377,\n",
       " 0.3465224968314322,\n",
       " 0.36695975918884666,\n",
       " 0.38165399239543724,\n",
       " 0.38620880861850443,\n",
       " 0.3827233840304182,\n",
       " 0.3696134347275032,\n",
       " 0.3676330798479088,\n",
       " 0.3599096958174905,\n",
       " 0.3493742078580482,\n",
       " 0.3502059569074778,\n",
       " 0.3523051330798479,\n",
       " 0.3514733840304182,\n",
       " 0.3493742078580482,\n",
       " 0.34897813688212925,\n",
       " 0.3464432826362484,\n",
       " 0.34521546261089986,\n",
       " 0.3458491761723701]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa473576",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}