vit_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from torch.optim.lr_scheduler import OneCycleLR\n",
    "from torchmetrics.functional import accuracy, precision\n",
    "import torchmetrics.functional as tf\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from pytorch_lightning.callbacks import TQDMProgressBar\n",
    "\n",
    "from data.cifar100 import CIFAR100DataModule\n",
    "from data.cifar10 import CIFAR10DataModule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vision Transformer Demo\n",
    "In this notebook, I present a simple implementation of the Vision Transformer (ViT) from \"An Image is Worth 16x16 Words: Transformers for Image Recognition At Scale.\" I have attempted to represent the essential components as faithfully as possible. This demo is much smaller-scale than the original, since I am training it on CIFAR-100 (with 60,000 images) instead of the 300-million JFT-300M dataset owned by Google.\n",
    "\n",
    "Though originally developed for NLP, the transformer architecture is gradually making its way into many different areas of deep learning, including image classification and labeling and even reinforcement learning. It's an amazingly versatile architecture and very powerful at representing whatever it's being used to model.\n",
    "\n",
    "As part of my effort to understand fundamental architectures and their applications better, I decided to implement the vision transformer (ViT) from the paper¹ directly, without referencing the official codebase. Here, I'll explain how it works (and how my version is implemented). I'll start with a brief review of how transformers work, but I won't get too deep into the weeds here since there are many other excellent guides to transformers (see The Illustrated Transformer for my favorite one). In addition, I'll cover my basic suggestions for using the process of implementing papers for learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set this to whatever folder you wish CIFAR-100 to be downloaded into\n",
    "CIFAR = \"/media/curttigges/project-files/datasets/cifar-100/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image Preparation\n",
    "This function is essential to the functioning of ViT. Essentially, we take images provided by the dataloader and cut them into a series of patches, which are then flattened into a single dimension as shown. This prepares them to be passed through a linear projection layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def img_to_patch(x, patch_size):\n",
    "    '''Transforms image into list of patches of the specified dimensions\n",
    "\n",
    "    Args:\n",
    "        x (Tensor): Tensor of dimensions B x C x H x W, representing a batch.\n",
    "        B=Batch size, C=Channel count.\n",
    "        patch_size (int): Size of one side of (square) patch.\n",
    "\n",
    "    Returns:\n",
    "        patch_seq (Tensor): List of patches of dimension B x N x [C * P ** 2],\n",
    "        where N is the number of patches and P is patch_size.\n",
    "\n",
    "    Notes:\n",
    "        May need to add padding\n",
    "    '''\n",
    "    B, C, H, W = x.shape\n",
    "\n",
    "    # reshape to B x C x H_count x H_patch x W_count x W_patch\n",
    "    x = x.reshape(B, C, H // patch_size, patch_size, W // patch_size, patch_size)\n",
    "    x = x.permute(0, 2, 4, 1, 3, 5)\n",
    "    x = x.flatten(1,2)\n",
    "    x = x.flatten(2, 4)\n",
    "    \n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer Encoder\n",
    "This is a fairly standard transformer encoder layer, which I have implemented myself (instead of using the PyTorch default) for demonstration purposes. Note that the paper uses this norm-first variant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ViTEncoder(nn.Module):\n",
    "    '''Basic transformer encoder, as specified in the paper\n",
    "\n",
    "    Args:\n",
    "        input_dim (int): Dimensions of transformer input (input embed size)\n",
    "        hidden_dim (int): Size of MLP head\n",
    "        num_heads (int): Number of self-attention heads\n",
    "        dropout (float): Probability of dropout\n",
    "    '''\n",
    "    def __init__(self, input_dim, hidden_dim, num_heads, dropout=0.0):\n",
    "        super().__init__()\n",
    "        self.norm1 = nn.LayerNorm(input_dim)\n",
    "        self.attn = nn.MultiheadAttention(input_dim, num_heads)\n",
    "        self.norm2 = nn.LayerNorm(input_dim)\n",
    "        self.fc1 = nn.Linear(input_dim, hidden_dim)\n",
    "        self.act = nn.GELU()\n",
    "        self.fc2 = nn.Linear(hidden_dim, input_dim)\n",
    "        self.drop1 = nn.Dropout(dropout)\n",
    "        self.drop2 = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.norm1(x)\n",
    "        out, _ = self.attn(out, out, out)\n",
    "        \n",
    "        # First residual connection\n",
    "        resid = x + out\n",
    "\n",
    "        # Pass through MLP layer\n",
    "        out = self.norm2(resid)\n",
    "        out = F.gelu(self.fc1(out))\n",
    "        out = self.drop1(out)\n",
    "        out = self.fc2(out)\n",
    "        out = self.drop2(out)\n",
    "\n",
    "        # Second residual connection\n",
    "        out = out + resid\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complete Model\n",
    "Here we can see the entirety of the model, which is surprisingly simple. It simply does the following in order:\n",
    "- Input is cut into patches and flattened.\n",
    "- Input is then passed through a linear projection layer and an activation to create an embedding.\n",
    "- A learnable class embedding is concatenated to the embedding.\n",
    "- Learnable positional embeddings are added to the result from above.\n",
    "- The result is transposed and sent through the transformer, which consists of a number of encoders.\n",
    "- Finally, only the class embedding is taken from the transformer output, and this is passed through a linear classification head (with two layers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ViTClassifier(nn.Module):\n",
    "    '''Encoder-only vision transformer\n",
    "\n",
    "    Args:\n",
    "        embed_dim (int): Size of embedding output from linear projection layer\n",
    "        hidden_dim (int): Size of MLP head\n",
    "        class_head_dim (int): Size of classification head\n",
    "        num_encoders (int): Number of encoder layers\n",
    "        num_heads (int): Number of self-attention heads\n",
    "        patch_size (int): Size of patches\n",
    "        num_patches (int): Total count of patches (patch sequence size) \n",
    "        dropout (float): Probability of dropout\n",
    "    '''\n",
    "    def __init__(\n",
    "        self, embed_size, hidden_size, class_head_dim, num_encoders, \n",
    "        num_heads, patch_size, num_patches, dropout):\n",
    "        super().__init__()\n",
    "\n",
    "        # Key parameters\n",
    "        self.patch_size = patch_size\n",
    "        self.num_patches = num_patches\n",
    "\n",
    "        # Initial projection of flattened patches into an embedding\n",
    "        self.input = nn.Linear(3*(patch_size**2), embed_size)\n",
    "        self.drop = nn.Dropout(dropout)\n",
    "\n",
    "        # Transformer with arbitrary number of encoders, heads, and hidden size\n",
    "        self.transformer = nn.Sequential(\n",
    "            *(ViTEncoder(embed_size, hidden_size, num_heads, dropout) for _ in range(num_encoders))\n",
    "        )\n",
    "        \n",
    "        # Classification head\n",
    "        self.fc1 = nn.Linear(embed_size, class_head_dim)\n",
    "        self.fc2 = nn.Linear(class_head_dim, 100)\n",
    "\n",
    "        # Learnable parameters for class and position embedding\n",
    "        self.class_embed = nn.Parameter(torch.randn(1, 1, embed_size))\n",
    "        self.pos_embed = nn.Parameter(torch.randn(1, 1 + num_patches, embed_size))\n",
    "\n",
    "    def forward(self, x):\n",
    "        # x will be in the shape B x N x C x P x P\n",
    "        x = img_to_patch(x, self.patch_size)       \n",
    "\n",
    "        # pass input through projection layer; shape is B x N x (C * P**2)\n",
    "        x = F.relu(self.input(x))\n",
    "        B, N, L = x.shape\n",
    "\n",
    "        # concatenate class embedding and add positional encoding\n",
    "        class_embed = self.class_embed.repeat(B, 1, 1)\n",
    "        x = torch.cat([class_embed, x], dim=1)\n",
    "        x = x + self.pos_embed[:, :N+1]\n",
    "        x = self.drop(x)\n",
    "\n",
    "        # apply transformer\n",
    "        x = x.transpose(0, 1) # result is N x B x (C * P**2)\n",
    "        x = self.transformer(x)\n",
    "        x = x[0] # grab the class embedding\n",
    "        \n",
    "        # pass through classification head\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "We use PyTorch Lightning for training, as this greatly simplifies and organizes the relevant code. The model is trained with the Adam optimizer with β1 = 0.9, β2 = 0.999, and a weight decay of 0.1. I use OneCycleLR with cosine decay to optimize the learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ViTTrainModule(pl.LightningModule):\n",
    "    '''Encoder-only vision transformer\n",
    "\n",
    "    Args:\n",
    "        embed_dim (int): Size of embedding output from linear projection layer\n",
    "        hidden_dim (int): Size of MLP head\n",
    "        class_head_dim (int): Size of classification head\n",
    "        num_encoders (int): Number of encoder layers\n",
    "        num_heads (int): Number of self-attention heads\n",
    "        patch_size (int): Size of patches\n",
    "        num_patches (int): Total count of patches (patch sequence size) \n",
    "        dropout (float): Probability of dropout\n",
    "        batch_size (int): Batch size (used for OneCycleLR)\n",
    "        learning_rate (float): Maximum learning rate\n",
    "        weight_decay (float): Optimizer weight decay\n",
    "    '''\n",
    "    def __init__(\n",
    "        self, \n",
    "        embed_dim, \n",
    "        hidden_dim, \n",
    "        class_head_dim, \n",
    "        num_encoders, \n",
    "        num_heads, \n",
    "        patch_size, \n",
    "        num_patches, \n",
    "        dropout, \n",
    "        batch_size, \n",
    "        learning_rate=0.001,\n",
    "        weight_decay=0.03):\n",
    "        super().__init__()\n",
    "\n",
    "        # Key parameters\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        # Transformer with arbitrary number of encoders, heads, and hidden size\n",
    "        self.model = ViTClassifier(\n",
    "            embed_dim,\n",
    "            hidden_dim,\n",
    "            class_head_dim,\n",
    "            num_encoders,\n",
    "            num_heads,\n",
    "            patch_size,\n",
    "            num_patches,\n",
    "            dropout\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)        \n",
    "        return x\n",
    "\n",
    "    def evaluate(self, batch, stage=None):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "        acc = accuracy(y_hat, y)\n",
    "\n",
    "        category_prec = precision(y_hat, y.type(torch.int), average='macro', num_classes=100)\n",
    "        category_recall = tf.recall(y_hat, y.type(torch.int), average='macro', num_classes=100)\n",
    "        category_f1 = tf.f1_score(y_hat, y.type(torch.int), average='macro', num_classes=100)\n",
    "\n",
    "        overall_prec = precision(y_hat, y.type(torch.int))\n",
    "        overall_recall = tf.recall(y_hat, y.type(torch.int))\n",
    "        overall_f1 = tf.f1_score(y_hat, y.type(torch.int))\n",
    "\n",
    "        if stage:\n",
    "            self.log(f\"{stage}_loss\", loss, prog_bar=True)\n",
    "            self.log(f\"{stage}_acc\", acc, prog_bar=True)\n",
    "\n",
    "            self.log(f\"{stage}_cat_prec\", category_prec, prog_bar=True)\n",
    "            self.log(f\"{stage}_cat_recall\", category_recall, prog_bar=True)\n",
    "            self.log(f\"{stage}_cat_f1\", category_f1, prog_bar=True)\n",
    "\n",
    "            self.log(f\"{stage}_ovr_prec\", overall_prec, prog_bar=True)\n",
    "            self.log(f\"{stage}_ovr_recall\", overall_recall, prog_bar=True)\n",
    "            self.log(f\"{stage}_ovr_f1\", overall_f1, prog_bar=True)\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self(x)\n",
    "        loss = F.cross_entropy(y_hat, y)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"val\")\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        self.evaluate(batch, \"test\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            self.parameters(), \n",
    "            lr=self.hparams.learning_rate,\n",
    "            betas=(0.9,0.999),\n",
    "            weight_decay=self.hparams.weight_decay)\n",
    "        \n",
    "        steps_per_epoch = 60000 // self.hparams.batch_size\n",
    "   \n",
    "        lr_scheduler_dict = {\n",
    "            \"scheduler\":OneCycleLR(\n",
    "                optimizer,\n",
    "                self.hparams.learning_rate,\n",
    "                epochs=self.trainer.max_epochs,\n",
    "                steps_per_epoch=steps_per_epoch,\n",
    "                anneal_strategy='cos'\n",
    "            ),\n",
    "            \"interval\":\"step\",\n",
    "        }\n",
    "        return {\"optimizer\":optimizer, \"lr_scheduler\":lr_scheduler_dict}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifications\n",
    "Here we specify the size and other attributes of the model. In the paper, the authors trained models of various sizes, but their base model consisted of 12 encoder layers, a hidden (embedding) size of 768, an MLP width of 3072, and 12 heads for the self-attention layer, which came to 86 million parameters in total. For this demo, I initialize a smaller model for ease of training, and because excessively large model don't do well on smaller datasets like CIFAR-100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_kwargs = {\n",
    "    \"embed_dim\":256, \n",
    "    \"hidden_dim\":512,\n",
    "    \"class_head_dim\":512, \n",
    "    \"num_encoders\":24,\n",
    "    \"num_heads\":8,\n",
    "    \"patch_size\":4,\n",
    "    \"num_patches\":64,\n",
    "    \"dropout\":0.1,\n",
    "    \"batch_size\":256,\n",
    "    \"learning_rate\":0.001,\n",
    "    \"weight_decay\":0.03\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialization\n",
    "Here we initialize our data module and model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n"
     ]
    }
   ],
   "source": [
    "cifar100 = CIFAR100DataModule(\n",
    "    batch_size=model_kwargs[\"batch_size\"], \n",
    "    num_workers=12, \n",
    "    data_dir=CIFAR)\n",
    "\n",
    "pl.seed_everything(42)\n",
    "model = ViTTrainModule(**model_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Instrumentation\n",
    "If you have a Weights & Biases account and want to monitor the training progress of the model, you can initialize this with the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m (\u001b[33mascendant\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.12.19 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.17"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/curttigges/projects/vit/wandb/run-20220624_190531-ox64b10f</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/ascendant/vit-cifar100/runs/ox64b10f\" target=\"_blank\">absurd-frog-57</a></strong> to <a href=\"https://wandb.ai/ascendant/vit-cifar100\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: logging graph, to disable use `wandb.watch(log_graph=False)`\n"
     ]
    }
   ],
   "source": [
    "wandb_logger = WandbLogger(project=\"vit-cifar100\")\n",
    "wandb_logger.watch(model, log=\"all\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "Finally, we initialize our trainer and run it. Be sure to comment out the WandB line if not using Weights & Biases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n",
      "Global seed set to 42\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name  | Type          | Params\n",
      "----------------------------------------\n",
      "0 | model | ViTClassifier | 12.9 M\n",
      "----------------------------------------\n",
      "12.9 M    Trainable params\n",
      "0         Non-trainable params\n",
      "12.9 M    Total params\n",
      "51.451    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 179: 100%|██████████| 196/196 [1:41:41<00:00, 31.13s/it, loss=0.456, v_num=b10f, val_loss=3.340, val_acc=0.498, val_cat_prec=0.469, val_cat_recall=0.462, val_cat_f1=0.432, val_ovr_prec=0.498, val_ovr_recall=0.498, val_ovr_f1=0.498]   \n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    max_epochs=180,\n",
    "    accelerator='gpu', \n",
    "    devices=1,\n",
    "    logger=wandb_logger, #comment out if not using WandB\n",
    "    callbacks=[TQDMProgressBar(refresh_rate=10)])\n",
    "    \n",
    "trainer.fit(model, datamodule=cifar100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "As you can see, the model is able to reach an accuracy of only 49.8%, which is far from SOTA performance. However, vision transformers typically only perform better than CNNs on enormous datasets. These are generally impractical for individuals to train (especially since Google has exclusive access to JFT-300M), so I have limited myself to CIFAR.\n",
    "\n",
    "It does perform better on CIFAR-10, as we can see below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_kwargs = {\n",
    "    \"embed_dim\":256, \n",
    "    \"hidden_dim\":512,\n",
    "    \"class_head_dim\":512, \n",
    "    \"num_encoders\":6,\n",
    "    \"num_heads\":8,\n",
    "    \"patch_size\":4,\n",
    "    \"num_patches\":64,\n",
    "    \"dropout\":0.2,\n",
    "    \"batch_size\":256,\n",
    "    \"learning_rate\":0.0003,\n",
    "    \"weight_decay\":0.03\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n"
     ]
    }
   ],
   "source": [
    "CIFAR10 = \"/media/curttigges/project-files/datasets/cifar-10/\"\n",
    "\n",
    "cifar100 = CIFAR10DataModule(\n",
    "    batch_size=model_kwargs[\"batch_size\"], \n",
    "    num_workers=12, \n",
    "    data_dir=CIFAR10)\n",
    "\n",
    "pl.seed_everything(42)\n",
    "model = ViTTrainModule(**model_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m (\u001b[33mascendant\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.12.19 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.12.17"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/home/curttigges/projects/vit/wandb/run-20220626_153609-2fd28mhy</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/ascendant/vit-cifar10/runs/2fd28mhy\" target=\"_blank\">solar-pine-3</a></strong> to <a href=\"https://wandb.ai/ascendant/vit-cifar10\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: logging graph, to disable use `wandb.watch(log_graph=False)`\n"
     ]
    }
   ],
   "source": [
    "wandb_logger = WandbLogger(project=\"vit-cifar10\")\n",
    "wandb_logger.watch(model, log=\"all\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n",
      "Global seed set to 42\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name  | Type          | Params\n",
      "----------------------------------------\n",
      "0 | model | ViTClassifier | 3.4 M \n",
      "----------------------------------------\n",
      "3.4 M     Trainable params\n",
      "0         Non-trainable params\n",
      "3.4 M     Total params\n",
      "13.500    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 179: 100%|██████████| 196/196 [28:17<00:00,  8.66s/it, loss=0.512, v_num=8mhy, val_loss=0.771, val_acc=0.788, val_cat_prec=0.786, val_cat_recall=0.786, val_cat_f1=0.782, val_ovr_prec=0.788, val_ovr_recall=0.788, val_ovr_f1=0.788]  \n"
     ]
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    max_epochs=180,\n",
    "    accelerator='gpu', \n",
    "    devices=1,\n",
    "    logger=wandb_logger, #comment out if not using WandB\n",
    "    callbacks=[TQDMProgressBar(refresh_rate=10)])\n",
    "    \n",
    "trainer.fit(model, datamodule=cifar100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On CIFAR-10, we get an accuracy of 78.8%. Not bad, but not as good as the 9x% that is SOTA. We will need a much, much larger dataset to exceed CNN models. Nevertheless, this illustrates that the basic approach works, and that transformers can be surprisingly good for computer vision!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "af51f29a878ae3a8f8e9f6c4ebe8e1dfd1996d87171925dc7d5b6a703abf9b9c"
  },
  "kernelspec": {
   "display_name": "Python 3.7.10 ('pytorch-dl')",
   "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.7.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}