Audio_ViT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0da1222",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers\n",
    "!pip install datasets\n",
    "!pip install pytorch_lightning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2f66b8",
   "metadata": {},
   "source": [
    "# Standard Vision Transformer\n",
    "\n",
    "Steps -\n",
    "1. Preprocessing\n",
    "2. Training\n",
    "3. Testing and Validation.\n",
    "4. Metrics\n",
    "\n",
    "## Preprocessing Steps:\n",
    "\n",
    "1. Creating the Hugging Face Dataset from the custom Created Mel Spectograms.\n",
    "2. Spliting data into Train, Test and Validation\n",
    "3. Creating the pixel values and Labels\n",
    "\n",
    "### Step 1 and 2:\n",
    "1. Creating the Hugging Face Dataset from the custom Created Mel Spectograms.\n",
    "2. Spliting data into Train, Test and Validation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af628ab8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "The directory at /home/ubuntu/audio_transformer/audio_transformer/data/train doesn't contain any data file",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-98aa44747a94>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mdatasets\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mload_dataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 3\u001b[0;31m \u001b[0mtrain_ds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mload_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"imagefolder\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata_dir\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"audio_transformer/data/train\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msplit\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"train\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      4\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/datasets/load.py\u001b[0m in \u001b[0;36mload_dataset\u001b[0;34m(path, name, data_dir, data_files, split, cache_dir, features, download_config, download_mode, ignore_verifications, keep_in_memory, save_infos, revision, use_auth_token, task, streaming, **config_kwargs)\u001b[0m\n\u001b[1;32m   1713\u001b[0m         \u001b[0mrevision\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mrevision\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1714\u001b[0m         \u001b[0muse_auth_token\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0muse_auth_token\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1715\u001b[0;31m         \u001b[0;34m**\u001b[0m\u001b[0mconfig_kwargs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1716\u001b[0m     )\n\u001b[1;32m   1717\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/datasets/load.py\u001b[0m in \u001b[0;36mload_dataset_builder\u001b[0;34m(path, name, data_dir, data_files, cache_dir, features, download_config, download_mode, revision, use_auth_token, **config_kwargs)\u001b[0m\n\u001b[1;32m   1534\u001b[0m         \u001b[0mdownload_mode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_mode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1535\u001b[0m         \u001b[0mdata_dir\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdata_dir\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1536\u001b[0;31m         \u001b[0mdata_files\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdata_files\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1537\u001b[0m     )\n\u001b[1;32m   1538\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/datasets/load.py\u001b[0m in \u001b[0;36mdataset_module_factory\u001b[0;34m(path, revision, download_config, download_mode, force_local_path, dynamic_modules_path, data_dir, data_files, **download_kwargs)\u001b[0m\n\u001b[1;32m   1193\u001b[0m             \u001b[0mdata_files\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdata_files\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1194\u001b[0m             \u001b[0mdownload_config\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_config\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1195\u001b[0;31m             \u001b[0mdownload_mode\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdownload_mode\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1196\u001b[0m         ).get_module()\n\u001b[1;32m   1197\u001b[0m     \u001b[0;31m# Try locally\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/datasets/load.py\u001b[0m in \u001b[0;36mget_module\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    829\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata_files\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    830\u001b[0m             \u001b[0;32melse\u001b[0m \u001b[0mget_patterns_locally\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata_dir\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresolve\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 831\u001b[0;31m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata_dir\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    832\u001b[0m             \u001b[0;32melse\u001b[0m \u001b[0mget_patterns_locally\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresolve\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    833\u001b[0m         )\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/datasets/data_files.py\u001b[0m in \u001b[0;36mget_patterns_locally\u001b[0;34m(base_path)\u001b[0m\n\u001b[1;32m    291\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0m_get_data_files_patterns\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresolver\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    292\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mFileNotFoundError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 293\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mFileNotFoundError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"The directory at {base_path} doesn't contain any data file\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    294\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    295\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: The directory at /home/ubuntu/audio_transformer/audio_transformer/data/train doesn't contain any data file"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "train_ds = load_dataset(\"imagefolder\", data_dir=\"audio_transformer/data/train\", split=\"train\")\n",
    "\n",
    "\n",
    "# split up training into training + validation\n",
    "splits = train_ds.train_test_split(test_size=0.3)\n",
    "train_val_ds = splits['train']\n",
    "test_ds = splits['test']\n",
    "\n",
    "splits = train_val_ds.train_test_split(test_size=0.1)\n",
    "train_ds = splits['train']\n",
    "val_ds = splits['test']\n",
    "\n",
    "\n",
    "print(len(train_ds),len(test_ds),len(val_ds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "936b5c8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 'air_conditioner',\n",
       " 1: 'car_horn',\n",
       " 2: 'children_playing',\n",
       " 3: 'dog_bark',\n",
       " 4: 'drilling',\n",
       " 5: 'engine_idling',\n",
       " 6: 'gun_shot',\n",
       " 7: 'jackhammer',\n",
       " 8: 'siren',\n",
       " 9: 'street_music'}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "with open('audio_transformer/data/label2id.json') as f:\n",
    "    label2id = json.load(f)\n",
    "\n",
    "with open('audio_transformer/data/id2label.json') as f:\n",
    "    id2label = json.load(f)\n",
    "\n",
    "id2label = {int(key):value for key,value in id2label.items()}\n",
    "id2label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dd65ce61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['image', 'classID', 'class'],\n",
       "    num_rows: 5500\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2c2ff10",
   "metadata": {},
   "source": [
    "### Step 3:\n",
    "We will now preprocess the data. The model requires 2 things: `pixel_values` and `labels`. \n",
    "\n",
    "We will perform data augmentaton **on-the-fly** using HuggingFace Datasets' `set_transform` method (docs can be found [here](https://huggingface.co/docs/datasets/package_reference/main_classes.html?highlight=set_transform#datasets.Dataset.set_transform)). This method is kind of a lazy `map`: the transform is only applied when examples are accessed. This is convenient for tokenizing or padding text, or augmenting images at training time for example, as we will do here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8e8883a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fc30f9a008f45af82039204b5804b41",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/160 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import ViTFeatureExtractor\n",
    "\n",
    "feature_extractor = ViTFeatureExtractor.from_pretrained(\"google/vit-base-patch16-224-in21k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d1b004c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import (CenterCrop, \n",
    "                                    Compose, \n",
    "                                    Normalize, \n",
    "                                    RandomHorizontalFlip,\n",
    "                                    RandomResizedCrop, \n",
    "                                    Resize, \n",
    "                                    ToTensor)\n",
    "\n",
    "normalize = Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std)\n",
    "_train_transforms = Compose(\n",
    "        [\n",
    "            RandomResizedCrop(feature_extractor.size),\n",
    "            RandomHorizontalFlip(),\n",
    "            ToTensor(),\n",
    "            normalize,\n",
    "        ]\n",
    "    )\n",
    "\n",
    "_val_transforms = Compose(\n",
    "        [\n",
    "            Resize(feature_extractor.size),\n",
    "            CenterCrop(feature_extractor.size),\n",
    "            ToTensor(),\n",
    "            normalize,\n",
    "        ]\n",
    "    )\n",
    "\n",
    "def train_transforms(examples):\n",
    "    examples['pixel_values'] = [_train_transforms(image.convert(\"RGB\")) for image in examples['image']]\n",
    "    return examples\n",
    "\n",
    "def val_transforms(examples):\n",
    "    examples['pixel_values'] = [_val_transforms(image.convert(\"RGB\")) for image in examples['image']]\n",
    "    return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9d3de87d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the transforms\n",
    "train_ds.set_transform(train_transforms)\n",
    "val_ds.set_transform(val_transforms)\n",
    "test_ds.set_transform(val_transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "821ccef8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'image': [<PIL.PngImagePlugin.PngImageFile image mode=RGBA size=496x369 at 0x7F63D0A4C550>,\n",
       "  <PIL.PngImagePlugin.PngImageFile image mode=RGBA size=496x369 at 0x7F63D0AA7710>],\n",
       " 'classID': [0, 7],\n",
       " 'class': ['air_conditioner', 'jackhammer'],\n",
       " 'pixel_values': [tensor([[[-0.3255, -0.2706, -0.2706,  ..., -0.4667, -0.6314, -0.6549],\n",
       "           [-0.3255, -0.2706, -0.2706,  ..., -0.4667, -0.6314, -0.6549],\n",
       "           [-0.3255, -0.2706, -0.2706,  ..., -0.4667, -0.6314, -0.6549],\n",
       "           ...,\n",
       "           [ 0.5529,  0.6078,  0.6078,  ...,  0.5608,  0.7020,  0.7176],\n",
       "           [ 0.6157,  0.6706,  0.6706,  ...,  0.5922,  0.6549,  0.6627],\n",
       "           [ 0.7882,  0.8588,  0.8588,  ...,  0.6784,  0.5216,  0.5059]],\n",
       "  \n",
       "          [[-0.8431, -0.8196, -0.8196,  ..., -0.8824, -0.8667, -0.8667],\n",
       "           [-0.8431, -0.8196, -0.8196,  ..., -0.8824, -0.8667, -0.8667],\n",
       "           [-0.8431, -0.8196, -0.8196,  ..., -0.8824, -0.8667, -0.8667],\n",
       "           ...,\n",
       "           [-0.5294, -0.4980, -0.4980,  ..., -0.5216, -0.4510, -0.4431],\n",
       "           [-0.4902, -0.4431, -0.4431,  ..., -0.5059, -0.4745, -0.4667],\n",
       "           [-0.3804, -0.2941, -0.2941,  ..., -0.4588, -0.5373, -0.5451]],\n",
       "  \n",
       "          [[-0.0196, -0.0039, -0.0039,  ..., -0.0745, -0.2314, -0.2549],\n",
       "           [-0.0196, -0.0039, -0.0039,  ..., -0.0745, -0.2314, -0.2549],\n",
       "           [-0.0196, -0.0039, -0.0039,  ..., -0.0745, -0.2314, -0.2549],\n",
       "           ...,\n",
       "           [-0.0902, -0.1137, -0.1137,  ..., -0.0980, -0.1608, -0.1686],\n",
       "           [-0.1216, -0.1529, -0.1529,  ..., -0.1137, -0.1373, -0.1451],\n",
       "           [-0.2078, -0.2549, -0.2549,  ..., -0.1529, -0.0824, -0.0745]]]),\n",
       "  tensor([[[ 0.7961,  0.7961,  0.7961,  ...,  0.8824,  0.8745,  0.8353],\n",
       "           [ 0.8039,  0.8039,  0.8039,  ...,  0.8824,  0.8745,  0.8353],\n",
       "           [ 0.8980,  0.8980,  0.8431,  ...,  0.8431,  0.8431,  0.8275],\n",
       "           ...,\n",
       "           [ 0.9922,  0.9922,  0.9765,  ...,  0.9922,  0.9922,  0.9922],\n",
       "           [ 0.9373,  0.9373,  0.9373,  ...,  0.9059,  0.9137,  0.9373],\n",
       "           [ 0.9294,  0.9294,  0.9294,  ...,  0.8980,  0.9059,  0.9294]],\n",
       "  \n",
       "          [[-0.3725, -0.3725, -0.3725,  ..., -0.2549, -0.2706, -0.3255],\n",
       "           [-0.3569, -0.3569, -0.3647,  ..., -0.2627, -0.2784, -0.3255],\n",
       "           [-0.2314, -0.2314, -0.3098,  ..., -0.3176, -0.3255, -0.3412],\n",
       "           ...,\n",
       "           [ 0.4824,  0.4824,  0.2000,  ...,  0.2627,  0.3098,  0.4824],\n",
       "           [-0.0824, -0.0824, -0.1059,  ..., -0.1843, -0.1608, -0.0824],\n",
       "           [-0.1373, -0.1373, -0.1373,  ..., -0.2314, -0.2078, -0.1373]],\n",
       "  \n",
       "          [[-0.2157, -0.2157, -0.2157,  ..., -0.2627, -0.2549, -0.2392],\n",
       "           [-0.2235, -0.2235, -0.2157,  ..., -0.2627, -0.2549, -0.2392],\n",
       "           [-0.2706, -0.2706, -0.2392,  ..., -0.2471, -0.2471, -0.2314],\n",
       "           ...,\n",
       "           [ 0.0196,  0.0196, -0.1529,  ..., -0.1373, -0.1059,  0.0196],\n",
       "           [-0.2471, -0.2471, -0.2706,  ..., -0.2549, -0.2549, -0.2471],\n",
       "           [-0.2784, -0.2784, -0.2784,  ..., -0.2706, -0.2706, -0.2784]]])]}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds[:2]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d4f037a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "\n",
    "def collate_fn(examples):\n",
    "    pixel_values = torch.stack([example[\"pixel_values\"] for example in examples])\n",
    "    labels = torch.tensor([example[\"classID\"] for example in examples])\n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "train_batch_size = 2\n",
    "eval_batch_size = 2\n",
    "\n",
    "train_dataloader = DataLoader(train_ds, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)\n",
    "val_dataloader = DataLoader(val_ds, collate_fn=collate_fn, batch_size=eval_batch_size)\n",
    "test_dataloader = DataLoader(test_ds, collate_fn=collate_fn, batch_size=eval_batch_size)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "850a223c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pixel_values torch.Size([2, 3, 224, 224])\n",
      "labels torch.Size([2])\n"
     ]
    }
   ],
   "source": [
    "batch = next(iter(train_dataloader))\n",
    "for k,v in batch.items():\n",
    "  if isinstance(v, torch.Tensor):\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "58997767",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert batch['pixel_values'].shape == (train_batch_size, 3, 224, 224)\n",
    "assert batch['labels'].shape == (train_batch_size,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "64ce31da",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 3, 224, 224])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(val_dataloader))['pixel_values'].shape\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a51acd",
   "metadata": {},
   "source": [
    "## Training:\n",
    "1. Defining Model\n",
    "2. Loading the Pretrained Model\n",
    "3. Actual Training.\n",
    "\n",
    "## Define the model:\n",
    "\n",
    "Here we define a `LightningModule`, which is very similar to a regular `nn.Module`, but with some additional functionalities.\n",
    "\n",
    "The model itself uses a linear layer on top of a pre-trained `ViTModel`. We place a linear layer on top of the last hidden state of the [CLS] token, which serves as a good representation of an entire image. We also add dropout for regularization.\n",
    "\n",
    "A resource that helped me in understanding PyTorch Lightning is the [documentation](https://pytorch-lightning.readthedocs.io/en/latest/index.html) as well as the [tutorial notebooks](https://github.com/PyTorchLightning/pytorch-lightning/tree/master/notebooks)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4ad3fc9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl\n",
    "from transformers import ViTForImageClassification, AdamW\n",
    "import torch.nn as nn\n",
    "\n",
    "class ViTLightningModule(pl.LightningModule):\n",
    "    def __init__(self, num_labels=10):\n",
    "        super(ViTLightningModule, self).__init__()\n",
    "        self.vit = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224-in21k',\n",
    "                                                              num_labels=10,\n",
    "                                                              id2label=id2label,\n",
    "                                                              label2id=label2id)\n",
    "\n",
    "    def forward(self, pixel_values):\n",
    "        outputs = self.vit(pixel_values=pixel_values)\n",
    "        return outputs.logits\n",
    "        \n",
    "    def common_step(self, batch, batch_idx):\n",
    "        pixel_values = batch['pixel_values']\n",
    "        labels = batch['labels']\n",
    "        logits = self(pixel_values)\n",
    "\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "        loss = criterion(logits, labels)\n",
    "        predictions = logits.argmax(-1)\n",
    "        correct = (predictions == labels).sum().item()\n",
    "        accuracy = correct/pixel_values.shape[0]\n",
    "\n",
    "        return loss, accuracy\n",
    "      \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
    "        # logs metrics for each training_step,\n",
    "        # and the average across the epoch\n",
    "        self.log(\"training_loss\", loss)\n",
    "        self.log(\"training_accuracy\", accuracy)\n",
    "\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
    "        self.log(\"validation_loss\", loss, on_epoch=True)\n",
    "        self.log(\"validation_accuracy\", accuracy, on_epoch=True)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        loss, accuracy = self.common_step(batch, batch_idx)     \n",
    "\n",
    "        return loss\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        # We could make the optimizer more fancy by adding a scheduler and specifying which parameters do\n",
    "        # not require weight_decay but just using AdamW out-of-the-box works fine\n",
    "        return AdamW(self.parameters(), lr=5e-5)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return train_dataloader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return val_dataloader\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return test_dataloader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57a329a0",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "\n",
    "Let's first start up Tensorboard (note that PyTorch Lightning logs to Tensorboard by default):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "77a2b5ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-8df26e520b0b2797\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-8df26e520b0b2797\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Start tensorboard.\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir lightning_logs/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "cce3cb8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57a2697dbfcd456dba1437bdfd35cf26",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/502 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bff3d37862a04e878edfe9a81a2f8921",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/330M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at google/vit-base-patch16-224-in21k were not used when initializing ViTForImageClassification: ['pooler.dense.bias', 'pooler.dense.weight']\n",
      "- This IS expected if you are initializing ViTForImageClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing ViTForImageClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of ViTForImageClassification were not initialized from the model checkpoint at google/vit-base-patch16-224-in21k and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "ename": "MisconfigurationException",
     "evalue": "You requested GPUs: [0]\n But your machine only has: []",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mMisconfigurationException\u001b[0m                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-25-e0a04acca46f>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mViTLightningModule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 14\u001b[0;31m \u001b[0mtrainer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTrainer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgpus\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mEarlyStopping\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmonitor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'validation_loss'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     15\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/pytorch_lightning/trainer/connectors/env_vars_connector.py\u001b[0m in \u001b[0;36minsert_env_defaults\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m         \u001b[0;31m# all args were already moved to kwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 38\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     39\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0minsert_env_defaults\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, logger, checkpoint_callback, enable_checkpointing, callbacks, default_root_dir, gradient_clip_val, gradient_clip_algorithm, process_position, num_nodes, num_processes, devices, gpus, auto_select_gpus, tpu_cores, ipus, log_gpu_memory, progress_bar_refresh_rate, enable_progress_bar, overfit_batches, track_grad_norm, check_val_every_n_epoch, fast_dev_run, accumulate_grad_batches, max_epochs, min_epochs, max_steps, min_steps, max_time, limit_train_batches, limit_val_batches, limit_test_batches, limit_predict_batches, val_check_interval, flush_logs_every_n_steps, log_every_n_steps, accelerator, strategy, sync_batchnorm, precision, enable_model_summary, weights_summary, weights_save_path, num_sanity_val_steps, resume_from_checkpoint, profiler, benchmark, deterministic, reload_dataloaders_every_n_epochs, reload_dataloaders_every_epoch, auto_lr_find, replace_sampler_ddp, detect_anomaly, auto_scale_batch_size, prepare_data_per_node, plugins, amp_backend, amp_level, move_metrics_to_cpu, multiple_trainloader_mode, stochastic_weight_avg, terminate_on_nan)\u001b[0m\n\u001b[1;32m    424\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mTrainerState\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    425\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m         \u001b[0mgpu_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtpu_cores\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parse_devices\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mauto_select_gpus\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtpu_cores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    427\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    428\u001b[0m         \u001b[0;31m# init connectors\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_parse_devices\u001b[0;34m(gpus, auto_select_gpus, tpu_cores)\u001b[0m\n\u001b[1;32m   1541\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1542\u001b[0m         \u001b[0;31m# TODO (@seannaren, @kaushikb11): Include IPU parsing logic here\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1543\u001b[0;31m         \u001b[0mgpu_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdevice_parser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparse_gpu_ids\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgpus\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1544\u001b[0m         \u001b[0mtpu_cores\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdevice_parser\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparse_tpu_cores\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtpu_cores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1545\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mgpu_ids\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtpu_cores\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/pytorch_lightning/utilities/device_parser.py\u001b[0m in \u001b[0;36mparse_gpu_ids\u001b[0;34m(gpus)\u001b[0m\n\u001b[1;32m     87\u001b[0m     \u001b[0m_check_unique\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgpus\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_sanitize_gpu_ids\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mgpus\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/aws_neuron_pytorch_p36/lib/python3.6/site-packages/pytorch_lightning/utilities/device_parser.py\u001b[0m in \u001b[0;36m_sanitize_gpu_ids\u001b[0;34m(gpus)\u001b[0m\n\u001b[1;32m    150\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mgpu\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mall_available_gpus\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    151\u001b[0m             raise MisconfigurationException(\n\u001b[0;32m--> 152\u001b[0;31m                 \u001b[0;34mf\"You requested GPUs: {gpus}\\n But your machine only has: {all_available_gpus}\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    153\u001b[0m             )\n\u001b[1;32m    154\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mgpus\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mMisconfigurationException\u001b[0m: You requested GPUs: [0]\n But your machine only has: []"
     ]
    }
   ],
   "source": [
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import EarlyStopping\n",
    "\n",
    "# for early stopping, see https://pytorch-lightning.readthedocs.io/en/1.0.0/early_stopping.html?highlight=early%20stopping\n",
    "early_stop_callback = EarlyStopping(\n",
    "    monitor='val_loss',\n",
    "    patience=3,\n",
    "    strict=False,\n",
    "    verbose=False,\n",
    "    mode='min'\n",
    ")\n",
    "\n",
    "model = ViTLightningModule()\n",
    "trainer = Trainer(gpus=1, callbacks=[EarlyStopping(monitor='validation_loss')])\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c3f7247",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_aws_neuron_pytorch_p36)",
   "language": "python",
   "name": "conda_aws_neuron_pytorch_p36"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}