ViT_classification_HF_trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af82b840",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from typing import List, Dict, Union, Optional\n",
    "\n",
    "from datasets import load_metric\n",
    "from transformers import ViTImageProcessor, AutoConfig, ViTModel\n",
    "from transformers import TrainingArguments, Trainer\n",
    "from transformers.modeling_outputs import ImageClassifierOutput\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2\"\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94e1234b",
   "metadata": {},
   "source": [
    "In this notebook, we'll walk through the process of building an image classification model using the Vision Transformer (ViT) architecture. We'll use the Hugging Face Transformers library to create and train the model.\n",
    "\n",
    "### Step 1: Preprocessing with ViTImageProcessor\n",
    "\n",
    "First, we need to set up the ViTImageProcessor to preprocess our image data. The processor loads the pre-trained configuration and tokenizer for the ViT model.\n",
    "\n",
    "### Step 2: Data Collation Function\n",
    "To prepare our data for training, we define a collation function collate_fn that stacks pixel values (images) and their corresponding labels into tensors within a dictionary. This function will be used during data loading.\n",
    "\n",
    "### Step 3: Classification Head\n",
    "We create a classification head for our ViT model. This head consists of linear layers and dropout for mapping the model's output features to class logits.\n",
    "\n",
    "### Step 4: Image Classification Model\n",
    "Now, we define the main image classification model. It consists of three parts: preprocessing, ViT backbone, and the classification head.\n",
    "\n",
    "### Step 5: Model Forward Pass\n",
    "The forward method of our image classification model takes pixel values as input and returns logits. It also computes the loss if labels are provided based on the problem type and number of classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb8e616",
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = ViTImageProcessor.from_pretrained(\"google/vit-base-patch16-224\", cache_dir=\"/archive/turganbay/.huggingface\")\n",
    "\n",
    "def collate_fn(examples: List[Union[torch.Tensor, int]]) -> Dict[str, torch.Tensor]:\n",
    "    \"\"\"\n",
    "    Collates a list of examples into a dictionary with pixel values and labels.\n",
    "\n",
    "    Args:\n",
    "        examples (List[Union[torch.Tensor, int]]): List of examples where each example is a tuple\n",
    "            containing a pixel value tensor (image) and its corresponding label (integer).\n",
    "\n",
    "    Returns:\n",
    "        Dict[str, torch.Tensor]: A dictionary containing \"pixel_values\" (stacked pixel values) and \"labels\" tensors.\n",
    "    \"\"\"\n",
    "    pixel_values = torch.stack([example[0] for example in examples])\n",
    "    labels = torch.tensor([example[1] for example in examples])\n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "\n",
    "class ClassificationHead(torch.nn.Module):\n",
    "    def __init__(self, config: AutoConfig, num_classes: int=100):\n",
    "        super().__init__()\n",
    "\n",
    "        self.config = config\n",
    "\n",
    "        self.dense = torch.nn.Linear(config.hidden_size, config.hidden_size)\n",
    "        self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)\n",
    "        self.out_proj = torch.nn.Linear(config.hidden_size, num_classes)\n",
    "\n",
    "    def init_weights(self):\n",
    "        self.dense.weight.data.normal_(mean=0.0, std=self.config.initializer_range)\n",
    "        if self.dense.bias is not None:\n",
    "            self.dense.bias.data.zero_()\n",
    "\n",
    "    def forward(self, features, **kwargs):\n",
    "        x = features[:, 0, :] # CLS token\n",
    "        x = self.dropout(x)\n",
    "        x = self.dense(x)\n",
    "        x = torch.tanh(x)\n",
    "        x = self.dropout(x)\n",
    "        x = self.out_proj(x)\n",
    "        return x\n",
    "\n",
    "class ImageClassification(nn.Module):\n",
    "    def __init__(self, num_classes: int, backbone: str) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "        self.config = AutoConfig.from_pretrained(backbone, cache_dir=\"/archive/turganbay/.huggingface\")\n",
    "        self.vit = ViTModel.from_pretrained(backbone, self.config, cache_dir=\"/archive/turganbay/.huggingface\")\n",
    "        self.classifier = ClassificationHead(self.config, num_classes)\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        pixel_values,\n",
    "        head_mask=None,\n",
    "        labels=None,\n",
    "        output_attentions=None,\n",
    "        output_hidden_states=None,\n",
    "        interpolate_pos_encoding=None,\n",
    "        return_dict=True,\n",
    "    ):\n",
    "        return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n",
    "\n",
    "        outputs = self.vit(\n",
    "            pixel_values,\n",
    "            head_mask=head_mask,\n",
    "            output_attentions=output_attentions,\n",
    "            output_hidden_states=output_hidden_states,\n",
    "            interpolate_pos_encoding=interpolate_pos_encoding,\n",
    "            return_dict=return_dict,\n",
    "        )\n",
    "\n",
    "        sequence_output = outputs[0]\n",
    "\n",
    "        logits = self.classifier(sequence_output)\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            labels = labels.to(logits.device)\n",
    "            if self.config.problem_type is None:\n",
    "                if self.num_classes == 1:\n",
    "                    loss_fct = nn.MSELoss()\n",
    "                    loss = loss_fct(logits.squeeze(), labels.squeeze())\n",
    "                elif self.num_classes > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):\n",
    "                    loss_fct = nn.CrossEntropyLoss()\n",
    "                    loss = loss_fct(logits.view(-1, self.num_classes), labels.view(-1))\n",
    "                else:\n",
    "                    loss_fct = nn.BCEWithLogitsLoss()\n",
    "                    loss = loss_fct(logits, labels)\n",
    "\n",
    "        if not return_dict:\n",
    "            output = (logits,) + outputs[1:]\n",
    "            return ((loss,) + output) if loss is not None else output\n",
    "\n",
    "        return ImageClassifierOutput(\n",
    "            loss=loss,\n",
    "            logits=logits,\n",
    "            hidden_states=outputs.hidden_states,\n",
    "            attentions=outputs.attentions,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a98e2266",
   "metadata": {},
   "source": [
    "### Step 6: Data Preparation\n",
    "\n",
    "In this step, we'll prepare the dataset for training. We'll use the CIFAR-100 dataset, which contains 100 different classes of images. We define a transformation pipeline for our dataset. This pipeline resizes the images to a specified size and converts them to tensors.\n",
    "\n",
    "#### Loading CIFAR-100 Dataset\n",
    "We load the CIFAR-100 dataset, specifying the data directory, setting train=True for the training split and train=False for the test split. The download=True flag ensures that the dataset is downloaded if it's not already available locally. We apply the previously defined transformation to the dataset. To create a train-validation split, we split the training dataset into two subsets: one for training and one for validation. We use the random_split function from PyTorch to achieve this.\n",
    "\n",
    "With our data prepared and split, we are ready to proceed with training and evaluation in the next steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ad8187",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = processor.size[\"height\"]\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize(size),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "train_dataset = datasets.CIFAR100(\"/archive/turganbay/cifar\", train=True, download=True, transform=transform)\n",
    "test_dataset = datasets.CIFAR100(\"/archive/turganbay/cifar\", train=False, download=True, transform=transform)\n",
    "N = len(train_dataset)\n",
    "T = int(N*0.9)\n",
    "train_dataset, val_dataset = random_split(train_dataset, [T, N-T])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c6708c9",
   "metadata": {},
   "source": [
    "### Step 7: Evaluation Metrics\n",
    "\n",
    "In this step, we'll define and compute evaluation metrics for our image classification model. We'll use commonly used metrics like accuracy, recall, precision, and F1 score. We start by loading metric functions for the evaluation. We initialize metric dictionaries for each metric we want to calculate: \"accuracy,\" \"recall,\" \"precision,\" and \"f1.\"\n",
    "\n",
    "Next, we define a function to compute these metrics given the evaluation predictions. The function takes eval_pred, which is a tuple containing logits and labels. It calculates the predicted class labels, computes the selected metrics, and returns the results as a dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f635ab6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = [\"accuracy\", \"recall\", \"precision\", \"f1\"]\n",
    "metric = {}\n",
    "for met in metrics:\n",
    "    metric[met] = load_metric(met)\n",
    "\n",
    "\n",
    "def compute_metrics(eval_pred: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float]:\n",
    "    \"\"\"\n",
    "    Compute evaluation metrics given prediction logits and true labels.\n",
    "\n",
    "    Args:\n",
    "        eval_pred (Tuple[np.ndarray, np.ndarray]): A tuple containing logits and true labels.\n",
    "\n",
    "    Returns:\n",
    "        Dict[str, float]: A dictionary containing computed evaluation metrics.\n",
    "    \"\"\"\n",
    "    logits, labels = eval_pred\n",
    "    predictions = np.argmax(logits, axis=-1)\n",
    "    metric_res = {}\n",
    "\n",
    "    for met in metrics:\n",
    "        if met != \"accuracy\":\n",
    "            metric_res[met] = metric[met].compute(predictions=predictions, references=labels, average=\"micro\")[met]\n",
    "        else:\n",
    "            metric_res[met] = metric[met].compute(predictions=predictions, references=labels)[met]\n",
    "\n",
    "    return metric_res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71118a7a",
   "metadata": {},
   "source": [
    "### Step 8: Model Training\n",
    "\n",
    "In this step, we'll set up the training process for our image classification model. We'll use the Hugging Face `Trainer` class for this purpose.\n",
    "\n",
    "First, we initialize our image classification model. We specify the number of classes (100 in this case) and the backbone architecture (Google Vision Transformer - `google/vit-base-patch16-224`).\n",
    "\n",
    "We configure the Hugging Face Trainer class to handle our training process. We specify various training arguments, including batch sizes, logging intervals, learning rate, and more.\n",
    "\n",
    "Finally, we start the training process using the configured Trainer. The trainer.train() method performs training for the specified number of steps (max_steps) and saves the best model checkpoint based on the evaluation metric.\n",
    "\n",
    "Now, your image classification model is ready to be trained!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c9e8eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ImageClassification(100, \"google/vit-base-patch16-224\")\n",
    "model.to(device)\n",
    "\n",
    "trainer = Trainer(\n",
    "    model,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=val_dataset,\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=processor,\n",
    "    args=TrainingArguments(\n",
    "        per_device_train_batch_size=64,\n",
    "        per_device_eval_batch_size=64,\n",
    "        logging_steps=200,\n",
    "        save_steps=200,\n",
    "        save_strategy=\"steps\",\n",
    "        evaluation_strategy=\"steps\",\n",
    "        warmup_steps=200,\n",
    "        max_steps=3000,\n",
    "        learning_rate=1e-5,\n",
    "        weight_decay=0.01,\n",
    "        optim=\"adamw_torch\",\n",
    "        load_best_model_at_end=True,\n",
    "        greater_is_better=False,\n",
    "        output_dir=f\"/archive/turganbay/vit_clf_models/model\",\n",
    "        report_to=\"wandb\",\n",
    "        run_name=\"vit_clf\",\n",
    "    ),\n",
    "    data_collator=collate_fn\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aa7a19c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f3ac08",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env_llm",
   "language": "python",
   "name": "env_llm"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}