finetuning_img_classifier_visual_transformer_lora.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LoRA Fine-tuning\n",
    "\n",
    "Fine-tune a base ViT model with two different datasets using LoRA. Load the base model and dynamically swap both LoRA adapters depending on the task to complete.\n",
    "\n",
    "## Why does this matter?\n",
    "\n",
    "A foundation model knows how to do many things, but it's not great at many tasks. We can fine-tune the model to produce specialized models that are very good at solving specific tasks.\n",
    "\n",
    "We'll use LoRA to fine-tune the foundation model and generate many, specialized adapters. We can load these adapters together with a model to dynamically transform its capabilities.\n",
    "\n",
    "When loading the model, we'll take the foundation model's original weights and apply the LoRA weight changes to it to get the fine-tuned model weights.\n",
    "\n",
    "The beauty of LoRA is that we don't need to fine-tune the entire matrix of weights. Instead, we can get away by fine-tuning two matrices of lower rank. These matrices, when multiplied together, will get us the weight updates we'll need to apply the foundation model to modify its capabilities.\n",
    "\n",
    "We can save a lot of memory when using LoRA to fine-tune models of different sizes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --quiet transformers accelerate evaluate datasets peft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use a Vision Transformer (ViT) model pre-trained on ImageNet-21k (14 million images, 21,843 classes) at resolution 224x224, and fine-tuned on ImageNet 2012 (1 million images, 1,000 classes) at resolution 224x224. [Here is the model card](https://huggingface.co/google/vit-base-patch16-224).\n",
    "\n",
    "This model has a size of 346 MB on disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint = \"google/vit-base-patch16-224-in21k\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating A Couple Of Helpful Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import torch\n",
    "from peft import PeftModel, LoraConfig, get_peft_model\n",
    "from transformers import AutoModelForImageClassification\n",
    "\n",
    "\n",
    "def print_model_size(path):\n",
    "    size = 0\n",
    "    for f in os.scandir(path):\n",
    "        size += os.path.getsize(f)\n",
    "\n",
    "    print(f\"Model size: {(size / 1e6):.2} MB\")\n",
    "\n",
    "\n",
    "def print_trainable_parameters(model, label):\n",
    "    parameters, trainable = 0, 0\n",
    "    \n",
    "    for _, p in model.named_parameters():\n",
    "        parameters += p.numel()\n",
    "        trainable += p.numel() if p.requires_grad else 0\n",
    "\n",
    "    print(f\"{label} trainable parameters: {trainable:,}/{parameters:,} ({100 * trainable / parameters:.2f}%)\")\n",
    "\n",
    "\n",
    "def split_dataset(dataset):\n",
    "    dataset_splits = dataset.train_test_split(test_size=0.1)\n",
    "    return dataset_splits.values()\n",
    "    \n",
    "\n",
    "def create_label_mappings(dataset):\n",
    "    label2id, id2label = dict(), dict()\n",
    "    for i, label in enumerate(dataset.features[\"label\"].names):\n",
    "        label2id[label] = i\n",
    "        id2label[i] = label \n",
    "\n",
    "    return label2id, id2label"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading and Preparing the Datasets\n",
    "\n",
    "We'll be loading two different datasets to fine-tune the base model:\n",
    "\n",
    "1. A dataset of pictures of food.\n",
    "2. A dataset of pictures of cats and dogs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# This is the food dataset\n",
    "dataset1 = load_dataset(\"food101\", split=\"train[:10000]\")\n",
    "\n",
    "# This is the datasets of pictures of cats and dogs.\n",
    "# Notice we need to rename the label column so we can\n",
    "# reuse the same code for both datasets.\n",
    "dataset2 = load_dataset(\"microsoft/cats_vs_dogs\", split=\"train\", trust_remote_code=True)\n",
    "dataset2 = dataset2.rename_column(\"labels\", \"label\")\n",
    "\n",
    "dataset1_train, dataset1_test = split_dataset(dataset1)\n",
    "dataset2_train, dataset2_test = split_dataset(dataset2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need these mappings to properly fine-tune the Vision Transformer model. You can find more information in the [`PretrainedConfig`](https://huggingface.co/docs/transformers/en/main_classes/configuration#transformers.PretrainedConfig) documentation, under the \"Parameters for fine-tuning tasks\" section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset1_label2id, dataset1_id2label = create_label_mappings(dataset1)\n",
    "dataset2_label2id, dataset2_id2label = create_label_mappings(dataset2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"model1\": {\n",
    "        \"train_data\": dataset1_train,\n",
    "        \"test_data\": dataset1_test,\n",
    "        \"label2id\": dataset1_label2id,\n",
    "        \"id2label\": dataset1_id2label,\n",
    "        \"epochs\": 5,\n",
    "        \"path\": \"./lora-model1\"\n",
    "    },\n",
    "    \"model2\": {\n",
    "        \"train_data\": dataset2_train,\n",
    "        \"test_data\": dataset2_test,\n",
    "        \"label2id\": dataset2_label2id,\n",
    "        \"id2label\": dataset2_id2label,\n",
    "        \"epochs\": 1,\n",
    "        \"path\": \"./lora-model2\"\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create an image processor automatically from the [preprocessor configuration](https://huggingface.co/google/vit-base-patch16-224/blob/main/preprocessor_config.json) specified by the base model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoImageProcessor\n",
    "\n",
    "image_processor = AutoImageProcessor.from_pretrained(model_checkpoint, use_fast=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now prepare the preprocessing pipeline to transform the images in our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import (\n",
    "    CenterCrop,\n",
    "    Compose,\n",
    "    Normalize,\n",
    "    Resize,\n",
    "    ToTensor,\n",
    ")\n",
    "\n",
    "preprocess_pipeline = Compose([\n",
    "    Resize(image_processor.size[\"height\"]),\n",
    "    CenterCrop(image_processor.size[\"height\"]),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=image_processor.image_mean, std=image_processor.image_std),\n",
    "])\n",
    "\n",
    "def preprocess(batch):\n",
    "    batch[\"pixel_values\"] = [\n",
    "        preprocess_pipeline(image.convert(\"RGB\")) for image in batch[\"image\"]\n",
    "    ]\n",
    "    return batch\n",
    "\n",
    "\n",
    "# Let's set the transform function to every train and test sets\n",
    "for cfg in config.values():\n",
    "    cfg[\"train_data\"].set_transform(preprocess)\n",
    "    cfg[\"test_data\"].set_transform(preprocess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-Tuning the Model\n",
    "\n",
    "These are functions that we'll need to fine-tune the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import evaluate\n",
    "import torch\n",
    "from peft import PeftModel, LoraConfig, get_peft_model\n",
    "from transformers import AutoModelForImageClassification\n",
    "\n",
    "\n",
    "metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "\n",
    "def data_collate(examples):\n",
    "    \"\"\"\n",
    "    Prepare a batch of examples from a list of elements of the\n",
    "    train or test datasets.\n",
    "    \"\"\"\n",
    "    pixel_values = torch.stack([example[\"pixel_values\"] for example in examples])\n",
    "    labels = torch.tensor([example[\"label\"] for example in examples])\n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    \"\"\"\n",
    "    Compute the model's accuracy on a batch of predictions.\n",
    "    \"\"\"\n",
    "    predictions = np.argmax(eval_pred.predictions, axis=1)\n",
    "    return metric.compute(predictions=predictions, references=eval_pred.label_ids)\n",
    "\n",
    "\n",
    "def get_base_model(label2id, id2label):\n",
    "    \"\"\"\n",
    "    Create an image classification base model from\n",
    "    the model checkpoint.\n",
    "    \"\"\"\n",
    "    return AutoModelForImageClassification.from_pretrained(\n",
    "        model_checkpoint,\n",
    "        label2id=label2id,\n",
    "        id2label=id2label,\n",
    "        ignore_mismatched_sizes=True, \n",
    "    )\n",
    "\n",
    "\n",
    "def build_lora_model(label2id, id2label):\n",
    "    \"\"\"Build the LoRA model to fine-tune the base model.\"\"\"\n",
    "    model = get_base_model(label2id, id2label)\n",
    "    print_trainable_parameters(model, label=\"Base model\")\n",
    "\n",
    "    config = LoraConfig(\n",
    "        r=16,\n",
    "        lora_alpha=16,\n",
    "        target_modules=[\"query\", \"value\"],\n",
    "        lora_dropout=0.1,\n",
    "        bias=\"none\",\n",
    "        modules_to_save=[\"classifier\"],\n",
    "    )\n",
    "\n",
    "    lora_model = get_peft_model(model, config)\n",
    "    print_trainable_parameters(lora_model, label=\"LoRA\")\n",
    "\n",
    "    return lora_model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now configure the fine-tuning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "batch_size = 128\n",
    "training_arguments = TrainingArguments(\n",
    "    output_dir=\"./model-checkpoints\",\n",
    "    remove_unused_columns=False,\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=5e-3,\n",
    "    per_device_train_batch_size=batch_size,\n",
    "    per_device_eval_batch_size=batch_size,\n",
    "    gradient_accumulation_steps=4,\n",
    "    fp16=True,\n",
    "    logging_steps=10,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\",\n",
    "    label_names=[\"labels\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now fine-tune both models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base model trainable parameters: 85,876,325/85,876,325 (100.00%)\n",
      "LoRA trainable parameters: 667,493/86,543,818 (0.77%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='85' max='85' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [85/85 05:34, Epoch 4/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>2.403800</td>\n",
       "      <td>0.286665</td>\n",
       "      <td>0.937000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.175100</td>\n",
       "      <td>0.222291</td>\n",
       "      <td>0.936000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.086300</td>\n",
       "      <td>0.191845</td>\n",
       "      <td>0.939000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.036800</td>\n",
       "      <td>0.189346</td>\n",
       "      <td>0.946000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='8' max='8' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [8/8 00:03]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation accuracy: 0.946\n",
      "Model size: 2.7 MB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base model trainable parameters: 85,800,194/85,800,194 (100.00%)\n",
      "LoRA trainable parameters: 591,362/86,391,556 (0.68%)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='41' max='41' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [41/41 02:27, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.008800</td>\n",
       "      <td>0.021229</td>\n",
       "      <td>0.994874</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='19' max='19' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [19/19 00:08]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Evaluation accuracy: 0.9948739854762921\n",
      "Model size: 2.4 MB\n"
     ]
    }
   ],
   "source": [
    "from transformers import Trainer\n",
    "\n",
    "for cfg in config.values():\n",
    "    training_arguments.num_train_epochs = cfg[\"epochs\"]\n",
    "    \n",
    "    trainer = Trainer(\n",
    "        build_lora_model(cfg[\"label2id\"], cfg[\"id2label\"]),\n",
    "        training_arguments,\n",
    "        train_dataset=cfg[\"train_data\"],\n",
    "        eval_dataset=cfg[\"test_data\"],\n",
    "        tokenizer=image_processor,\n",
    "        compute_metrics=compute_metrics,\n",
    "        data_collator=data_collate,\n",
    "    )\n",
    "\n",
    "    results = trainer.train()\n",
    "    evaluation_results = trainer.evaluate(cfg['test_data'])\n",
    "    print(f\"Evaluation accuracy: {evaluation_results['eval_accuracy']}\")\n",
    "\n",
    "    # We can now save the fine-tuned model to disk.\n",
    "    trainer.save_model(cfg[\"path\"])\n",
    "    print_model_size(cfg[\"path\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Inference\n",
    "\n",
    "Let's start by defining a couple of functions that will help us build the inference model and run predictions using it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_inference_model(label2id, id2label, lora_adapter_path):\n",
    "    \"\"\"Build the model that will be use to run inference.\"\"\"\n",
    "\n",
    "    # Let's load the base model\n",
    "    model = get_base_model(label2id, id2label)\n",
    "\n",
    "    # Now, we can create the inference model combining the base model\n",
    "    # with the fine-tuned LoRA adapter.\n",
    "    return PeftModel.from_pretrained(model, lora_adapter_path)\n",
    "\n",
    "\n",
    "def predict(image, model, image_processor):\n",
    "    \"\"\"Predict the class represented by the supplied image.\"\"\"\n",
    "    \n",
    "    encoding = image_processor(image.convert(\"RGB\"), return_tensors=\"pt\")\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**encoding)\n",
    "        logits = outputs.logits\n",
    "\n",
    "    class_index = logits.argmax(-1).item()\n",
    "    return model.config.id2label[class_index]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to create two inference models, one using each of the LoRA adapters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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",
      "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"
     ]
    }
   ],
   "source": [
    "for cfg in config.values():\n",
    "    cfg[\"inference_model\"] = build_inference_model(cfg[\"label2id\"], cfg[\"id2label\"], cfg[\"path\"]) \n",
    "    cfg[\"image_processor\"] = AutoImageProcessor.from_pretrained(cfg[\"path\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a list of sample images and the model that we need to use to "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = [\n",
    "    {\n",
    "        \"image\": \"https://www.allrecipes.com/thmb/AtViolcfVtInHgq_mRtv4tPZASQ=/1500x0/filters:no_upscale():max_bytes(150000):strip_icc()/ALR-187822-baked-chicken-wings-4x3-5c7b4624c8554f3da5aabb7d3a91a209.jpg\",\n",
    "        \"model\": \"model1\",\n",
    "    },\n",
    "    {\n",
    "        \"image\": \"https://wallpapers.com/images/featured/kitty-cat-pictures-nzlg8fu5sqx1m6qj.jpg\",\n",
    "        \"model\": \"model2\",\n",
    "    },\n",
    "    {\n",
    "        \"image\": \"https://i.natgeofe.com/n/5f35194b-af37-4f45-a14d-60925b280986/NationalGeographic_2731043_3x4.jpg\",\n",
    "        \"model\": \"model2\",\n",
    "    },\n",
    "    {\n",
    "        \"image\": \"https://www.simplyrecipes.com/thmb/KE6iMblr3R2Db6oE8HdyVsFSj2A=/1500x0/filters:no_upscale():max_bytes(150000):strip_icc()/__opt__aboutcom__coeus__resources__content_migration__simply_recipes__uploads__2019__09__easy-pepperoni-pizza-lead-3-1024x682-583b275444104ef189d693a64df625da.jpg\",\n",
    "        \"model\": \"model1\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now run predictions on every sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prediction: chicken_wings\n",
      "Prediction: cat\n",
      "Prediction: dog\n",
      "Prediction: pizza\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "for sample in samples:\n",
    "    image = Image.open(requests.get(sample[\"image\"], stream=True).raw)\n",
    "    \n",
    "    inference_model = config[sample[\"model\"]][\"inference_model\"]\n",
    "    image_processor = config[sample[\"model\"]][\"image_processor\"]\n",
    "\n",
    "    prediction = predict(image, inference_model, image_processor)\n",
    "    print(f\"Prediction: {prediction}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}