notebooks/load-jax-weights-vitb16.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74243911",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "796f93bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf \n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from pprint import pformat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "864619a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "from vit.configs import base_config\n",
    "from vit.layers import mha\n",
    "from vit.models import ViTClassifierExtended"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c306012f",
   "metadata": {},
   "source": [
    "## Select the master dataframe from [AugReg paper](https://arxiv.org/abs/2106.10270)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86505e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "with tf.io.gfile.GFile(\"gs://vit_models/augreg/index.csv\") as f:\n",
    "    df = pd.read_csv(f)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715577b7",
   "metadata": {},
   "source": [
    "## Pick a checkpoint\n",
    "\n",
    "**Criteria**\n",
    "\n",
    "* B16 architecture\n",
    "* Resolution 224\n",
    "* Pacth size 16\n",
    "* Best top-1 accuracy on ImageNet-1k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f19ffd6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "b16s = df.query(\n",
    "    'ds==\"i21k\" & adapt_resolution==224 & adapt_ds==\"imagenet2012\" & name==\"B/16\"'\n",
    ").sort_values(\"adapt_final_test\", ascending=False)\n",
    "b16s.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "973cf5ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_b16_i1k_checkpoint = str(b16s.iloc[0][\"adapt_filename\"])\n",
    "b16s.iloc[0][\"adapt_filename\"], b16s.iloc[0][\"adapt_final_test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac493235",
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = best_b16_i1k_checkpoint\n",
    "\n",
    "path = f\"gs://vit_models/augreg/{filename}.npz\"\n",
    "\n",
    "print(f\"{tf.io.gfile.stat(path).length / 1024 / 1024:.1f} MiB - {path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481e588a",
   "metadata": {},
   "source": [
    "## Copy over the checkpoint and load it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33d6353e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !gsutil cp {path} .\n",
    "local_path = path.split(\"//\")[-1].split(\"/\")[-1]\n",
    "local_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3f38c1c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "with open(local_path, \"rb\") as f:\n",
    "    params_jax = np.load(f)\n",
    "    params_jax = dict(zip(params_jax.keys(), params_jax.values()))\n",
    "\n",
    "# print(pformat(list(params_jax.keys())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4dfb22",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_jax[\"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/kernel\"].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c69a90e",
   "metadata": {},
   "source": [
    "## Instantiate a ViT model in TF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b292d89e",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = base_config.get_config()\n",
    "with config.unlocked():\n",
    "    config.num_classes = 1000\n",
    "\n",
    "config.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb69ebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure it works.\n",
    "vit_b16_model = ViTClassifierExtended(config)\n",
    "vit_b16_model(tf.ones((1, 224, 224, 3)))[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b302148",
   "metadata": {},
   "source": [
    "## Copy the projection layer params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de00ba49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Projection.\n",
    "\n",
    "vit_b16_model.layers[0].layers[0].kernel.assign(\n",
    "    tf.Variable(params_jax[\"embedding/kernel\"])\n",
    ")\n",
    "vit_b16_model.layers[0].layers[0].bias.assign(tf.Variable(params_jax[\"embedding/bias\"]))\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e40b35b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[0].layers[0].kernel.numpy(), params_jax[\"embedding/kernel\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10fba4ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[0].layers[0].bias.numpy(), params_jax[\"embedding/bias\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b0d5a28",
   "metadata": {},
   "source": [
    "## Copy the positional embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9a5760",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Positional embedding.\n",
    "\n",
    "vit_b16_model.positional_embedding.assign(\n",
    "    tf.Variable(params_jax[\"Transformer/posembed_input/pos_embedding\"])\n",
    ")\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95922f8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.positional_embedding.numpy(),\n",
    "    params_jax[\"Transformer/posembed_input/pos_embedding\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85b11ac7",
   "metadata": {},
   "source": [
    "## Copy the `cls_token`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e96d7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cls token.\n",
    "\n",
    "vit_b16_model.cls_token.assign(tf.Variable(params_jax[\"cls\"]))\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40705101",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(vit_b16_model.cls_token.numpy(), params_jax[\"cls\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50623ba7",
   "metadata": {},
   "source": [
    "## Copy the final Layer Norm params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36922a95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Final layer norm layer.\n",
    "vit_b16_model.layers[-2].gamma.assign(\n",
    "    tf.Variable(params_jax[\"Transformer/encoder_norm/scale\"])\n",
    ")\n",
    "vit_b16_model.layers[-2].beta.assign(\n",
    "    tf.Variable(params_jax[\"Transformer/encoder_norm/bias\"])\n",
    ")\n",
    "\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3519f8a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[-2].gamma.numpy(), params_jax[\"Transformer/encoder_norm/scale\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "927b71d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[-2].beta.numpy(), params_jax[\"Transformer/encoder_norm/bias\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f14687",
   "metadata": {},
   "source": [
    "## Copy head layer params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bdba3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Head layer.\n",
    "\n",
    "vit_b16_model.layers[-1].kernel.assign(tf.Variable(params_jax[\"head/kernel\"]))\n",
    "vit_b16_model.layers[-1].bias.assign(tf.Variable(params_jax[\"head/bias\"]))\n",
    "print(\" \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37cd4eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[-1].kernel.numpy(), params_jax[\"head/kernel\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35f7c642",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(\n",
    "    vit_b16_model.layers[-1].bias.numpy(), params_jax[\"head/bias\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9d3cec2",
   "metadata": {},
   "source": [
    "## Copy the Transformer params"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd00a40",
   "metadata": {},
   "source": [
    "**Structure of a single Transformer encoder block in the JAX model**:\n",
    "\n",
    "\n",
    "```md\n",
    " 'Transformer/encoderblock_0/LayerNorm_0/bias',\n",
    " 'Transformer/encoderblock_0/LayerNorm_0/scale',\n",
    " 'Transformer/encoderblock_0/LayerNorm_2/bias',\n",
    " 'Transformer/encoderblock_0/LayerNorm_2/scale',\n",
    " 'Transformer/encoderblock_0/MlpBlock_3/Dense_0/bias',\n",
    " 'Transformer/encoderblock_0/MlpBlock_3/Dense_0/kernel',\n",
    " 'Transformer/encoderblock_0/MlpBlock_3/Dense_1/bias',\n",
    " 'Transformer/encoderblock_0/MlpBlock_3/Dense_1/kernel',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/bias',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/kernel',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/bias',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/kernel',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/bias',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/kernel',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/bias',\n",
    " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/kernel',\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e70d7b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def modify_attention_block(tf_component, jax_component, params_jax, config):\n",
    "    tf_component.kernel.assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"{jax_component}/kernel\"].reshape(config.projection_dim, -1)\n",
    "        )\n",
    "    )\n",
    "    tf_component.bias.assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"{jax_component}/bias\"].reshape(-1)\n",
    "        )\n",
    "    )\n",
    "    return tf_component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce251063",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 0\n",
    "for outer_layer in vit_b16_model.layers:\n",
    "    if isinstance(outer_layer, tf.keras.Model) and outer_layer.name != \"projection\":\n",
    "        tf_block = vit_b16_model.get_layer(outer_layer.name)\n",
    "        jax_block_name = f\"encoderblock_{idx}\"\n",
    "\n",
    "        # LayerNorm layers.\n",
    "        layer_norm_idx = 0\n",
    "        for layer in tf_block.layers:\n",
    "            if isinstance(layer, tf.keras.layers.LayerNormalization):\n",
    "                layer_norm_jax_prefix = (\n",
    "                    f\"Transformer/{jax_block_name}/LayerNorm_{layer_norm_idx}\"\n",
    "                )\n",
    "                layer.gamma.assign(\n",
    "                    tf.Variable(params_jax[f\"{layer_norm_jax_prefix}/scale\"])\n",
    "                )\n",
    "                layer.beta.assign(\n",
    "                    tf.Variable(params_jax[f\"{layer_norm_jax_prefix}/bias\"])\n",
    "                )\n",
    "                layer_norm_idx += 2\n",
    "\n",
    "        # FFN layers.\n",
    "        ffn_layer_idx = 0\n",
    "        for layer in tf_block.layers:\n",
    "            if isinstance(layer, tf.keras.layers.Dense):\n",
    "                dense_layer_jax_prefix = (\n",
    "                    f\"Transformer/{jax_block_name}/MlpBlock_3/Dense_{ffn_layer_idx}\"\n",
    "                )\n",
    "                layer.kernel.assign(\n",
    "                    tf.Variable(params_jax[f\"{dense_layer_jax_prefix}/kernel\"])\n",
    "                )\n",
    "                layer.bias.assign(\n",
    "                    tf.Variable(params_jax[f\"{dense_layer_jax_prefix}/bias\"])\n",
    "                )\n",
    "                ffn_layer_idx += 1\n",
    "\n",
    "        # Attention layer.\n",
    "        for layer in tf_block.layers:\n",
    "            attn_layer_jax_prefix = (\n",
    "                f\"Transformer/{jax_block_name}/MultiHeadDotProductAttention_1\"\n",
    "            )\n",
    "            if isinstance(layer, mha.TFViTAttention):\n",
    "                # Key\n",
    "                layer.self_attention.key = modify_attention_block(\n",
    "                    layer.self_attention.key,\n",
    "                    f\"{attn_layer_jax_prefix}/key\",\n",
    "                    params_jax,\n",
    "                    config,\n",
    "                )\n",
    "                # Query\n",
    "                layer.self_attention.query = modify_attention_block(\n",
    "                    layer.self_attention.query,\n",
    "                    f\"{attn_layer_jax_prefix}/query\",\n",
    "                    params_jax,\n",
    "                    config,\n",
    "                )\n",
    "                # Value\n",
    "                layer.self_attention.value = modify_attention_block(\n",
    "                    layer.self_attention.value,\n",
    "                    f\"{attn_layer_jax_prefix}/value\",\n",
    "                    params_jax,\n",
    "                    config,\n",
    "                )\n",
    "                # Final dense projection\n",
    "                layer.dense_output.dense.kernel.assign(\n",
    "                    tf.Variable(\n",
    "                        params_jax[f\"{attn_layer_jax_prefix}/out/kernel\"].reshape(\n",
    "                            -1, config.projection_dim\n",
    "                        )\n",
    "                    )\n",
    "                )\n",
    "                layer.dense_output.dense.bias.assign(\n",
    "                    tf.Variable(params_jax[f\"{attn_layer_jax_prefix}/out/bias\"])\n",
    "                )\n",
    "\n",
    "        idx += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e4354e",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 0\n",
    "for outer_layer in vit_b16_model.layers:\n",
    "    if isinstance(outer_layer, tf.keras.Model) and outer_layer.name != \"projection\":\n",
    "        tf_block = vit_b16_model.get_layer(outer_layer.name)\n",
    "        jax_block_name = f\"encoderblock_{idx}\"\n",
    "\n",
    "        # Layer norm.\n",
    "        layer_norm_idx = 0\n",
    "        for layer in tf_block.layers:\n",
    "            if isinstance(layer, tf.keras.layers.LayerNormalization):\n",
    "\n",
    "                layer_norm_jax_prefix = (\n",
    "                    f\"Transformer/{jax_block_name}/LayerNorm_{layer_norm_idx}\"\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.gamma.numpy(), params_jax[f\"{layer_norm_jax_prefix}/scale\"]\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.beta.numpy(), params_jax[f\"{layer_norm_jax_prefix}/bias\"]\n",
    "                )\n",
    "                layer_norm_idx += 2\n",
    "\n",
    "        # FFN layers.\n",
    "        ffn_layer_idx = 0\n",
    "        for layer in tf_block.layers:\n",
    "            if isinstance(layer, tf.keras.layers.Dense):\n",
    "                dense_layer_jax_prefix = (\n",
    "                    f\"Transformer/{jax_block_name}/MlpBlock_3/Dense_{ffn_layer_idx}\"\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.kernel.numpy(), params_jax[f\"{dense_layer_jax_prefix}/kernel\"]\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.bias.numpy(), params_jax[f\"{dense_layer_jax_prefix}/bias\"]\n",
    "                )\n",
    "                ffn_layer_idx += 1\n",
    "\n",
    "        # Attention layers.\n",
    "        for layer in tf_block.layers:\n",
    "            attn_layer_jax_prefix = (\n",
    "                f\"Transformer/{jax_block_name}/MultiHeadDotProductAttention_1\"\n",
    "            )\n",
    "            if isinstance(layer, mha.TFViTAttention):\n",
    "\n",
    "                # Key\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.key.kernel.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/key/kernel\"].reshape(\n",
    "                        config.projection_dim, -1\n",
    "                    ),\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.key.bias.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/key/bias\"].reshape(-1),\n",
    "                )\n",
    "                # Query\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.query.kernel.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/query/kernel\"].reshape(\n",
    "                        config.projection_dim, -1\n",
    "                    ),\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.query.bias.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/query/bias\"].reshape(-1),\n",
    "                )\n",
    "                # Value\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.value.kernel.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/value/kernel\"].reshape(\n",
    "                        config.projection_dim, -1\n",
    "                    ),\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.self_attention.value.bias.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/value/bias\"].reshape(-1),\n",
    "                )\n",
    "\n",
    "                # Final dense projection\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.dense_output.dense.kernel.numpy(),\n",
    "                    params_jax[f\"{attn_layer_jax_prefix}/out/kernel\"].reshape(\n",
    "                        -1, config.projection_dim\n",
    "                    ),\n",
    "                )\n",
    "                np.testing.assert_allclose(\n",
    "                    layer.dense_output.dense.bias.numpy(),\n",
    "                    tf.Variable(params_jax[f\"{attn_layer_jax_prefix}/out/bias\"]),\n",
    "                )\n",
    "\n",
    "        idx += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed49b697",
   "metadata": {},
   "source": [
    "## Verification\n",
    "\n",
    "Largely taken from here: https://github.com/sayakpaul/BiT-jax2tf/blob/main/convert_jax_weights_tf.ipynb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69977e76",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from PIL import Image\n",
    "from io import BytesIO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62ee1d08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_image(image):\n",
    "    image = np.array(image)\n",
    "    image_resized = tf.image.resize(image, (224, 224))\n",
    "    image_resized = tf.cast(image_resized, tf.float32)\n",
    "    image_resized = (image_resized - 127.5) / 127.5\n",
    "    return tf.expand_dims(image_resized, 0).numpy()\n",
    "\n",
    "def load_image_from_url(url):\n",
    "    response = requests.get(url)\n",
    "    image = Image.open(BytesIO(response.content))\n",
    "    image = preprocess_image(image)\n",
    "    return image\n",
    "\n",
    "!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fa58197",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"ilsvrc2012_wordnet_lemmas.txt\", \"r\") as f:\n",
    "    lines = f.readlines()\n",
    "imagenet_int_to_str = [line.rstrip() for line in lines]\n",
    "\n",
    "img_url = \"https://p0.pikrepo.com/preview/853/907/close-up-photo-of-gray-elephant.jpg\"\n",
    "image = load_image_from_url(img_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba3feda",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions = vit_b16_model.predict(image)\n",
    "logits = predictions[0]\n",
    "predicted_label = imagenet_int_to_str[int(np.argmax(logits))]\n",
    "expected_label = \"Indian_elephant, Elephas_maximus\"\n",
    "assert (\n",
    "    predicted_label == expected_label\n",
    "), f\"Expected {expected_label} but was {predicted_label}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e75cd753",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_b16_model.save(\"vit_b16_patch16_224\")"
   ]
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}