training-deit-for-dfdc.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training sample model with DeiT: Data-efficient Image Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Useful soruces:\n",
    "\n",
    "* [DeiT classifier for Cassava dataset](https://www.kaggle.com/huseynlilkin/cnn-or-transformer-pytorch-xla-tpu-for-cassava/edit)\n",
    "* [Kaggle Utilty Script](https://www.kaggle.com/huseynlilkin/kaggle-pytorch-utility-script/edit?rvi=1)\n",
    "* [Face-net Pytorch for face detection](https://www.kaggle.com/huseynlilkin/guide-to-mtcnn-in-facenet-pytorch/edit)\n",
    "* [DeepFake starter kit](https://www.kaggle.com/huseynlilkin/deepfake-starter-kit/edit)\n",
    "* [Has nice helper functions](https://www.kaggle.com/huseynlilkin/my-deep-fake-solution/edit)\n",
    "* [LR reduce - good webiste](https://www.deeplearningwizard.com/deep_learning/boosting_models_pytorch/lr_scheduling/#reduce-on-loss-plateau-decay-patience0-factor01)\n",
    "* [Selim's code](https://github.com/selimsef/dfdc_deepfake_challenge)\n",
    "* [DeiT github page](https://github.com/facebookresearch/deit)\n",
    "* [AffinityPropagation for clustering](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.AffinityPropagation.html#sklearn.cluster.AffinityPropagation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.770481Z",
     "start_time": "2021-03-26T15:29:25.229482Z"
    },
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_kg_hide-output": true,
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import glob\n",
    "import time\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.models as models\n",
    "from torch.optim import Adam, SGD\n",
    "from torch.nn.parameter import Parameter\n",
    "from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler\n",
    "from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau\n",
    "from torch.cuda.amp import autocast, GradScaler\n",
    "from torch.nn.utils import clip_grad_norm_\n",
    "\n",
    "from albumentations import Compose, RandomBrightnessContrast, \\\n",
    "    HorizontalFlip, HueSaturationValue, OneOf, ToGray, \\\n",
    "    ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur, Resize, Normalize\n",
    "\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "from albumentations import ImageOnlyTransform\n",
    "\n",
    "import config"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Detection with DEiT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.772476Z",
     "start_time": "2021-03-26T15:29:25.229Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "# CONFIG\n",
    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "MODEL_NAME = 'deit_base_patch16_224' # other model names ['deit_base_patch16_224', 'vit_base_patch16_384', 'resnext50_32x4d', 'tf_efficientnet_b3_ns']\n",
    "LOAD_PRETRAINED = True\n",
    "TARGET_SIZE = 2\n",
    "LOAD_CHECKPOINT = False\n",
    "\n",
    "LR = 1e-3\n",
    "# lr scheduler\n",
    "MODE = 'min'\n",
    "FACTOR = 0.1\n",
    "PATIENCE = 1\n",
    "\n",
    "EPOCHS = 20\n",
    "BATCH_SIZE = 32\n",
    "NUM_WORKERS = 0 #os.cpu_count() - 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.773478Z",
     "start_time": "2021-03-26T15:29:25.230Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "class DFDCDataset(Dataset):\n",
    "    def __init__(self, df, transform=None):\n",
    "        self.df = df\n",
    "        self.names = df['name'].values\n",
    "        self.label_names = df['label'].values\n",
    "        self.labels = np.where(self.label_names == 'REAL', 1, 0)\n",
    "        self.transform = transform\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        label = torch.tensor(self.labels[idx]).long()\n",
    "        \n",
    "        file_name = self.names[idx]\n",
    "        file_path = f'{config.root}/train_faces/{file_name}'\n",
    "        image = cv2.imread(file_path)\n",
    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        if self.transform:\n",
    "            augmented = self.transform(image=image)\n",
    "            image = augmented['image']\n",
    "        \n",
    "        return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.774477Z",
     "start_time": "2021-03-26T15:29:25.232Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "class DFDCDatasetNPZ(Dataset):\n",
    "    \"\"\"\n",
    "    Loads data from numpy npz files\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, npz_path, transform=None):\n",
    "        self.npz = np.load(npz_path, allow_pickle=True)\n",
    "        self.data = self.npz['data']\n",
    "        self.label_names = self.npz['labels']\n",
    "        self.labels = np.where(self.label_names == 'REAL', 1, 0)\n",
    "        self.transform = transform\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        label = torch.tensor(self.labels[idx]).long()\n",
    "        image = np.asarray(self.data[idx])\n",
    "        \n",
    "        if self.transform:\n",
    "            augmented = self.transform(image=image)\n",
    "            image = augmented['image']\n",
    "            \n",
    "        return image, label\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-21T15:57:44.588059Z",
     "start_time": "2021-03-21T15:57:44.582060Z"
    }
   },
   "outputs": [],
   "source": [
    "# Original transform functions from Selim's code\n",
    "def create_train_transforms_by_selim(size=300):\n",
    "    return Compose([\n",
    "        ImageCompression(quality_lower=60, quality_upper=100, p=0.5),\n",
    "        GaussNoise(p=0.1),\n",
    "        GaussianBlur(blur_limit=3, p=0.05),\n",
    "        HorizontalFlip(),\n",
    "        OneOf([\n",
    "            IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),\n",
    "            IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_LINEAR),\n",
    "            IsotropicResize(max_side=size, interpolation_down=cv2.INTER_LINEAR, interpolation_up=cv2.INTER_LINEAR),\n",
    "        ], p=1),\n",
    "        PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),\n",
    "        OneOf([RandomBrightnessContrast(), FancyPCA(), HueSaturationValue()], p=0.7),\n",
    "        ToGray(p=0.2),\n",
    "        ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5),\n",
    "    ]\n",
    "    )\n",
    "\n",
    "\n",
    "def create_val_transforms_by_selim(size=300):\n",
    "    return Compose([\n",
    "        IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),\n",
    "        PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.775478Z",
     "start_time": "2021-03-26T15:29:25.234Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "def create_train_transforms(size=224):\n",
    "    return Compose([\n",
    "        Resize(224, 224),\n",
    "        ImageCompression(quality_lower=60, quality_upper=100, p=0.5),\n",
    "        GaussNoise(p=0.1),\n",
    "        GaussianBlur(blur_limit=3, p=0.05),\n",
    "        HorizontalFlip(),\n",
    "        # TODO: IsotropicResize\n",
    "        PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),\n",
    "        OneOf([RandomBrightnessContrast(), HueSaturationValue()], p=0.7), # FancyPCA() is missing\n",
    "        ToGray(p=0.2),\n",
    "        ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5),\n",
    "        Normalize(\n",
    "            mean = [0.485, 0.456, 0.406],\n",
    "            std = [0.229, 0.224, 0.225]\n",
    "        ),\n",
    "        ToTensorV2()\n",
    "    ])\n",
    "\n",
    "def create_val_transforms(size=224):\n",
    "    return Compose([\n",
    "        Resize(224, 224),\n",
    "        PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT),\n",
    "        Normalize(\n",
    "            mean = [0.485, 0.456, 0.406],\n",
    "            std = [0.229, 0.224, 0.225]\n",
    "        ),\n",
    "        ToTensorV2()\n",
    "    ])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.776478Z",
     "start_time": "2021-03-26T15:29:25.235Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "class DeiT(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.model = torch.hub.load('facebookresearch/deit:main', MODEL_NAME, pretrained=LOAD_PRETRAINED)\n",
    "        n_features = self.model.head.in_features\n",
    "        self.model.head = nn.Linear(n_features, TARGET_SIZE)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "    \n",
    "class EfficientNet(nn.Module):\n",
    "    def __init__(self, version='b0'):\n",
    "        super().__init__()\n",
    "        self.model = timm.create_model(f'tf_efficientnet_{version}', pretrained=LOAD_PRETRAINED)\n",
    "        n_features = self.model.classifier.in_features\n",
    "        self.model.classifier = nn.Linear(n_features, TARGET_SIZE)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.model(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- **TODO: CHECK HOW THEY TRAINED DEIT ORIGINALLY**\n",
    "- **TODO: NEXT TIME, DOWNLOAD OTHER NOTEBOOK AND OPEN LOCALLY**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.777478Z",
     "start_time": "2021-03-26T15:29:25.237Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computers and stores the average oand current value\"\"\"\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "        \n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "    \n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.778479Z",
     "start_time": "2021-03-26T15:29:25.238Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "def save_checkpoint(model, optimizer, scheduler, scaler, epoch, loss, name='demo_local_train.tar'):\n",
    "    path = 'D:/DFDC/models/' + name\n",
    "    torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'scheduler_state_dict': scheduler.state_dict(),\n",
    "            'scaler_state_dict': scaler.state_dict(),\n",
    "            'loss': loss,\n",
    "            }, path)\n",
    "    \n",
    "\n",
    "def load_checkpoint(path, model, optimizer, scheduler, scaler):\n",
    "    checkpoint = torch.load(path)\n",
    "    \n",
    "    model.load_state_dict(checkpoint['model_state_dict'])\n",
    "    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
    "    scaler.load_state_dict(checkpoint['scaler_state_dict'])\n",
    "    \n",
    "    epoch = checkpoint['epoch']\n",
    "    loss = checkpoint['loss']\n",
    "    \n",
    "    return model, optimizer, scheduler, scaler, epoch, loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-26T15:29:27.779480Z",
     "start_time": "2021-03-26T15:29:25.239Z"
    },
    "init_cell": true
   },
   "outputs": [],
   "source": [
    "def get_validation_loader(path_to_chunk):\n",
    "    validation_dataset = DFDCDatasetNPZ(path_to_chunk, create_val_transforms())\n",
    "    validation_loader = DataLoader(validation_dataset, \n",
    "                         batch_size = BATCH_SIZE,\n",
    "                         shuffle = False,\n",
    "                         num_workers = NUM_WORKERS,\n",
    "                         pin_memory = True,\n",
    "                         drop_last = True)\n",
    "    \n",
    "    return validation_loader\n",
    "    \n",
    "    \n",
    "def get_train_loader(path_to_chunk):\n",
    "    train_dataset = DFDCDatasetNPZ(path_to_chunk, create_train_transforms())\n",
    "    train_loader = DataLoader(train_dataset, \n",
    "                             batch_size = BATCH_SIZE,\n",
    "                             shuffle = True,\n",
    "                             num_workers = NUM_WORKERS,\n",
    "                             pin_memory = True,\n",
    "                             drop_last = True)\n",
    "\n",
    "    return train_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-21T16:31:38.802721Z",
     "start_time": "2021-03-21T16:30:24.729270Z"
    },
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "train_loss = []\n",
    "validation_loss = []\n",
    "\n",
    "model = EfficientNet('b7')\n",
    "model.to(DEVICE)\n",
    "optimizer = Adam(model.parameters(), lr=LR)\n",
    "scheduler = ReduceLROnPlateau(optimizer,  mode=MODE, factor=FACTOR, patience=PATIENCE, verbose=True)\n",
    "criterion = nn.CrossEntropyLoss(weight=torch.FloatTensor([2/10, 8/10]).cuda())\n",
    "scaler = GradScaler()\n",
    "chunks = glob.glob(os.path.dirname(config.CHUNK_PATH)+'/*')\n",
    "\n",
    "# using last chunk as validation\n",
    "validation_chunk_path = chunks.pop()\n",
    "print(f\"Chunk {os.path.basename(validation_chunk_path)} was used as validation\")\n",
    "validation_loader = get_validation_loader(validation_chunk_path)\n",
    "\n",
    "if LOAD_CHECKPOINT:\n",
    "    model, optimizer, scheduler, scaler, epoch, loss = load_checkpoint('D:/DFDC/models/demo_local_train.tar', model, optimizer, scheduler, scaler)\n",
    "    \n",
    "for epoch in range(EPOCHS):\n",
    "    t = time.time()\n",
    "    model.train()\n",
    "    for chunk in chunks:\n",
    "        print(f\"Train starts for : EPOCH: {epoch} CHUNK: {os.path.basename(chunk)}\")\n",
    "        train_loader = get_train_loader(chunk)\n",
    "        for step, (images, labels) in enumerate(train_loader):\n",
    "            images = images.to(DEVICE)\n",
    "            labels = labels.to(DEVICE)\n",
    "            batch_size = labels.size(0)\n",
    "\n",
    "            with autocast():\n",
    "                y_preds = model(images)\n",
    "                loss = criterion(y_preds, labels)\n",
    "                train_loss.append(loss.item())\n",
    "                scaler.scale(loss).backward()\n",
    "                # TODO: should we do gradient clipping here?\n",
    "                # grad_norm = clip_grad_norm_(model.parameters, 1e-7)\n",
    "                # TODO: what is a gradient accumulation?\n",
    "                scaler.step(optimizer)\n",
    "                scaler.update()\n",
    "                optimizer.zero_grad()\n",
    "    \n",
    "    save_checkpoint(model, optimizer, scheduler, scaler, epoch, loss, name='D:/DFDC/checkpoints/model-checkpoint-cfg1.tar')\n",
    "    print(f'Train : Epoch {epoch} completed in {time.time() - t:.2f} secs')\n",
    "    print(f'Last train loss: {train_loss[-1]} AVG train loss: {np.mean(train_loss)}')\n",
    "    \n",
    "    # validation\n",
    "    model.eval()\n",
    "    tmp_loss = []\n",
    "    for step, (images, labels) in enumerate(validation_loader):\n",
    "        images = images.to(DEVICE)\n",
    "        labels = labels.to(DEVICE)\n",
    "        batch_size = labels.size(0)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            y_preds = model(images)\n",
    "        v_loss = criterion(y_preds, labels)\n",
    "        tmp_loss.append(v_loss.item())\n",
    "    \n",
    "    cur_val_loss = np.mean(tmp_loss)\n",
    "    validation_loss.append(cur_val_loss)\n",
    "    print(f\"Validation loss: {cur_val_loss} average validation loss: {np.mean(validation_loss)}\")\n",
    "    prev_loss = -99 if len(validation_loss) < 2 else validation_loss[-2]\n",
    "    \n",
    "    if cur_val_loss < prev_loss:\n",
    "        print(f\"Average validation loss improved from {prev_loss} to {cur_val_loss}. Saving checkpoint\")\n",
    "        save_checkpoint(model, optimizer, scheduler, scaler, epoch, loss, name='D:/DFDC/checkpoints/model-best-cfg1.tar')\n",
    "    \n",
    "    scheduler.step(np.mean(validation_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "20 epochs 2h 40m "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-03-21T15:57:47.039031Z",
     "start_time": "2021-03-21T15:57:44.568Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(24,12))\n",
    "print(f\"Training completed: average training loss: {np.mean(train_loss)}, average validation loss: {np.mean(validation_loss)}\")\n",
    "plt.plot(train_loss)\n",
    "plt.plot(validation_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Initialization Cell",
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "332px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": "20",
    "lenType": 16,
    "lenVar": "80"
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}