retrieval_regvlm.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "H1r7bAUV0X-7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/ml/miniconda3/envs/py3.11/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n",
            "[nltk_data] Downloading package stopwords to /home/ml/nltk_data...\n",
            "[nltk_data]   Package stopwords is already up-to-date!\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mmadhava20217\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "'ecaxnutp'"
            ]
          },
          "execution_count": 1,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import transformers\n",
        "import torch\n",
        "import torchvision\n",
        "\n",
        "from tqdm import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "\n",
        "from torchinfo import summary\n",
        "import os\n",
        "import glob\n",
        "\n",
        "import tokenizers\n",
        "import itertools\n",
        "\n",
        "import random\n",
        "import math\n",
        "import copy\n",
        "from timm.scheduler import CosineLRScheduler\n",
        "\n",
        "from utils.RegVLM import RegVLM\n",
        "from utils.dataset import re_train_dataset, re_eval_dataset\n",
        "\n",
        "import wandb\n",
        "import nltk\n",
        "\n",
        "os.environ['TOKENIZERS_PARALLELISM'] = 'true'\n",
        "nltk.download('stopwords')\n",
        "\n",
        "device = 'cuda:0'\n",
        "\n",
        "DEVICE = torch.device(device) if torch.cuda.is_available() else torch.device('cpu')\n",
        "model_name = 'ViT-S BERT-S (fixed everything)'\n",
        "\n",
        "MU = 1.0\n",
        "NU = 1.0\n",
        "\n",
        "algo = f'RegVLM - {MU} - {NU}'\n",
        "# fix the seed for reproducibility\n",
        "seed = 6969\n",
        "torch.manual_seed(seed)\n",
        "np.random.seed(seed)\n",
        "random.seed(seed)\n",
        "torch.backends.cudnn.benchmark = True\n",
        "os.environ['WANDB_NOTEBOOK_NAME']  = 'retrieval_regvlm.ipynb'\n",
        "\n",
        "id = wandb.util.generate_id()\n",
        "wandb.login()\n",
        "\n",
        "NUM_WORKERS = 8\n",
        "torch.set_num_threads(12)\n",
        "\n",
        "# set earlier ID\n",
        "# id = 'kkwc4hhd'\n",
        "id = 'ecaxnutp'         # mu = 1., nu = 1.\n",
        "\n",
        "id"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "MAX_LENGTH = 30\n",
        "BATCH_SIZE = 144\n",
        "DIMENSION = 224\n",
        "EPOCHS = 15\n",
        "warmup_epochs = 3\n",
        "\n",
        "INTERVAL = 5\n",
        "\n",
        "EVAL_AT = 128\n",
        "\n",
        "lr = 1e-4\n",
        "init_lr = 1e-6\n",
        "min_lr = 1e-5\n",
        "decay = 0.01\n",
        "beta1 = 0.9\n",
        "beta2 = 0.999\n",
        "\n",
        "n_layers = 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "3YBytuU6NIii"
      },
      "outputs": [],
      "source": [
        "weights_path = f'Models/{model_name}/{algo}/checkpoint_final.pth'\n",
        "\n",
        "MODEL_SAVE_PATH = MODEL_SAVE_PATH = f'Finetuning/{model_name}/{algo}/checkpoint'\n",
        "if os.path.exists(os.path.dirname(MODEL_SAVE_PATH)) == False:\n",
        "    os.makedirs(os.path.dirname(MODEL_SAVE_PATH))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Wii4D6Dnwmey"
      },
      "outputs": [],
      "source": [
        "@torch.no_grad()\n",
        "def evaluation(model, data_loader, tokenizer, device, k=10, max_len = 30):\n",
        "    # test; k for top-k; tokenizer is model.bert\n",
        "    model.eval()\n",
        "\n",
        "    texts = data_loader.dataset.text\n",
        "    num_text = len(texts)\n",
        "    text_bs = 64\n",
        "    text_feats = []\n",
        "    text_embeds = []\n",
        "    text_atts = []\n",
        "    for i in range(0, num_text, text_bs):\n",
        "        text = texts[i: min(num_text, i+text_bs)]\n",
        "        text_input = tokenizer(text, padding='max_length', truncation=True, max_length=max_len, return_tensors=\"pt\").to(device)\n",
        "        text_output = model.bert(text_input.input_ids, text_input.attention_mask)\n",
        "        text_feat = text_output['last_hidden_state']                        #unnormalized\n",
        "        text_embed = model.txt_proj(model.pooler(text_feat.transpose(1,2)))\n",
        "        text_embeds.append(text_embed)\n",
        "        text_feats.append(text_feat)\n",
        "        text_atts.append(text_input.attention_mask)\n",
        "    text_embeds = torch.cat(text_embeds,dim=0)\n",
        "    text_feats = torch.cat(text_feats,dim=0)\n",
        "    text_atts = torch.cat(text_atts,dim=0)\n",
        "\n",
        "    image_feats = []\n",
        "    image_embeds = []\n",
        "    for image, img_id in data_loader:\n",
        "        image = image.to(device)\n",
        "        image_feat = model.vit(image)['last_hidden_state']                  #unnormalized \n",
        "        image_embed = model.img_proj(model.pooler(image_feat.transpose(1, 2)))\n",
        "\n",
        "        image_feats.append(image_feat)\n",
        "        image_embeds.append(image_embed)\n",
        "\n",
        "    image_feats = torch.cat(image_feats,dim=0)\n",
        "    image_embeds = torch.cat(image_embeds,dim=0)\n",
        "\n",
        "    sims_matrix = image_embeds @ text_embeds.t()\n",
        "    score_matrix_i2t = torch.full((len(data_loader.dataset.image),len(texts)),-100.0).to(device)\n",
        "\n",
        "\n",
        "    for i,sims in enumerate(sims_matrix):\n",
        "        topk_sim, topk_idx = sims.topk(k=k, dim=0)\n",
        "\n",
        "        encoder_output = image_feats[i].repeat(k,1,1)\n",
        "        # encoder_att = torch.ones(encoder_output.size()[:-1],dtype=torch.long).to(device)\n",
        "        output = model.fusion(encoder_output,\n",
        "                            text_feats[topk_idx],\n",
        "                            text_atts[topk_idx])['last_hidden_state']\n",
        "        \n",
        "        score = model.itm__head(output[:,0,:])#[:,1]          # take output for prediction head 1\n",
        "        score_matrix_i2t[i,topk_idx] = score.squeeze(1)\n",
        "\n",
        "    sims_matrix = sims_matrix.t()\n",
        "    score_matrix_t2i = torch.full((len(texts),len(data_loader.dataset.image)),-100.0).to(device)\n",
        "\n",
        "    for i,sims in enumerate(sims_matrix):\n",
        "        topk_sim, topk_idx = sims.topk(k=k, dim=0)\n",
        "        encoder_output = image_feats[topk_idx]\n",
        "        # encoder_att = torch.ones(encoder_output.size()[:-1],dtype=torch.long).to(device)\n",
        "        output = model.fusion(encoder_output,\n",
        "                                    text_feats[i].repeat(k,1,1),\n",
        "                                    text_atts[i].repeat(k,1))['last_hidden_state']\n",
        "        score = model.itm__head(output[:,0,:])#[:,1]\n",
        "        score_matrix_t2i[i,topk_idx] = score.squeeze(1)\n",
        "\n",
        "    return score_matrix_i2t.cpu().numpy(), score_matrix_t2i.cpu().numpy()\n",
        "\n",
        "\n",
        "\n",
        "@torch.no_grad()\n",
        "def itm_eval(scores_i2t, scores_t2i, txt2img, img2txt):\n",
        "\n",
        "    #Images->Text\n",
        "    ranks = np.zeros(scores_i2t.shape[0])\n",
        "    for index,score in enumerate(scores_i2t):\n",
        "        inds = np.argsort(score)[::-1]\n",
        "        # Score\n",
        "        rank = 1e20\n",
        "        for i in img2txt[index]:\n",
        "            tmp = np.where(inds == i)[0][0]\n",
        "            if tmp < rank:\n",
        "                rank = tmp\n",
        "        ranks[index] = rank\n",
        "\n",
        "    # Compute metrics\n",
        "    tr1 = 100.0 * len(np.where(ranks < 1)[0]) / len(ranks)\n",
        "    tr5 = 100.0 * len(np.where(ranks < 5)[0]) / len(ranks)\n",
        "    tr10 = 100.0 * len(np.where(ranks < 10)[0]) / len(ranks)\n",
        "\n",
        "    #Text->Images\n",
        "    ranks = np.zeros(scores_t2i.shape[0])\n",
        "\n",
        "    for index,score in enumerate(scores_t2i):\n",
        "        inds = np.argsort(score)[::-1]\n",
        "        ranks[index] = np.where(inds == txt2img[index])[0][0]\n",
        "\n",
        "    # Compute metrics\n",
        "    ir1 = 100.0 * len(np.where(ranks < 1)[0]) / len(ranks)\n",
        "    ir5 = 100.0 * len(np.where(ranks < 5)[0]) / len(ranks)\n",
        "    ir10 = 100.0 * len(np.where(ranks < 10)[0]) / len(ranks)\n",
        "\n",
        "    tr_mean = (tr1 + tr5 + tr10) / 3\n",
        "    ir_mean = (ir1 + ir5 + ir10) / 3\n",
        "    r_mean = (tr_mean + ir_mean) / 2\n",
        "\n",
        "    eval_result =  {'txt_r1': tr1,\n",
        "                    'txt_r5': tr5,\n",
        "                    'txt_r10': tr10,\n",
        "                    'txt_r_mean': tr_mean,\n",
        "                    'img_r1': ir1,\n",
        "                    'img_r5': ir5,\n",
        "                    'img_r10': ir10,\n",
        "                    'img_r_mean': ir_mean,\n",
        "                    'r_mean': r_mean}\n",
        "    return eval_result\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "aQhQuMVd04K8"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision.transforms import v2\n",
        "from PIL import Image\n",
        "\n",
        "def create_dataset(dataset, config):\n",
        "\n",
        "    ## image transforms\n",
        "    train_transform = v2.Compose([\n",
        "        v2.ToImage(),\n",
        "        v2.ToDtype(torch.int8, scale = True),\n",
        "        v2.Resize(size = (DIMENSION, DIMENSION), antialias = False),\n",
        "        v2.RandAugment(),\n",
        "        v2.ToDtype(torch.float32, scale=True),\n",
        "        v2.Normalize(\n",
        "            mean = [0.5, 0.5, 0.5],\n",
        "            std =  [0.5, 0.5, 0.5]\n",
        "        )\n",
        "    ])\n",
        "    test_transform = v2.Compose([\n",
        "        v2.ToImage(),\n",
        "        v2.ToDtype(torch.int8, scale = True),\n",
        "        v2.Resize(size = (DIMENSION, DIMENSION), antialias = False),\n",
        "        v2.ToDtype(torch.float32, scale=True),\n",
        "        v2.Normalize(\n",
        "            mean = [0.5, 0.5, 0.5],\n",
        "            std =  [0.5, 0.5, 0.5]\n",
        "        )\n",
        "        ])\n",
        "\n",
        "    if dataset=='re':\n",
        "        train_dataset = re_train_dataset(config['train_file'], train_transform, config['image_root'])\n",
        "        val_dataset = re_eval_dataset(config['val_file'], test_transform, config['image_root'])\n",
        "        test_dataset = re_eval_dataset(config['test_file'], test_transform, config['image_root'])\n",
        "        return train_dataset, val_dataset, test_dataset\n",
        "\n",
        "\n",
        "def create_loader(datasets, samplers, batch_size, num_workers, is_trains, collate_fns):\n",
        "    loaders = []\n",
        "    for dataset,sampler,bs,n_worker,is_train,collate_fn in zip(datasets,samplers,batch_size,num_workers,is_trains,collate_fns):\n",
        "        if is_train:\n",
        "            shuffle = (sampler is None)\n",
        "            drop_last = True\n",
        "        else:\n",
        "            shuffle = False\n",
        "            drop_last = False\n",
        "        loader = DataLoader(\n",
        "            dataset,\n",
        "            batch_size=bs,\n",
        "            num_workers=n_worker,\n",
        "            pin_memory=True,\n",
        "            sampler=sampler,\n",
        "            shuffle=shuffle,\n",
        "            collate_fn=collate_fn,\n",
        "            drop_last=drop_last,\n",
        "        )\n",
        "        loaders.append(loader)\n",
        "    return loaders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "config = {'train_file': ['Jsons/flickr30k_train.json'],\n",
        "          'val_file': 'Jsons/flickr30k_val.json',\n",
        "          'test_file': 'Jsons/flickr30k_test.json',\n",
        "          'image_root': './',\n",
        "          'image_res': DIMENSION}\n",
        "\n",
        "\n",
        "train_dataset, val_dataset, test_dataset = create_dataset('re', config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [],
      "source": [
        "train_dataloader = torch.utils.data.DataLoader(train_dataset,\n",
        "                                               batch_size = BATCH_SIZE,\n",
        "                                               num_workers = NUM_WORKERS,\n",
        "                                               shuffle = True,\n",
        "                                               drop_last = True)\n",
        "val_loader = torch.utils.data.DataLoader(val_dataset,\n",
        "                                         batch_size = BATCH_SIZE,\n",
        "                                         num_workers = NUM_WORKERS,\n",
        "                                         shuffle = False,\n",
        "                                         drop_last = False)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset,\n",
        "                                         batch_size = BATCH_SIZE,\n",
        "                                         num_workers = NUM_WORKERS,\n",
        "                                         shuffle = False,\n",
        "                                         drop_last = False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Some weights of ViTForMaskedImageModeling were not initialized from the model checkpoint at WinKawaks/vit-small-patch16-224 and are newly initialized: ['decoder.0.bias', 'decoder.0.weight', 'vit.embeddings.mask_token']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ]
        }
      ],
      "source": [
        "vit_model = transformers.ViTForMaskedImageModeling.from_pretrained('WinKawaks/vit-small-patch16-224').to(DEVICE)\n",
        "bert_model = transformers.BertForMaskedLM.from_pretrained(\"prajjwal1/bert-small\").to(DEVICE)\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(\"prajjwal1/bert-small\")\n",
        "\n",
        "model = RegVLM(\n",
        "            vit = vit_model,\n",
        "            bert = bert_model,\n",
        "            vit_num_patches= 196,\n",
        "            vit_emb_dim=384,\n",
        "            bert_emb_dim=512,\n",
        "            bert_layers=n_layers,\n",
        "            vocab_size=tokenizer.vocab_size,\n",
        "            mask_token_id= tokenizer.mask_token_id,\n",
        "            # cls_token_id=tokenizer.cls_token_id\n",
        "            ).train().to(DEVICE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {},
      "outputs": [],
      "source": [
        "chkpt = torch.load(weights_path, map_location=DEVICE)\n",
        "model.load_state_dict(chkpt)\n",
        "\n",
        "model = model.to(DEVICE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {},
      "outputs": [],
      "source": [
        "#optimiser\n",
        "optim = torch.optim.AdamW(model.parameters(),\n",
        "                          lr = lr,\n",
        "                          weight_decay = decay,\n",
        "                          betas = [beta1, beta2],\n",
        "                          )\n",
        "\n",
        "epoch_steps = math.ceil(len(train_dataset)/BATCH_SIZE)\n",
        "num_steps = int(EPOCHS * epoch_steps)\n",
        "warmup_steps = int(warmup_epochs * epoch_steps)\n",
        "\n",
        "lr_scheduler = CosineLRScheduler(\n",
        "        optim,\n",
        "        t_initial=num_steps,\n",
        "        # t_mul=1.,\n",
        "        lr_min=min_lr,\n",
        "        warmup_lr_init = init_lr,\n",
        "        warmup_t=warmup_steps,\n",
        "        cycle_limit=1,\n",
        "        t_in_epochs=False,\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/html": [
              "wandb version 0.16.6 is available!  To upgrade, please run:\n",
              " $ pip install wandb --upgrade"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "Tracking run with wandb version 0.16.4"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "Run data is saved locally in <code>/home/ml/ML Projects/Denoising MAMO/wandb/run-20240421_142103-ecaxnutp</code>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "Resuming run <strong><a href='https://wandb.ai/madhava20217/MAMO%20-%20Finetuning/runs/ecaxnutp' target=\"_blank\">RegVLM - 1.0 - 1.0 - ViT-S, BERT-S</a></strong> to <a href='https://wandb.ai/madhava20217/MAMO%20-%20Finetuning' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              " View project at <a href='https://wandb.ai/madhava20217/MAMO%20-%20Finetuning' target=\"_blank\">https://wandb.ai/madhava20217/MAMO%20-%20Finetuning</a>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              " View run at <a href='https://wandb.ai/madhava20217/MAMO%20-%20Finetuning/runs/ecaxnutp' target=\"_blank\">https://wandb.ai/madhava20217/MAMO%20-%20Finetuning/runs/ecaxnutp</a>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "text/html": [
              "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src='https://wandb.ai/madhava20217/MAMO%20-%20Finetuning/runs/ecaxnutp?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
            ],
            "text/plain": [
              "<wandb.sdk.wandb_run.Run at 0x7f4bf9520790>"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# wandB init\n",
        "wandb.init(\n",
        "    id = id,# id,\n",
        "    resume =  'allow',\n",
        "    project = 'MAMO - Finetuning',\n",
        "    name = f'{algo} - ViT-S, BERT-S',\n",
        "\n",
        "    config = {\n",
        "        'architecture': model_name,\n",
        "        'dataset':'ImageNet1K',\n",
        "        'warmup_epochs': warmup_epochs,\n",
        "        'epochs' : EPOCHS,\n",
        "        'batch_size': BATCH_SIZE,\n",
        "        'masking_ratio_img' : 0.25,\n",
        "        'masking_ratio_itxt' : 0.75,\n",
        "        'mask_patch_size': 196,\n",
        "        'image_size' : DIMENSION,\n",
        "        'optim_params':{\n",
        "            'optim': 'AdamW',\n",
        "            'beta1': beta1,\n",
        "            'beta2': beta2,\n",
        "            'weight_decay': decay,\n",
        "            'learning_rate': lr,\n",
        "        },\n",
        "        'accumulation_iters': 1,\n",
        "        'patch_size_mask' : 32,\n",
        "        'eval size': EVAL_AT,\n",
        "        'Mu' : MU,\n",
        "        'Nu': NU\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {},
      "outputs": [],
      "source": [
        "import re\n",
        "nums = [re.match(r'.*checkpoint_(.*).pth', x).group(1) for x in glob.glob(MODEL_SAVE_PATH+'*.pth')]\n",
        "# if len(nums) > 0:\n",
        "#     nums.remove(\"final\")\n",
        "nums = [int(x) for x in nums]\n",
        "\n",
        "CHKPT = -1\n",
        "\n",
        "if len(nums) != 0:\n",
        "    CHKPT = max(nums)\n",
        "\n",
        "    load_path = '{}_{}.pth'.format(MODEL_SAVE_PATH, CHKPT)\n",
        "    chkpt = torch.load(load_path, map_location = {'cuda:1': device, \n",
        "                                                  'cuda:0': device})\n",
        "\n",
        "    model.load_state_dict(chkpt['model'])\n",
        "    optim.load_state_dict(chkpt['optimizer'])\n",
        "    # lr_scheduler.load_state_dict(chkpt['scheduler_state_dict'])\n",
        "    \n",
        "    print(load_path)\n",
        "    \n",
        "    print(\"loaded earlier settings\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "gAJCklE-n0jS"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "  0%|          | 0/1006 [00:00<?, ?it/s]"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Train Loss: 0.2526365371483932: 100%|██████████| 1006/1006 [11:45<00:00,  1.43it/s] \n",
            "Train Loss: 0.2956937146382351: 100%|██████████| 1006/1006 [11:32<00:00,  1.45it/s] \n",
            "Train Loss: 0.35547990285260067: 100%|██████████| 1006/1006 [11:37<00:00,  1.44it/s]\n",
            "Train Loss: 0.3120167626893781: 100%|██████████| 1006/1006 [11:23<00:00,  1.47it/s] \n",
            "Train Loss: 0.28317336071556654: 100%|██████████| 1006/1006 [11:22<00:00,  1.48it/s]\n",
            "Train Loss: 0.2558037112235549: 100%|██████████| 1006/1006 [11:22<00:00,  1.47it/s] \n",
            "Train Loss: 0.22945741174472256: 100%|██████████| 1006/1006 [11:18<00:00,  1.48it/s]\n",
            "Train Loss: 0.2027101660699067: 100%|██████████| 1006/1006 [11:23<00:00,  1.47it/s] \n",
            "Train Loss: 0.17804646554216003: 100%|██████████| 1006/1006 [11:24<00:00,  1.47it/s]\n",
            "Train Loss: 0.16237262795284893: 100%|██████████| 1006/1006 [11:22<00:00,  1.47it/s]\n",
            "Train Loss: 0.1455998471003048: 100%|██████████| 1006/1006 [11:15<00:00,  1.49it/s] \n",
            "Train Loss: 0.1305813566627903: 100%|██████████| 1006/1006 [11:15<00:00,  1.49it/s] \n",
            "Train Loss: 0.11861978343338012: 100%|██████████| 1006/1006 [11:22<00:00,  1.47it/s]\n",
            "Train Loss: 0.10876594763434312: 100%|██████████| 1006/1006 [11:36<00:00,  1.44it/s]\n",
            "Train Loss: 0.10316871030526477: 100%|██████████| 1006/1006 [11:39<00:00,  1.44it/s]\n",
            "Train Loss: 0.10111581215868236: 100%|██████████| 1006/1006 [11:12<00:00,  1.50it/s]\n",
            "Train Loss: 0.09871250603747948: 100%|██████████| 1006/1006 [11:19<00:00,  1.48it/s]\n",
            "Train Loss: 0.09514689859684194: 100%|██████████| 1006/1006 [11:11<00:00,  1.50it/s]\n"
          ]
        }
      ],
      "source": [
        "scaler = torch.cuda.amp.grad_scaler.GradScaler()\n",
        "itm_loss_fn = torch.nn.BCEWithLogitsLoss()\n",
        "\n",
        "for epoch in range(CHKPT+1, EPOCHS + warmup_epochs):\n",
        "    num_samples = 0\n",
        "    ft_loss = 0\n",
        "    # net all losses\n",
        "    net_itc_loss = 0\n",
        "    net_itm_loss = 0\n",
        "    for idx, data in (pbar := tqdm(enumerate(train_dataloader), total = len(train_dataloader))):\n",
        "        img, txt, img_idx= data\n",
        "        text_input = tokenizer(txt, padding='max_length', truncation=True, max_length=35, return_tensors=\"pt\").to(device)\n",
        "        txt, attn_mask = text_input.input_ids, text_input.attention_mask\n",
        "        # vision\n",
        "        img = img.to(DEVICE)\n",
        "\n",
        "        # language\n",
        "        txt = txt.to(DEVICE)\n",
        "        attn_mask = attn_mask.to(DEVICE)\n",
        "\n",
        "\n",
        "\n",
        "        # masked modeling real training\n",
        "        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):                           # casting to f16\n",
        "            # forward step for online network\n",
        "            img_rep, txt_rep, joint_rep, img_txt_matching = model(img,\n",
        "                                                                txt,\n",
        "                                                                attn_mask,\n",
        "                                                                retrieval = True)\n",
        "            # ITC loss\n",
        "            sim, itc_loss = model.get_itc_loss(img_rep, txt_rep)\n",
        "            \n",
        "            #itm loss\n",
        "            # sample for each image and each text separately\n",
        "            neg_txt, neg_img = model.get_samples(sim)\n",
        "            \n",
        "            itm_labels = torch.cat([torch.ones(len(img)),torch.zeros(2*len(img))],\n",
        "                               dim=0).float().unsqueeze(1).to(DEVICE)\n",
        "            # stack \n",
        "            itm_img_feats = torch.vstack([img_rep, img_rep[neg_img]])\n",
        "            itm_txt_feats = torch.vstack([txt_rep[neg_txt], txt_rep])\n",
        "            itm_txt_attn = torch.vstack([attn_mask[neg_txt], attn_mask])\n",
        "\n",
        "            neg_mamo_reps = model.fusion(itm_img_feats, itm_txt_feats, itm_txt_attn)['last_hidden_state'][:, 0, :]\n",
        "            neg_itm_outputs = model.itm__head(neg_mamo_reps)\n",
        "            \n",
        "            itm_outputs = torch.vstack([img_txt_matching, neg_itm_outputs])\n",
        "            \n",
        "            \n",
        "            # softmax probabilities\n",
        "            itm_loss = itm_loss_fn(itm_outputs, itm_labels)\n",
        "\n",
        "            # TOTAL LOSS\n",
        "            net_loss = (itc_loss) + (itm_loss)\n",
        "\n",
        "        scaler.scale(net_loss).backward()\n",
        "\n",
        "        # BACKPROP\n",
        "        scaler.step(optim)        # fp16\n",
        "        scaler.update()           # fp16\n",
        "        optim.zero_grad(set_to_none = True)\n",
        "        lr_scheduler.step_update(epoch * epoch_steps + idx)\n",
        "\n",
        "        # update and calc loss\n",
        "        num_samples+=1\n",
        "\n",
        "        net_itc_loss+= itc_loss.item()\n",
        "        net_itm_loss+= itm_loss.item()\n",
        "        ft_loss+= net_loss.item()\n",
        "        pbar.set_description(f\"Train Loss: {ft_loss/num_samples}\")\n",
        "\n",
        "    train_stats = {'train_loss': ft_loss/num_samples,\n",
        "                   'itc_loss': net_itc_loss/num_samples,\n",
        "                   'itm_loss': net_itm_loss/num_samples}    \n",
        "\n",
        "\n",
        "    val_result = {}\n",
        "    if (epoch - warmup_epochs + 1) % INTERVAL == 0:\n",
        "        # VALIDATION\n",
        "        score_val_i2t, score_val_t2i, = evaluation(model, val_loader, tokenizer, DEVICE, k=EVAL_AT, max_len = MAX_LENGTH)\n",
        "        val_result = itm_eval(score_val_i2t, score_val_t2i, val_loader.dataset.txt2img, val_loader.dataset.img2txt)\n",
        "\n",
        "\n",
        "    log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},\n",
        "                    **{f'val_{k}': v for k, v in val_result.items()},\n",
        "                    'epoch': epoch,\n",
        "                }\n",
        "    \n",
        "    save_path = '{}_{}.pth'.format(MODEL_SAVE_PATH, epoch)\n",
        "    save_obj = {\n",
        "        'model': model.state_dict(),\n",
        "        'optimizer': optim.state_dict(),\n",
        "        # 'lr_scheduler': lr_scheduler.state_dict(),\n",
        "        'epoch': epoch,\n",
        "    }\n",
        "    torch.save(save_obj, save_path)\n",
        "    if (epoch-warmup_epochs+1) % 15 == 0:\n",
        "        wandb.save(save_path)\n",
        "    wandb.log(log_stats)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['/home/ml/ML Projects/Denoising MAMO/wandb/run-20240421_142103-ecaxnutp/files/Finetuning/ViT-S BERT-S (fixed everything)/RegVLM - 1.0 - 1.0/checkpoint_final.pth']"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "save_path = '{}_{}.pth'.format(MODEL_SAVE_PATH, 'final')\n",
        "torch.save(model.state_dict(), save_path)\n",
        "wandb.save(save_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [],
      "source": [
        "# testing\n",
        "score_test_i2t, score_test_t2i = evaluation(model, test_loader, tokenizer, DEVICE, k=EVAL_AT, max_len = MAX_LENGTH)\n",
        "test_result = itm_eval(score_test_i2t, score_test_t2i, test_loader.dataset.txt2img, test_loader.dataset.img2txt)\n",
        "\n",
        "log_stats = {**{f'test_{k}': v for k, v in test_result.items()}}\n",
        "wandb.log(log_stats)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "wandb: WARNING Ensure read and write access to run files dir: /home/ml/ML Projects/Denoising MAMO/wandb/run-20240421_142103-ecaxnutp/files, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables\n",
            "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[32m\u001b[41mERROR\u001b[0m Control-C detected -- Run data was not synced\n"
          ]
        }
      ],
      "source": [
        "wandb.finish()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "py3.11",
      "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.11.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}