train_vit_h_224.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import glob\n",
    "import albumentations as A\n",
    "from albumentations.pytorch.transforms import ToTensorV2\n",
    "import utilities\n",
    "import cv2\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import timm\n",
    "import math\n",
    "from transformers import (get_linear_schedule_with_warmup, \n",
    "                          get_cosine_schedule_with_warmup, \n",
    "                          get_cosine_with_hard_restarts_schedule_with_warmup,\n",
    "                          get_constant_schedule_with_warmup)\n",
    "from tqdm import tqdm\n",
    "import faiss\n",
    "import random\n",
    "import gc\n",
    "import transformers\n",
    "from transformers import CLIPProcessor, CLIPVisionModel,  CLIPVisionConfig\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from pytorch_metric_learning import losses\n",
    "import open_clip\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CFG:\n",
    "    model_name = 'ViT-H-14' \n",
    "    model_data = 'laion2b_s32b_b79k'\n",
    "    samples_per_class = 50\n",
    "    min_samples = 4\n",
    "    image_size = 224 \n",
    "    seed = 5\n",
    "    workers = 6\n",
    "    train_batch_size = 8\n",
    "    valid_batch_size = 32 \n",
    "    emb_size = 512\n",
    "    vit_bb_lr = {'10': 1.25e-6, '20': 2.5e-6, '26': 5e-6, '32': 10e-6} \n",
    "    vit_bb_wd = 1e-3\n",
    "    hd_lr = 3e-4\n",
    "    hd_wd = 1e-5\n",
    "    autocast = True\n",
    "    n_warmup_steps = 1000\n",
    "    n_epochs = 1\n",
    "    device = torch.device('cuda')\n",
    "    s=30.\n",
    "    m=.45\n",
    "    m_min=.05\n",
    "    acc_steps = 4\n",
    "    global_step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "utilities.set_seed(CFG.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('RN50', 'openai'),\n",
       " ('RN50', 'yfcc15m'),\n",
       " ('RN50', 'cc12m'),\n",
       " ('RN50-quickgelu', 'openai'),\n",
       " ('RN50-quickgelu', 'yfcc15m'),\n",
       " ('RN50-quickgelu', 'cc12m'),\n",
       " ('RN101', 'openai'),\n",
       " ('RN101', 'yfcc15m'),\n",
       " ('RN101-quickgelu', 'openai'),\n",
       " ('RN101-quickgelu', 'yfcc15m'),\n",
       " ('RN50x4', 'openai'),\n",
       " ('RN50x16', 'openai'),\n",
       " ('RN50x64', 'openai'),\n",
       " ('ViT-B-32', 'openai'),\n",
       " ('ViT-B-32', 'laion400m_e31'),\n",
       " ('ViT-B-32', 'laion400m_e32'),\n",
       " ('ViT-B-32', 'laion2b_e16'),\n",
       " ('ViT-B-32', 'laion2b_s34b_b79k'),\n",
       " ('ViT-B-32-quickgelu', 'openai'),\n",
       " ('ViT-B-32-quickgelu', 'laion400m_e31'),\n",
       " ('ViT-B-32-quickgelu', 'laion400m_e32'),\n",
       " ('ViT-B-16', 'openai'),\n",
       " ('ViT-B-16', 'laion400m_e31'),\n",
       " ('ViT-B-16', 'laion400m_e32'),\n",
       " ('ViT-B-16-plus-240', 'laion400m_e31'),\n",
       " ('ViT-B-16-plus-240', 'laion400m_e32'),\n",
       " ('ViT-L-14', 'openai'),\n",
       " ('ViT-L-14', 'laion400m_e31'),\n",
       " ('ViT-L-14', 'laion400m_e32'),\n",
       " ('ViT-L-14', 'laion2b_s32b_b82k'),\n",
       " ('ViT-L-14-336', 'openai'),\n",
       " ('ViT-H-14', 'laion2b_s32b_b79k'),\n",
       " ('ViT-g-14', 'laion2b_s12b_b42k')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "open_clip.list_pretrained()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_backbone, model_transforms, _ = open_clip.create_model_and_transforms(CFG.model_name, pretrained=CFG.model_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "landmarks = random.sample(glob.glob('../data/landmark-2020/*'), 9691)\n",
    "products = random.sample(glob.glob('../data/products/*'), 9691)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_globs = []\n",
    "train_globs.append(landmarks)\n",
    "train_globs.append(products)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_samples(paths):\n",
    "\n",
    "    value_counts = []\n",
    "    folder_count = 0\n",
    "    maps = []\n",
    "    for i, var in enumerate(paths):\n",
    "        dataset_maps = []\n",
    "        for j, folder_path in enumerate(var): \n",
    "            folder_contents = glob.glob(folder_path + '/*')\n",
    "            length = len(folder_contents)\n",
    "            if length >= CFG.min_samples:\n",
    "                folder_size = 0\n",
    "                for file_path in folder_contents[:CFG.samples_per_class]:\n",
    "                        folder_size += 1\n",
    "                        dataset_maps.append((file_path, folder_count))\n",
    "                folder_count += 1\n",
    "                value_counts.append(folder_size)\n",
    "        maps.append(dataset_maps)\n",
    "\n",
    "    total = 0\n",
    "    for i, var in enumerate(maps):\n",
    "        length = len(var)\n",
    "        total += length\n",
    "        print('samples in dataset', length)\n",
    "\n",
    "    for i, var in enumerate(maps):\n",
    "        print('percentage of samples of dataset', len(var)/total)\n",
    "\n",
    "    total_samples = []\n",
    "    for i, dataset_map in enumerate(maps):\n",
    "        for j, map in enumerate(dataset_map):\n",
    "            total_samples.append(map)\n",
    "    \n",
    "    return total_samples, folder_count, np.array(value_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train:\n",
      "samples in dataset 145877\n",
      "samples in dataset 139875\n",
      "percentage of samples of dataset 0.5105021137209889\n",
      "percentage of samples of dataset 0.4894978862790112\n"
     ]
    }
   ],
   "source": [
    "print('train:')\n",
    "train_samples, train_classes, value_counts = get_samples(train_globs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "17642\n"
     ]
    }
   ],
   "source": [
    "CFG.n_classes = train_classes\n",
    "print(CFG.n_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "17641\n",
      "17642\n"
     ]
    }
   ],
   "source": [
    "prev_num = 0\n",
    "switches = 0\n",
    "count = 0\n",
    "for var in train_samples:\n",
    "    current_num = var[-1]\n",
    "\n",
    "    if current_num != prev_num:\n",
    "        assert count >= CFG.min_samples\n",
    "        count = 0\n",
    "        switches+= 1\n",
    "    count+=1\n",
    "    \n",
    "    prev_num = current_num\n",
    "print(switches)\n",
    "print(train_classes)\n",
    "assert switches+1 == train_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20111\n"
     ]
    }
   ],
   "source": [
    "valid_paths = glob.glob('../data/objectNET-4-of-10/*')\n",
    "valid_samples = [(file_path, i) for i, folder_path in enumerate(valid_paths) for file_path in glob.glob(folder_path + '/*')]\n",
    "print(len(valid_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "285752\n"
     ]
    }
   ],
   "source": [
    "print(len(train_samples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImageNet_DS(Dataset):\n",
    "    def __init__(self, map, transforms):\n",
    "        self.map=map\n",
    "        self.transforms=transforms\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        image_path, label = self.map[index] \n",
    "        \n",
    "        label = torch.tensor(label, dtype=torch.long)\n",
    "\n",
    "        images = self.transforms(image_path, CFG.image_size)\n",
    "\n",
    "        images = transforms.functional.resize(images, size=[CFG.image_size, CFG.image_size]) \n",
    "\n",
    "        return {\n",
    "            'images': images,\n",
    "            'labels': label\n",
    "        }\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    def __init__(self, hidden_size):\n",
    "        super(Head, self).__init__()\n",
    "\n",
    "        self.emb = nn.Linear(hidden_size, CFG.emb_size, bias=False)\n",
    "        self.arc = utilities.ArcMarginProduct_subcenter(CFG.emb_size, CFG.n_classes)\n",
    "        self.dropout = utilities.Multisample_Dropout()\n",
    "\n",
    "    def forward(self, x):\n",
    "        embeddings = self.dropout(x, self.emb)\n",
    "        \n",
    "        output = self.arc(embeddings)\n",
    "\n",
    "        return output, F.normalize(embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, vit_backbone):\n",
    "        super(Model, self).__init__()\n",
    "\n",
    "        self.vit_backbone = vit_backbone\n",
    "\n",
    "        self.head = Head(1024)\n",
    "\n",
    "    def forward(self, images):\n",
    "\n",
    "        x = transforms.functional.resize(images, size=[CFG.image_size, CFG.image_size]) \n",
    "        x = x/255\n",
    "        x = transforms.functional.normalize(x,  \n",
    "                                             mean=model_transforms.transforms[-1].mean, \n",
    "                                             std=model_transforms.transforms[-1].std)\n",
    "\n",
    "        x = self.vit_backbone.encode_image(x)\n",
    "        \n",
    "        return self.head(x)\n",
    "\n",
    "    def get_parameters(self):\n",
    "\n",
    "        parameter_settings = [] \n",
    "        parameter_settings.extend(self.get_parameter_section([(n, p) for n, p in self.vit_backbone.named_parameters()], lr=CFG.vit_bb_lr, wd=CFG.vit_bb_wd)) \n",
    "\n",
    "        parameter_settings.extend(self.get_parameter_section([(n, p) for n, p in self.head.named_parameters()], lr=CFG.hd_lr, wd=CFG.hd_wd)) \n",
    "\n",
    "        return parameter_settings\n",
    "\n",
    "    def get_parameter_section(self, parameters, lr=None, wd=None): \n",
    "        parameter_settings = []\n",
    "\n",
    "\n",
    "        lr_is_dict = isinstance(lr, dict)\n",
    "        wd_is_dict = isinstance(wd, dict)\n",
    "\n",
    "        layer_no = None\n",
    "        for no, (n,p) in enumerate(parameters):\n",
    "            \n",
    "            for split in n.split('.'):\n",
    "                if split.isnumeric():\n",
    "                    layer_no = int(split)\n",
    "            \n",
    "            if not layer_no:\n",
    "                layer_no = 0\n",
    "            \n",
    "            if lr_is_dict:\n",
    "                for k,v in lr.items():\n",
    "                    if layer_no < int(k):\n",
    "                        temp_lr = v\n",
    "                        break\n",
    "            else:\n",
    "                temp_lr = lr\n",
    "\n",
    "            if wd_is_dict:\n",
    "                for k,v in wd.items():\n",
    "                    if layer_no < int(k):\n",
    "                        temp_wd = v\n",
    "                        break\n",
    "            else:\n",
    "                temp_wd = wd\n",
    "\n",
    "            weight_decay = 0.0 if 'bias' in n else temp_wd\n",
    "\n",
    "            parameter_setting = {\"params\" : p, \"lr\" : temp_lr, \"weight_decay\" : temp_wd}\n",
    "\n",
    "            parameter_settings.append(parameter_setting)\n",
    "\n",
    "            #print(f'no {no} | params {n} | lr {temp_lr} | weight_decay {weight_decay} | requires_grad {p.requires_grad}')\n",
    "\n",
    "        return parameter_settings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = ImageNet_DS(train_samples, utilities.transforms_auto_augment)\n",
    "valid_dataset = ImageNet_DS(valid_samples, utilities.transforms_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size = CFG.train_batch_size, num_workers=CFG.workers, shuffle=True, drop_last=True)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size = CFG.valid_batch_size, num_workers=CFG.workers, shuffle=False, drop_last=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ArcFace_criterion(logits_m, target, margins):\n",
    "    arc = utilities.ArcFaceLossAdaptiveMargin(margins=margins, s=CFG.s)\n",
    "    loss_m = arc(logits_m, target, CFG.n_classes)\n",
    "    return loss_m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, optimizer, scaler, scheduler, epoch):\n",
    "    model.train()\n",
    "    loss_metrics = utilities.AverageMeter()\n",
    "    criterion = ArcFace_criterion\n",
    "\n",
    "    tmp = np.sqrt(1 / np.sqrt(value_counts))\n",
    "    margins = (tmp - tmp.min()) / (tmp.max() - tmp.min()) * CFG.m + CFG.m_min\n",
    "        \n",
    "    bar = tqdm(train_loader)\n",
    "    for step, data in enumerate(bar):\n",
    "        step += 1\n",
    "        images = data['images'].to(CFG.device, dtype=torch.float)\n",
    "        labels = data['labels'].to(CFG.device)\n",
    "        batch_size = labels.size(0)\n",
    "\n",
    "        with torch.cuda.amp.autocast(enabled=CFG.autocast):\n",
    "            outputs, features = model(images)\n",
    "\n",
    "        loss = criterion(outputs, labels, margins)\n",
    "        loss_metrics.update(loss.item(), batch_size)\n",
    "        loss = loss / CFG.acc_steps\n",
    "        scaler.scale(loss).backward()\n",
    "\n",
    "        if step % CFG.acc_steps == 0 or step == len(bar):\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "            optimizer.zero_grad()\n",
    "            scheduler.step()\n",
    "            CFG.global_step += 1\n",
    "            \n",
    "        lrs = utilities.get_lr_groups(optimizer.param_groups)\n",
    "\n",
    "        loss_avg = loss_metrics.avg\n",
    "\n",
    "        bar.set_postfix(loss=loss_avg, epoch=epoch, lrs=lrs, step=CFG.global_step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(model, valid_loader, epoch):\n",
    "    with torch.no_grad():\n",
    "        model.eval() \n",
    "\n",
    "        all_embeddings = []\n",
    "        all_labels = [] \n",
    "\n",
    "        bar = tqdm(valid_loader)\n",
    "        for i, data in enumerate(bar):\n",
    "            images = data['images'].to(CFG.device, dtype=torch.float)\n",
    "            labels = data['labels'].to(CFG.device)\n",
    "\n",
    "            outputs, embeddings = model(images)\n",
    "\n",
    "            all_embeddings.append(embeddings.detach().cpu().numpy())\n",
    "            all_labels.append(labels.detach().cpu().numpy())\n",
    "\n",
    "            bar.set_postfix(epoch=epoch)\n",
    "\n",
    "    all_embeddings = np.concatenate(all_embeddings, axis=0)\n",
    "    all_labels = np.concatenate(all_labels, axis=0)\n",
    "\n",
    "    return all_embeddings, all_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(vit_backbone).to(CFG.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model.get_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 3/35719 [00:02<7:11:16,  1.38it/s, epoch=0, loss=15.3, lrs=['0.000000e+00'], step=0] /home/ivan/miniconda3/lib/python3.7/site-packages/torch/optim/lr_scheduler.py:136: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n",
      "  0%|          | 37/35719 [00:11<2:33:22,  3.88it/s, epoch=0, loss=15.7, lrs=['1.125000e-08', '2.250000e-08', '4.500000e-08', '9.000000e-08', '2.700000e-06'], step=9]"
     ]
    }
   ],
   "source": [
    "optimizer = torch.optim.AdamW(model.get_parameters())\n",
    " \n",
    "scaler = torch.cuda.amp.GradScaler(enabled=CFG.autocast)\n",
    "\n",
    "steps_per_epoch = math.ceil(len(train_loader) / CFG.acc_steps)\n",
    "\n",
    "num_training_steps = math.ceil(CFG.n_epochs * steps_per_epoch)\n",
    "\n",
    "scheduler = get_cosine_schedule_with_warmup(optimizer,\n",
    "                                            num_training_steps=num_training_steps,\n",
    "                                            num_warmup_steps=CFG.n_warmup_steps)   \n",
    "\n",
    "CFG.global_step = 0                   \n",
    "for epoch in range(math.ceil(CFG.n_epochs)):\n",
    "    \n",
    "    train(model, train_loader, optimizer, scaler, scheduler, epoch)\n",
    "    embeddings, labels = eval(model, valid_loader, epoch)\n",
    "\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache() \n",
    "\n",
    "    scores, indices = utilities.get_similiarity(embeddings, 6)\n",
    "    indices = indices[:,1:] \n",
    "    labels, indices = labels.tolist(), indices.tolist() \n",
    "    preds = utilities.convert_indices_to_labels(indices, labels)\n",
    "    score = utilities.map_per_set(labels, preds)\n",
    "    print('score : ', score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = CFG.model_name.replace('/','-')\n",
    "torch.save(model.state_dict(), f'../models/{model_name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('base')",
   "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.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e4ee870ab444af8a8689fba9fdb6a16993f9af4d6f8c51486b98fd7ee4129479"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}