multilabel_classification/pytorch_network_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train a neural network for multi-label classification on the CelebA dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cleanlab/examples/blob/master/multilabel_classification/pytorch_network_training.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to train a Pytorch neural network for image tagging and use the model to produce out-of-sample predicted class probabilities for each image in the dataset. These are required inputs to find label errors in multi-label classification datasets with cleanlab. Here we consider a subset of the [CelebA](https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) dataset, where each image may be tagged with one or more of the following tags: `['Wearing_Hat', 'Wearing_Necklace', 'Wearing_Necktie', 'Eyeglasses', 'No_Beard', 'Smiling']`, depending on which ones apply to the person depicted.\n",
    "\n",
    "This notebook only shows how to train the network and use it to produce `pred_probs`, using them to find label issues is demonstrated in our other [example](https://github.com/cleanlab/examples/) notebook on [Find Label Errors in Multi-Label Classification Data (CelebA Image Tagging)](https://github.com/cleanlab/examples/blob/multilabel_tutorial/multilabel_classification/image_tagging.ipynb). Here we fit a state-of-the-art neural network initialized from a pretrained [TIMM](https://timm.fast.ai/) network backbone. You can use this same code to obtain a multi-label classifier (i.e. image tagging model) for *any* image dataset.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please install the dependencies specified in this [requirements.txt](https://github.com/cleanlab/examples/blob/master/multilabel_classification/requirements.txt) file before running the notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next you need to download the dataset. The below code attempts to download the dataset from the Google Drive where its authors have stored it, but this Drive only allows limited programmatic downloads per day. If the Google Drive download script fails, please just manually download the following links from the [CelebA dataset webpage](https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html) and save them in the current working directory:\n",
    " * [Images](https://drive.google.com/uc?id=0B7EVK8r0v71pZjFTYXZWM3FlRnM)\n",
    " * [Labels](https://drive.google.com/uc?id=0B7EVK8r0v71pblRyaVFSWGxPY0U)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gdown\n",
    "import os\n",
    "def download_drive(url,output):\n",
    "    filename = gdown.download(url, output, quiet=False)\n",
    "    if filename is None:\n",
    "        print(f\"Downloading {url} failed, please download it from browser and paste it in {os.getcwd()}\")\n",
    "url = 'https://drive.google.com/uc?id=0B7EVK8r0v71pblRyaVFSWGxPY0U'\n",
    "output = 'list_attr_celeba.txt'\n",
    "if not os.path.exists(output):\n",
    "    download_drive(url,output)\n",
    "url = 'https://drive.google.com/uc?id=0B7EVK8r0v71pZjFTYXZWM3FlRnM' # Usually errors out, Download them manually from the link below\n",
    "output = 'img_align_celeba.zip'\n",
    "if not os.path.exists(output):\n",
    "    download_drive(url,output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember you can just manually download the data from the link above if this code failed. Once you have downloaded the zipped data file, unzip the folder which contains a bunch of individual image files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip -qq img_align_celeba.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we import the other required dependencies (make sure you have installed these packages)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from collections import defaultdict\n",
    "from collections import Counter\n",
    "from types import SimpleNamespace\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.model_selection import StratifiedKFold\n",
    "from PIL import Image\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.utils.data as data\n",
    "from timm.optim import create_optimizer\n",
    "from timm.models import create_model\n",
    "from timm.data import resolve_data_config\n",
    "from timm.data.loader import create_loader\n",
    "from timm.utils import CheckpointSaver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we load the dataset, and preprocess it to only keep the classes of interest (i.e. *image tags*) listed above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA =(\"image_id \" + open(\"list_attr_celeba.txt\").read()[7:]).splitlines()\n",
    "\n",
    "from collections import defaultdict\n",
    "q = DATA[0]\n",
    "d_data = defaultdict(list)\n",
    "ls = q.split()\n",
    "for i in q.split():\n",
    "    d_data[i] = []\n",
    "for j in DATA[1:]:\n",
    "    labels = j.split()\n",
    "    for k in range(0,len(labels)):\n",
    "        if k==0:\n",
    "            d_data[ls[k]].append(labels[k])\n",
    "        else:\n",
    "            # map -1 entries -> 0\n",
    "            ps = int((int(labels[k])+1)/2)\n",
    "            d_data[ls[k]].append(ps)\n",
    "            \n",
    "dat = pd.DataFrame.from_dict(d_data)\n",
    "\n",
    "\n",
    "selected = ['image_id',\n",
    "'Eyeglasses',\n",
    " 'Wearing_Earrings',\n",
    " 'Wearing_Hat',\n",
    " 'Wearing_Necklace',\n",
    " 'Wearing_Necktie',\n",
    " 'No_Beard',\n",
    " 'Smiling']\n",
    "\n",
    "def is_label(row):\n",
    "    for s in selected[1:]:\n",
    "        if row[s]!=0:\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "def get_loc(i):\n",
    "    return os.path.join(os.getcwd(),'img_align_celeba/')+i\n",
    "\n",
    "dat_label = dat.apply(is_label,axis=1)\n",
    "dat_selected = dat[dat_label][selected]\n",
    "dat_selected['image_path'] = dat_selected['image_id'].map(lambda x:get_loc(x))\n",
    "selected[0] = 'image_path'\n",
    "\n",
    "df = dat_selected[selected]\n",
    "set_lab = {}\n",
    "for i,row in df.iterrows():\n",
    "    q = str(row.tolist()[1:])\n",
    "    if q not in set_lab:\n",
    "        set_lab[(str(q))]=len(set_lab)\n",
    "\n",
    "# Here we drop a couple rare class-combinations just to simplify stratified data splitting\n",
    "def get_lab(row):\n",
    "    q = str(row.tolist()[1:])\n",
    "    return set_lab[q]\n",
    "\n",
    "df['unique_label'] = df.apply(get_lab,axis=1)\n",
    "cnt = Counter(df['unique_label'])\n",
    "\n",
    "def drop(val):\n",
    "    if cnt[val]>10:\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "is_drop = df['unique_label'].apply(lambda x:drop(x))\n",
    "df = df[is_drop]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resulting DataFrame contains the ID of each image and the classes (i.e. *tags*) that apply to this image.\n",
    "\n",
    "We define a general class of Pytorch neural networks for multi-label image classification. You can use this class for training models on *any* image dataset, and this class can utilize *any* pretrained [TIMM](https://timm.fast.ai/) network backbone. The `MultiLabelModel` below adds an appropriate output layer on top of the network backbone and then fine-tunes the entire network jointly on your multi-label classification dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiLabelModel(nn.Module):\n",
    "    \"\"\" \n",
    "    Pytorch network for multi-label classification that can utilize any TIMM network backbone.\n",
    "    Some of this code is inspired by: https://github.com/yang-ruixin/PyTorch-Image-Models-Multi-Label-Classification\n",
    "    Note this network uses Sigmoid output activations because the predicted probabilities do not need to sum to 1\n",
    "    for multi-label classification, in which each image may belong to multiple classes rather than only one.\n",
    "    \"\"\"\n",
    "    def __init__(self, model, n_classes, class_weights=None, verbose = False):\n",
    "        super().__init__()\n",
    "        self.base_model = model  # network backbone can be any TIMM model\n",
    "        self.num_classes = n_classes\n",
    "        self.verbose = verbose\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.base_model(x)\n",
    "\n",
    "        x = torch.flatten(x, 1)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def get_loss(self, loss_fn, output, target):\n",
    "\n",
    "        return loss_fn(output, target)\n",
    "\n",
    "    def validate(self, loader):\n",
    "        self.eval();\n",
    "        with torch.no_grad():\n",
    "            total_loss = 0\n",
    "            m = nn.Sigmoid()\n",
    "            labels = []\n",
    "            preds = []\n",
    "            for batch_idx, (input, target) in enumerate(loader):\n",
    "                input = input.cuda()\n",
    "                labels.append(target.detach().cpu())\n",
    "                target = target.float().cuda()\n",
    "                output = m(self(input))\n",
    "                loss = self.get_loss(loss_fn, output, target)\n",
    "\n",
    "                total_loss += loss.item()\n",
    "                pred_model = (output > 0.5).detach().cpu()\n",
    "                preds.append(pred_model)\n",
    "\n",
    "            num_of_batches_per_epoch = len(loader)\n",
    "            avg_loss = total_loss / num_of_batches_per_epoch\n",
    "            print(\"VALIDATION DATA STATS\")\n",
    "\n",
    "            print(\"AVERAGE LOSS:\", avg_loss)\n",
    "            preds = torch.cat(preds).int()\n",
    "            labels = torch.cat(labels).int()\n",
    "            acc_score = accuracy_score(labels, preds)\n",
    "            print(\"MULTILABEL accuracy score:\", acc_score)\n",
    "            per_class = []\n",
    "            for i in range(len(preds.T)):\n",
    "                per_class.append(accuracy_score(labels.T[i], preds.T[i]))\n",
    "            print(dataset_train.label_names)\n",
    "            print(per_class)\n",
    "            print('\\n\\n')\n",
    "        return avg_loss\n",
    "\n",
    "    def predict_proba(self, loader):\n",
    "        self.eval();\n",
    "        with torch.no_grad():\n",
    "            m = nn.Sigmoid()\n",
    "            preds = []\n",
    "            for batch_idx, (input, target) in enumerate(loader):\n",
    "                input = input.cuda()\n",
    "                output = m(self(input))\n",
    "                pred_model = output.detach().cpu()\n",
    "                preds.append(pred_model)\n",
    "            preds = torch.cat(preds)\n",
    "        return preds\n",
    "\n",
    "    def train_one_epoch(\n",
    "        self,\n",
    "        loader,\n",
    "        optimizer,\n",
    "        loss_fn,\n",
    "    ):\n",
    "        sta = time.time()\n",
    "        second_order = hasattr(optimizer, \"is_second_order\") and optimizer.is_second_order\n",
    "        self.train()\n",
    "        total_loss = 0\n",
    "        m = nn.Sigmoid()\n",
    "        labels = []\n",
    "        preds = []\n",
    "        ct = 0\n",
    "        for batch_idx, (input, target) in enumerate(loader):\n",
    "            input = input.cuda()\n",
    "            ct += 1\n",
    "            labels.append(target.detach().cpu())\n",
    "            target = target.float().cuda()\n",
    "            output = m(self(input))\n",
    "            loss = self.get_loss(loss_fn, output, target)\n",
    "            total_loss += loss.item()\n",
    "            pred_model = (output > 0.5).detach().cpu()\n",
    "            preds.append(pred_model)\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward(create_graph=second_order)\n",
    "            optimizer.step()\n",
    "            if ct % 80 == 0 and self.verbose:\n",
    "                print(\"LOSS:\", loss.item())\n",
    "        num_of_batches_per_epoch = len(loader)\n",
    "        avg_loss = total_loss / num_of_batches_per_epoch\n",
    "        print(\"TRAINING DATA STATS\")\n",
    "        print(\"AVERAGE LOSS:\", avg_loss)\n",
    "        preds = torch.cat(preds).int()\n",
    "        labels = torch.cat(labels).int()\n",
    "        acc_score = accuracy_score(labels, preds)\n",
    "        print(\"MULTILABEL accuracy score:\", acc_score)\n",
    "        per_class = []\n",
    "        for i in range(len(preds.T)):\n",
    "            per_class.append(accuracy_score(labels.T[i], preds.T[i]))\n",
    "        print(dataset_train.label_names)\n",
    "        print(per_class)\n",
    "        print('\\n\\n')\n",
    "        sto = time.time()\n",
    "        print(\"training time\", sto - sta)\n",
    "        return avg_loss\n",
    "    \n",
    "\n",
    "    def fit(self, loader_train, load_val, num_epochs=10):\n",
    "        if os.path.exists(\"weights_model\"):\n",
    "            print(\"removing weights directory\")\n",
    "            os.system('rm -rf weights_model')\n",
    "        os.mkdir(\"weights_model\")\n",
    "        args = SimpleNamespace()\n",
    "        args.weight_decay = 0\n",
    "        args.lr = 1e-4\n",
    "        args.opt = 'adam'\n",
    "        args.momentum = 0.9\n",
    "        args.sched = \"step\"\n",
    "\n",
    "        optimizer = create_optimizer(args, self)\n",
    "        saver = CheckpointSaver(\n",
    "            model=self,\n",
    "            optimizer=optimizer,\n",
    "            checkpoint_dir=\"weights_model\"\n",
    "        )\n",
    "        errs = []\n",
    "        num_of_data_train = len(loader_train.dataset.data)\n",
    "        for epoch in range(0, num_epochs):\n",
    "            loss_train = self.train_one_epoch(\n",
    "                loader_train,\n",
    "                optimizer,\n",
    "                loss_fn,\n",
    "            )\n",
    "            loss_val = self.validate(loader_val)\n",
    "            errs.append([loss_train, loss_val])\n",
    "            saver.save_checkpoint(epoch, metric=loss_val)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also create a wrapper class for training our `MultiLabel` model on multi-label classification image datasets. You can easily apply this to your own datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetMultiLabel(data.Dataset):\n",
    "    def __init__(\n",
    "            self,\n",
    "            annotation_path=None,\n",
    "            df = None,\n",
    "            transform=None):\n",
    "\n",
    "        super().__init__()\n",
    "        self.transform = transform\n",
    "        self.data = []\n",
    "        self.labels = []\n",
    "        self.label_names = []\n",
    "        if annotation_path is None:\n",
    "            assert df is not None\n",
    "        else:\n",
    "            df = pd.read_csv(annotation_path)\n",
    "        \n",
    "        cols = df.columns\n",
    "        self.label_names = list(cols[1:-1])\n",
    "        for i,row in df.iterrows():\n",
    "            lb = []\n",
    "            for j in cols:\n",
    "                if j=='unique_label':\n",
    "                    continue\n",
    "                if j=='image_path':\n",
    "                    self.data.append(row[j])\n",
    "                else:\n",
    "                    lb.append(float(row[j]))\n",
    "            self.labels.append(lb)\n",
    "                \n",
    "    def __getitem__(self, idx):\n",
    "        img_path = self.data[idx]\n",
    "\n",
    "        img = Image.open(img_path)\n",
    "        if self.transform:\n",
    "            img = self.transform(img)\n",
    "\n",
    "        labels = self.labels[idx]\n",
    "\n",
    "        return img, labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reset_weights(m):\n",
    "  '''\n",
    "  Re-initializes model weights, eg. between cross-validation folds. \n",
    "  '''\n",
    "  for layer in m.children():\n",
    "    if hasattr(layer, 'reset_parameters'):\n",
    "        layer.reset_parameters()\n",
    "\n",
    "def create_df(pred_probs, dataset):\n",
    "    \"\"\"\n",
    "    Returns a dataframe with image_path and predicted probabilities for each image.\n",
    "    \"\"\"\n",
    "    ls = dataset_val.label_names\n",
    "    cl = defaultdict(list)\n",
    "    cl['image_path'] = dataset.data\n",
    "    for i in range(0,len(ls)):\n",
    "        cl[ls[i]] = pred_val.T[i].tolist()\n",
    "    return pd.DataFrame.from_dict(cl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a `DatasetMultiLabel` and `MultiLabelModel` for the Celeb-A dataset. Here we use the [efficientnet_b0](https://rwightman.github.io/pytorch-image-models/models/tf-efficientnet/) backbone for our neural network, but you can easily use any other TIMM backbone. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = DatasetMultiLabel(df = df)\n",
    "\n",
    "loss_fn = nn.BCELoss()\n",
    "\n",
    "model = create_model(\n",
    "    'efficientnet_b0',  # you can replace this with any TIMM backbone\n",
    "    num_classes=len(dataset.labels[0]),\n",
    ")\n",
    "data_config = resolve_data_config(\n",
    "       args = {}, model=model\n",
    "    )\n",
    "\n",
    "model = MultiLabelModel(\n",
    "        model,\n",
    "        n_classes=len(dataset.labels[0]),\n",
    "    ).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We train this network using K-fold cross validation. This allows us to obtain **out-of-sample** predictions for each image in the dataset (i.e. predictions from a copy of the model which never saw this image during training). Out-of-sample predictions are less prone to overfitting, and thus better suited for finding label issues. From each fold of cross-validation, we store the predicted class probabilities for the images that were out-of-sample in a DataFrame `df_pred`. These predictions are subsequently used for finding label issues in cleanlab's Tutorial on [Multi-Label Classification](https://docs.cleanlab.ai/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_splits = 4  # number of cross-validation splits (higher values will take longer but give better results)\n",
    "skf = StratifiedKFold(n_splits=num_splits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ct = 1\n",
    "for train_index, test_index in skf.split(df,df['unique_label']):\n",
    "    if ct!=1:\n",
    "        model.apply(reset_weights);\n",
    "    dataset_train = DatasetMultiLabel(df = df.iloc[train_index])\n",
    "    dataset_val = DatasetMultiLabel(df = df.iloc[test_index])\n",
    "    loader_train = create_loader(\n",
    "        dataset_train,\n",
    "        input_size=data_config[\"input_size\"],\n",
    "        batch_size=64,\n",
    "        is_training=True,\n",
    "        mean=data_config[\"mean\"],\n",
    "        std=data_config[\"std\"],\n",
    "       interpolation=data_config[\"interpolation\"],\n",
    "    )\n",
    "    loader_val = create_loader(\n",
    "        dataset_val,\n",
    "        input_size=data_config[\"input_size\"],\n",
    "        batch_size=64,\n",
    "        is_training=False,\n",
    "        mean=data_config[\"mean\"],\n",
    "        std=data_config[\"std\"],\n",
    "        interpolation=data_config[\"interpolation\"],\n",
    "\n",
    "    )\n",
    "    model.fit(loader_train,loader_val,num_epochs=40)\n",
    "    checkpoint = torch.load(\"weights_model/model_best.pth.tar\")\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    pred_val = model.predict_proba(loader_val)\n",
    "    df_pred = create_df(pred_val,dataset_val)\n",
    "    df_pred.to_csv(str(ct)+\"_fold.csv\",index=False)\n",
    "    ct+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfl = []\n",
    "for i in range(1,num_splits+1):\n",
    "    dfl.append(pd.read_csv(str(i)+\"_fold.csv\"))\n",
    "\n",
    "cols = dfl[0].columns[1:]\n",
    "df_pred = pd.concat(dfl,axis=0)\n",
    "df_pred['image_path'] = (df_pred['image_path'].map(lambda x:x.split('/')[-1]))\n",
    "df_pred.set_index('image_path').to_csv(\"pred_probs.csv\")"
   ]
  }
 ],
 "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "00885e89789f58e60dbba52a405dc834aaf92411914fde0d391f9b48289a0610"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}