notebooks/evaluate_example_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "048e9a63-4c0e-4c4e-8d90-34f2af8049f9",
   "metadata": {},
   "source": [
    "# Evaluate Popular Vision Language Models on SPEC\n",
    "In [our paper](https://arxiv.org/abs/2312.00081), we evaluated four popular VLMs using our SPEC dataset, namely: [CLIP](https://arxiv.org/abs/2103.00020), [BLIP](https://arxiv.org/abs/2201.12086), [FLAVA](https://arxiv.org/abs/2112.04482), and [CoCa](https://arxiv.org/abs/2205.01917). \\\n",
    "This notebook will guide readers to reproduce these results step by step, let's go!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c22764-d4af-4ab9-8460-d6e4dfc79301",
   "metadata": {},
   "source": [
    "## 1. Prepare the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!git clone https://github.com/wjpoom/SPEC.git\n",
    "%cd SPEC\n",
    "!pip install -e . --quiet"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "id": "fb66cde2-cca5-45c1-ac3d-706b3ff4662d",
   "metadata": {},
   "source": [
    "## 2. Import Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c5a66c29-7726-4126-bdeb-0a20a1eeeac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import zipfile\n",
    "import os\n",
    "import torch\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from spec import get_data, get_model\n",
    "from huggingface_hub import hf_hub_download"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8d7636e-5a5e-462e-b99a-45feccf1e3ba",
   "metadata": {},
   "source": [
    "## 3. Prepare the testing dataset\n",
    "We store the data on HuggingFace. Before starting, you need to download and decompress the data as following："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "856550c0-da71-4c09-9533-2df5bd75ba19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify the path to save the downloaded and extracted the data\n",
    "data_root = 'data'\n",
    "# download *.zip files\n",
    "hf_hub_download(repo_id='wjpoom/SPEC', repo_type='dataset', filename='data.zip', local_dir=data_root)\n",
    "# extract *.zip files\n",
    "with zipfile.ZipFile(os.path.join(data_root, 'data.zip'), 'r') as zip_ref:\n",
    "    zip_ref.extractall(os.path.join(data_root))\n",
    "# remove the *.zip files\n",
    "os.remove(os.path.join(data_root, 'data.zip'))\n",
    "print(f'The SPEC dataset is prepared at: {data_root}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "860f7834-e471-4f9a-9d4b-b4a152ffc133",
   "metadata": {},
   "source": [
    "## 4. Let's Evaluate VLMs on SPEC dataset!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06a98b7-2b4d-47d3-a13b-794dca37a403",
   "metadata": {},
   "source": [
    "### 4.1 Evaluate CLIP\n",
    "We use the `ViT/B-32` variant of [CLIP](https://arxiv.org/abs/2103.00020) with weights resumed from the checkpoint release by OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8e687eb6-08a7-47d8-9927-30b174ea7a32",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:07<00:00,  2.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Image2Text Accuracy: 57.00 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:10<00:00,  1.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Text2Image Accuracy: 52.00 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:12<00:00,  2.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Image2Text Accuracy: 27.10 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [00:36<00:00,  1.15s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Text2Image Accuracy: 26.75 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:08<00:00,  2.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Image2Text Accuracy: 44.27 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:22<00:00,  1.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Text2Image Accuracy: 36.27 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:08<00:00,  2.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Image2Text Accuracy: 34.07 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:22<00:00,  1.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Text2Image Accuracy: 32.47 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <count>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [00:43<00:00,  1.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count subset: Image2Text Accuracy: 25.27 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <count>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [02:59<00:00,  2.52s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count subset: Text2Image Accuracy: 23.62 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <absolute_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [00:45<00:00,  1.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_spatial subset: Image2Text Accuracy: 12.64 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <absolute_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [02:55<00:00,  2.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_spatial subset: Text2Image Accuracy: 12.20 %\n",
      "\n",
      "############# finished the evaluation on all selected subsets ###############\n",
      "average of all subset: Image2Text Accuracy: 33.39 %\n",
      "average of all subset: Text2Image Accuracy: 30.55 %\n",
      "result saved to clip_openai_evaluate_result.pth.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "model_cache_dir = 'models/clip' # specify the path to save the downloaded model checkpoint\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model, image_preprocess = get_model(model_name='clip', cache_dir=model_cache_dir, device=device)\n",
    "# load datasets\n",
    "subset_names = ['absolute_size', 'relative_size', 'absolute_spatial', 'relative_spatial', 'existence', 'count']\n",
    "subsets = get_data(data_root=data_root, subset_names=subset_names, image_preprocess=image_preprocess, batch_size=64, num_workers=8)\n",
    "# evaluate\n",
    "result = {}\n",
    "i2t_acc = 0.\n",
    "t2i_acc = 0.\n",
    "subset_num = 0\n",
    "for subset_name, dataloaders in subsets.items():\n",
    "    subset_result = model.evaluate(subset_name=subset_name, dataloaders=dataloaders)\n",
    "    result[subset_name] = subset_result\n",
    "    i2t_acc += subset_result['accuracy']['i2t_accuracy']\n",
    "    t2i_acc += subset_result['accuracy']['t2i_accuracy']\n",
    "    subset_num += 1\n",
    "# print and save results\n",
    "print(f'\\n############# finished the evaluation on all selected subsets ###############')\n",
    "print(f'average of all subset: Image2Text Accuracy: {i2t_acc/subset_num:.2f} %')\n",
    "print(f'average of all subset: Text2Image Accuracy: {t2i_acc/subset_num:.2f} %')\n",
    "out_path = 'results'  # specify the path to save the evaluation results\n",
    "os.makedirs(out_path, exist_ok=True)\n",
    "out_fn = f\"clip_result.pth\"   # specify the filename according to the model you used\n",
    "torch.save(result, os.path.join(out_path, out_fn))\n",
    "print(f'result saved to {out_fn}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0614e38d-646b-4d7e-a6e7-0a6b05307da9",
   "metadata": {},
   "source": [
    "### 4.2 Evaluate BLIP\n",
    "We use the `ViT-B` variant of [BLIP](https://arxiv.org/abs/2201.12086) with weights resumed from the checkpoint released in this [link](https://github.com/salesforce/BLIP), which is finetuned on COCO for image-text retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ec215787-f3ea-4890-af47-7e83abc3a9a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from ~/.cache/blip/blip-coco-base.pth\n",
      "missing keys:\n",
      "[]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:36<00:00,  2.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Image2Text Accuracy: 55.50 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [00:38<00:00,  2.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Text2Image Accuracy: 50.10 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [01:11<00:00,  2.23s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Image2Text Accuracy: 30.65 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [02:17<00:00,  4.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Text2Image Accuracy: 29.60 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:55<00:00,  2.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Image2Text Accuracy: 43.20 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:20<00:00,  3.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Text2Image Accuracy: 43.07 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [00:54<00:00,  2.26s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Image2Text Accuracy: 34.33 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:20<00:00,  3.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Text2Image Accuracy: 33.27 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <count>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [02:44<00:00,  2.32s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count subset: Image2Text Accuracy: 36.87 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <count>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [10:56<00:00,  9.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count subset: Text2Image Accuracy: 37.40 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <absolute_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [02:53<00:00,  2.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_spatial subset: Image2Text Accuracy: 12.07 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <absolute_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [10:56<00:00,  9.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_spatial subset: Text2Image Accuracy: 11.58 %\n",
      "\n",
      "############# finished the evaluation on all selected subsets ###############\n",
      "average of all subset: Image2Text Accuracy: 35.44 %\n",
      "average of all subset: Text2Image Accuracy: 34.17 %\n",
      "result saved to blip_evaluate_result.pth.\n"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "model_cache_dir = 'models/blip' # specify the path to save the downloaded model checkpoint\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model, image_preprocess = get_model(model_name='blip', cache_dir=model_cache_dir, device=device)\n",
    "# load datasets\n",
    "subset_names = ['absolute_size', 'relative_size', 'absolute_spatial', 'relative_spatial', 'existence', 'count']\n",
    "subsets = get_data(data_root=data_root, subset_names=subset_names, image_preprocess=image_preprocess, batch_size=64, num_workers=8)\n",
    "# evaluate\n",
    "result = {}\n",
    "i2t_acc = 0.\n",
    "t2i_acc = 0.\n",
    "subset_num = 0\n",
    "for subset_name, dataloaders in subsets.items():\n",
    "    subset_result = model.evaluate(subset_name=subset_name, dataloaders=dataloaders)\n",
    "    result[subset_name] = subset_result\n",
    "    i2t_acc += subset_result['accuracy']['i2t_accuracy']\n",
    "    t2i_acc += subset_result['accuracy']['t2i_accuracy']\n",
    "    subset_num += 1\n",
    "# print and save results\n",
    "print(f'\\n############# finished the evaluation on all selected subsets ###############')\n",
    "print(f'average of all subset: Image2Text Accuracy: {i2t_acc/subset_num:.2f} %')\n",
    "print(f'average of all subset: Text2Image Accuracy: {t2i_acc/subset_num:.2f} %')\n",
    "out_path = 'results'  # specify the path to save the evaluation results\n",
    "os.makedirs(out_path, exist_ok=True)\n",
    "out_fn = f\"blip_result.pth\"   # specify the filename according to the model you used\n",
    "torch.save(result, os.path.join(out_path, out_fn))\n",
    "print(f'result saved to {out_fn}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33b537d7-3343-4aa5-a632-ef07ed8cfa6d",
   "metadata": {},
   "source": [
    "### 4.3 Evaluate FLAVA\n",
    "We use the `full` version of [FLAVA](https://arxiv.org/abs/2112.04482) with weights resumed from this [checkpoint](https://huggingface.co/facebook/flava-full)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc335f88-09f7-445b-b548-23bac460e052",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [01:14<00:00,  4.65s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Image2Text Accuracy: 57.90 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <existence>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 16/16 [02:07<00:00,  7.99s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "existence subset: Text2Image Accuracy: 51.80 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [02:33<00:00,  4.78s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Image2Text Accuracy: 25.80 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_spatial>: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 32/32 [07:53<00:00, 14.80s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_spatial subset: Text2Image Accuracy: 25.85 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:46<00:00,  4.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Image2Text Accuracy: 37.07 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <absolute_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [04:31<00:00, 11.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absolute_size subset: Text2Image Accuracy: 36.67 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [01:46<00:00,  4.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Image2Text Accuracy: 33.53 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <relative_size>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 24/24 [04:44<00:00, 11.86s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "relative_size subset: Text2Image Accuracy: 33.07 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Image to Text retrieval on <count>: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 71/71 [05:46<00:00,  4.88s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "count subset: Image2Text Accuracy: 14.00 %\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Text to Image retrieval on <count>:   1%|█▍                                                                                                      | 1/71 [01:43<2:00:24, 103.21s/it]"
     ]
    }
   ],
   "source": [
    "# load model\n",
    "model_cache_dir = 'models/flava' # specify the path to save the downloaded model checkpoint\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model, image_preprocess = get_model(model_name='flava', cache_dir=model_cache_dir, device=device)\n",
    "# load datasets\n",
    "subset_names = ['absolute_size', 'relative_size', 'absolute_spatial', 'relative_spatial', 'existence', 'count']\n",
    "subsets = get_data(data_root=data_root, subset_names=subset_names, image_preprocess=image_preprocess, batch_size=64, num_workers=8)\n",
    "# evaluate\n",
    "result = {}\n",
    "i2t_acc = 0.\n",
    "t2i_acc = 0.\n",
    "subset_num = 0\n",
    "for subset_name, dataloaders in subsets.items():\n",
    "    subset_result = model.evaluate(subset_name=subset_name, dataloaders=dataloaders)\n",
    "    result[subset_name] = subset_result\n",
    "    i2t_acc += subset_result['accuracy']['i2t_accuracy']\n",
    "    t2i_acc += subset_result['accuracy']['t2i_accuracy']\n",
    "    subset_num += 1\n",
    "# print and save results\n",
    "print(f'\\n############# finished the evaluation on all selected subsets ###############')\n",
    "print(f'average of all subset: Image2Text Accuracy: {i2t_acc/subset_num:.2f} %')\n",
    "print(f'average of all subset: Text2Image Accuracy: {t2i_acc/subset_num:.2f} %')\n",
    "out_path = 'results'  # specify the path to save the evluation results\n",
    "os.makedirs(out_path, exist_ok=True)\n",
    "out_fn = f\"flava_result.pth\"   # specify the filename according to the model you used\n",
    "torch.save(result, os.path.join(out_path, out_fn))\n",
    "print(f'result saved to {out_fn}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3528ed5-6ab2-456f-b712-0f7250d9d557",
   "metadata": {},
   "source": [
    "### 4.4 Evaluate CoCa\n",
    "We used the `ViT/B-32` variant of [CoCa](https://arxiv.org/abs/2205.01917) model with weights resumed from the [checkpoint](https://github.com/mlfoundations/open_clip) that pretrained on LAION-2B dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fde7225-1399-45a5-b59d-a71921f140be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "model_cache_dir = 'models/coca' # specify the path to save the downloaded model checkpoint\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "model, image_preprocess = get_model(model_name='coca', cache_dir=model_cache_dir, device=device)\n",
    "# load datasets\n",
    "subset_names = ['absolute_size', 'relative_size', 'absolute_spatial', 'relative_spatial', 'existence', 'count']\n",
    "subsets = get_data(data_root=data_root, subset_names=subset_names, image_preprocess=image_preprocess, batch_size=64, num_workers=8)\n",
    "# evaluate\n",
    "result = {}\n",
    "i2t_acc = 0.\n",
    "t2i_acc = 0.\n",
    "subset_num = 0\n",
    "for subset_name, dataloaders in subsets.items():\n",
    "    subset_result = model.evaluate(subset_name=subset_name, dataloaders=dataloaders)\n",
    "    result[subset_name] = subset_result\n",
    "    i2t_acc += subset_result['accuracy']['i2t_accuracy']\n",
    "    t2i_acc += subset_result['accuracy']['t2i_accuracy']\n",
    "    subset_num += 1\n",
    "# print and save results\n",
    "print(f'\\n############# finished the evaluation on all selected subsets ###############')\n",
    "print(f'average of all subset: Image2Text Accuracy: {i2t_acc/subset_num:.2f} %')\n",
    "print(f'average of all subset: Text2Image Accuracy: {t2i_acc/subset_num:.2f} %')\n",
    "out_path = 'results'  # specify the path to save the evluation results\n",
    "os.makedirs(out_path, exist_ok=True)\n",
    "out_fn = f\"coca_result.pth\"   # specify the filename according to the model you used\n",
    "torch.save(result, os.path.join(out_path, out_fn))\n",
    "print(f'result saved to {out_fn}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46eaa47-3b2a-4e14-b366-8a567f7e1e54",
   "metadata": {},
   "source": [
    "## What's Next?\n",
    "Want to test your own visual language model on SPEC? We have provided a [tutorial](https://github.com/wjpoom/SPEC/blob/main/docs/evaluate_custom_model.md) to help evaluate custom models, feel free to have a try."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spec",
   "language": "python",
   "name": "spec"
  },
  "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.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}