pretrain_vicvlm_pixel.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "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",
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l4qpp8th\n"
     ]
    }
   ],
   "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",
    "import wandb \n",
    "import nltk\n",
    "nltk.download('stopwords')\n",
    "\n",
    "# from utils.RegVLM import RegVLM\n",
    "from utils.hog import HOGLayer\n",
    "\n",
    "from utils.dataset import pretrain_dataset\n",
    "from utils.mim_utils import create_masked_image\n",
    "\n",
    "device = 'cuda:0'\n",
    "\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",
    "\n",
    "MU = 5.\n",
    "NU = 3.\n",
    "LAMBDA = 2.\n",
    "\n",
    "NUM_WORKERS = 8\n",
    "torch.set_num_threads(12)\n",
    "\n",
    "\n",
    "id = wandb.util.generate_id()\n",
    "wandb.login()\n",
    "\n",
    "# set earlier ID\n",
    "# id = '3nxfr2f6'\n",
    "\n",
    "\n",
    "algo = f'VicVLM (no MLM and MIM Reg) - {MU} - {NU} - {LAMBDA}'\n",
    "print(id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_JSON = 'Jsons/flickr30k_train.json'\n",
    "MODEL_SAVE_PATH = f'Models/{model_name}/{algo}/checkpoint'\n",
    "\n",
    "lr = 2.5e-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",
    "warmup_epochs = 3\n",
    "EPOCHS = 12\n",
    "\n",
    "BATCH_SIZE = 96\n",
    "\n",
    "MASKING_RATIO_IMG = 0.75\n",
    "MASKING_RATIO_TXT = 0.25\n",
    "\n",
    "HOG_BINS = 9\n",
    "\n",
    "\n",
    "ALPHA = 0.995               # EWMA\n",
    "\n",
    "\n",
    "n_layers = 2\n",
    "\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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DIMENSION = 224\n",
    "\n",
    "MAX_LEN = 30\n",
    "\n",
    "# ViT config\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(\"prajjwal1/bert-small\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms.v2 as v2\n",
    "\n",
    "## image transforms\n",
    "img_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.RandomVerticalFlip(),\n",
    "    # v2.RandomHorizontalFlip(),\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",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = pretrain_dataset(\n",
    "               ann_file = [DATASET_JSON],\n",
    "               transform = img_transform,\n",
    "               tokenizer = tokenizer,\n",
    "               max_words = MAX_LEN,\n",
    "               input_size = DIMENSION,\n",
    "               mask_patch_size = 32,\n",
    "               model_patch_size = 16,\n",
    "               masking_ratio = MASKING_RATIO_IMG,\n",
    "               txt_masking_ratio = MASKING_RATIO_TXT,\n",
    "               mask_token = tokenizer.mask_token,\n",
    "               mask_token_id = tokenizer.mask_token_id,\n",
    "               max_length = MAX_LEN + 5,\n",
    "               )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataloader\n",
    "\n",
    "pretrain_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    batch_size = BATCH_SIZE,\n",
    "    pin_memory = True,\n",
    "    num_workers = NUM_WORKERS,\n",
    "    shuffle = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deleteEncodingLayers(model, num_layers_to_keep):  # must pass in the full bert model\n",
    "    oldModuleList = model.encoder.layer\n",
    "    newModuleList = torch.nn.ModuleList()\n",
    "\n",
    "    # Now iterate over all layers, only keepign only the relevant layers.\n",
    "    for i in range(0, num_layers_to_keep):\n",
    "        newModuleList.append(oldModuleList[i])\n",
    "\n",
    "    # create a copy of the model, modify it with the new list, and return\n",
    "    copyOfModel = copy.deepcopy(model)\n",
    "    copyOfModel.encoder.layer = newModuleList\n",
    "\n",
    "    return copyOfModel\n",
    "\n",
    "def deleteLaterEncodingLayers(model, num_layers_to_keep):  # must pass in the full bert model\n",
    "    oldModuleList = model.encoder.layer\n",
    "    newModuleList = torch.nn.ModuleList()\n",
    "\n",
    "    # Now iterate over all layers, only keepign only the relevant layers.\n",
    "    for i in range(num_layers_to_keep, 0, -1):\n",
    "        newModuleList.append(oldModuleList[-i])\n",
    "\n",
    "    # create a copy of the model, modify it with the new list, and return\n",
    "    copyOfModel = copy.deepcopy(model)\n",
    "    copyOfModel.encoder.layer = newModuleList\n",
    "\n",
    "    return copyOfModel\n",
    "\n",
    "class RegVLM_Mixer(torch.nn.Module):\n",
    "    def __init__(self, base_bert,\n",
    "                 n_layers = 2,\n",
    "                 n_visual_tokens = 196,\n",
    "                 vision_embedding_dim = 384,\n",
    "                 emb_dims = 512,\n",
    "                 cls_token_id = 101):\n",
    "        # prepare decoder\n",
    "        super().__init__()\n",
    "        self.cls_token_id = cls_token_id\n",
    "        self.embedding_module = base_bert.embeddings\n",
    "        self.n_visual_tokens = n_visual_tokens\n",
    "        self.vision_emb_dim = vision_embedding_dim\n",
    "        self.base_model = deleteLaterEncodingLayers(base_bert, n_layers).encoder\n",
    "        \n",
    "        self.pooler = torch.nn.AdaptiveAvgPool1d(1)\n",
    "        self.emb_dimension = emb_dims\n",
    "        \n",
    "        if self.vision_emb_dim == self.emb_dimension:\n",
    "            self.dimension_caster = torch.nn.Identity()\n",
    "        else:\n",
    "            self.dimension_caster = torch.nn.Linear(self.vision_emb_dim, self.emb_dimension, bias = False)  # no bias here\n",
    "        \n",
    "        \n",
    "    def forward(self, vision_embedding, text_embedding, text_attn_mask):\n",
    "        # assert len(vision_embedding) == len(text_embedding)\n",
    "        n_batch = len(vision_embedding)\n",
    "        \n",
    "        cls_emb = self.embedding_module(torch.tensor([[self.cls_token_id]]*n_batch, \n",
    "                                                     device = vision_embedding.device),\n",
    "                                        torch.tensor([[1]]*n_batch,\n",
    "                                                     device = vision_embedding.device))\n",
    "        \n",
    "        # normalize dimensions\n",
    "        new_vision_emb = self.dimension_caster(vision_embedding[:, 1:, :])   # remove cls token here\n",
    "        \n",
    "        # concatenate\n",
    "        concatenated_emb = torch.cat([cls_emb, new_vision_emb, text_embedding], dim = 1)\n",
    "        \n",
    "        # create attention mask\n",
    "        vision_attention_mask = torch.ones(n_batch, self.n_visual_tokens + 1).to(text_attn_mask.device) # add a cls token here\n",
    "        attn_mask = torch.cat([vision_attention_mask, text_attn_mask], dim = 1)\n",
    "        \n",
    "        attn_mask = attn_mask[:, None, None, :]\n",
    "        \n",
    "        # forward\n",
    "        return self.base_model(concatenated_emb, attn_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VicVLM(torch.nn.Module):\n",
    "    def __init__(self,\n",
    "                 vit,\n",
    "                 bert,\n",
    "                 vit_num_patches = 196,\n",
    "                 vit_emb_dim = 384,\n",
    "                 bert_emb_dim = 512,\n",
    "                 bert_layers = 2,\n",
    "                 vocab_size = 30522,\n",
    "                 mask_token_id = 103,\n",
    "                 cls_token_id = 101,\n",
    "                 tau = None,\n",
    "                 lamda = 25,\n",
    "                 mu = 25,\n",
    "                 nu = 1.,\n",
    "                 eps = 1e-4):\n",
    "       super().__init__()\n",
    "       self.vit = vit.vit\n",
    "       self.mim_reconstruction_joint = vit.decoder\n",
    "       self.mim_reconstruction = copy.deepcopy(vit.decoder)\n",
    "       self.bert = bert.base_model\n",
    "       self.bert = deleteEncodingLayers(self.bert.base_model, bert_layers)\n",
    "       self.fusion = RegVLM_Mixer(bert.base_model,\n",
    "                              n_layers = bert_layers,\n",
    "                              n_visual_tokens=vit_num_patches,\n",
    "                              vision_embedding_dim=vit_emb_dim,\n",
    "                              emb_dims = bert_emb_dim,\n",
    "                              cls_token_id = cls_token_id)\n",
    "       \n",
    "       self.lamda = lamda\n",
    "       self.mew = mu\n",
    "       self.nu = nu\n",
    "       self.eps = eps\n",
    "              \n",
    "       # vit patches data\n",
    "       self.vit_num_patches = vit_num_patches\n",
    "       \n",
    "       # vocab size\n",
    "       self.vocab_size = vocab_size\n",
    "       # mask token\n",
    "       self.mask_token_id = mask_token_id\n",
    "       \n",
    "       # learnable temperature parameter\n",
    "       self.tau = torch.nn.Parameter(torch.FloatTensor([0.07]))      # uniform in range 1 to 5\n",
    "       if tau is not None:\n",
    "           self.tau = torch.nn.Parameter(torch.FloatTensor([tau]))      # uniform in range 1 to 5\n",
    "       \n",
    "       self.tau.requires_grad = True\n",
    "\n",
    "       # joint representation\n",
    "       self.pooler = torch.nn.Sequential(\n",
    "           torch.nn.AdaptiveAvgPool1d(1),\n",
    "           torch.nn.Flatten()\n",
    "       )\n",
    "       self.img_proj = torch.nn.Linear(vit_emb_dim, min(vit_emb_dim, bert_emb_dim))\n",
    "       self.txt_proj = torch.nn.Linear(bert_emb_dim, min(vit_emb_dim, bert_emb_dim))\n",
    "\n",
    "       \n",
    "       # masked representation modeling\n",
    "       self.mrm_proj = torch.nn.Sequential(\n",
    "            torch.nn.Linear(bert_emb_dim, bert_emb_dim),\n",
    "            torch.nn.Tanh(),\n",
    "       )\n",
    "       \n",
    "       # head for masked image modeling\n",
    "       self.mim_proj = torch.nn.Sequential(\n",
    "           torch.nn.Linear(bert_emb_dim, vit_emb_dim),\n",
    "       )\n",
    "        \n",
    "       # head for masked language modeling\n",
    "       self.mlm_head_joint = copy.deepcopy(bert.cls)\n",
    "       self.mlm_head = bert.cls\n",
    "       \n",
    "       self.itc_head = torch.nn.Linear(bert_emb_dim, bert_emb_dim)\n",
    "       \n",
    "       self.itm__head = torch.nn.Sequential(\n",
    "           torch.nn.Linear(bert_emb_dim, bert_emb_dim),\n",
    "           torch.nn.LeakyReLU(),\n",
    "           torch.nn.Linear(bert_emb_dim, 1)\n",
    "       )\n",
    "    \n",
    "       \n",
    "       \n",
    "    def forward(self, image, text, attn_mask,\n",
    "                masked_pos = None,\n",
    "                masked_text = None,\n",
    "                image_text_matching = False,\n",
    "                retrieval = False,\n",
    "                ):\n",
    "        \n",
    "        if retrieval is True:\n",
    "            img_rep = self.vit(image)['last_hidden_state']\n",
    "            txt_rep = self.bert(text, attn_mask)['last_hidden_state']\n",
    "            joint_rep = self.fusion(img_rep, txt_rep, attn_mask)['last_hidden_state']\n",
    "            \n",
    "            \n",
    "            # img_rep = self.img_proj(self.pooler(img_rep.transpose(1,2)))\n",
    "            # txt_rep = self.txt_proj(self.pooler(txt_rep.transpose(1,2)))\n",
    "            return img_rep, txt_rep, joint_rep, self.itm__head(joint_rep[:, 0, :])\n",
    "        \n",
    "        if image_text_matching == True:\n",
    "            img_rep = self.vit(image)['last_hidden_state']\n",
    "            txt_rep = self.bert(text, attn_mask)['last_hidden_state']\n",
    "            joint_rep = self.fusion(img_rep, txt_rep, attn_mask)['last_hidden_state']\n",
    "            \n",
    "            return img_rep, txt_rep, joint_rep, self.itm__head(joint_rep[:, 0, :])\n",
    "        \n",
    "        else:\n",
    "            # return mask_img-clean_txt, clean_img,-mask_txt, \n",
    "            img_rep = self.vit(image)['last_hidden_state']              # clean image\n",
    "            txt_rep = self.bert(text, attn_mask)['last_hidden_state']   # clean text\n",
    "            \n",
    "            mask_img_rep = self.vit(image, bool_masked_pos = masked_pos)['last_hidden_state']\n",
    "            mask_txt_rep = self.bert(masked_text, attn_mask)['last_hidden_state']\n",
    "            \n",
    "            # multimodal prediction\n",
    "            c_img_m_txt = self.fusion(img_rep, mask_txt_rep, attn_mask)['last_hidden_state']\n",
    "            m_img_c_txt = self.fusion(mask_img_rep, txt_rep, attn_mask)['last_hidden_state']\n",
    "            \n",
    "            # pure txt\n",
    "            txt_prediction = self.mlm_head(mask_txt_rep)\n",
    "            \n",
    "            # pure fusion\n",
    "            img_txt_joint = self.fusion(img_rep, txt_rep, attn_mask)['last_hidden_state']\n",
    "        \n",
    "            return (c_img_m_txt, m_img_c_txt, img_txt_joint, mask_img_rep, txt_prediction, img_rep, txt_rep)\n",
    "\n",
    "\n",
    "    def get_mim_loss(self, img_rep, img, bool_masked_pos):\n",
    "        sequence_output = img_rep\n",
    "\n",
    "        # Reshape to (batch_size, num_channels, height, width)\n",
    "        sequence_output = sequence_output[:, 1:]\n",
    "        batch_size, sequence_length, num_channels = sequence_output.shape\n",
    "        height = width = math.floor(sequence_length**0.5)\n",
    "        sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width)\n",
    "\n",
    "        # Reconstruct pixel values\n",
    "        reconstructed_pixel_values = self.mim_reconstruction(sequence_output)\n",
    "\n",
    "        masked_im_loss = None\n",
    "        if bool_masked_pos is not None:\n",
    "            size = 224 // 16\n",
    "            bool_masked_pos = bool_masked_pos.reshape(-1, size, size)\n",
    "            mask = (\n",
    "                bool_masked_pos.repeat_interleave(16, 1)\n",
    "                .repeat_interleave(16, 2)\n",
    "                .unsqueeze(1)\n",
    "                .contiguous()\n",
    "            )\n",
    "            reconstruction_loss = torch.nn.functional.l1_loss(img, reconstructed_pixel_values, reduction=\"none\")\n",
    "            masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / 3\n",
    "\n",
    "        return masked_im_loss\n",
    "    \n",
    "    def get_joint_mim_loss(self, joint_rep, img, bool_masked_pos):\n",
    "        joint_rep = joint_rep[:, :self.vit_num_patches+1, :]\n",
    "        joint_rep = self.mim_proj(joint_rep)                    # bs x vit_npatch x vit dim\n",
    "        sequence_output = joint_rep\n",
    "\n",
    "        # Reshape to (batch_size, num_channels, height, width)\n",
    "        sequence_output = sequence_output[:, 1:]\n",
    "        batch_size, sequence_length, num_channels = sequence_output.shape\n",
    "        height = width = math.floor(sequence_length**0.5)\n",
    "        sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width)\n",
    "\n",
    "        # Reconstruct pixel values\n",
    "        reconstructed_pixel_values = self.mim_reconstruction_joint(sequence_output)\n",
    "\n",
    "        masked_im_loss_joint = None\n",
    "        if bool_masked_pos is not None:\n",
    "            size = 224 // 16\n",
    "            bool_masked_pos = bool_masked_pos.reshape(-1, size, size)\n",
    "            mask = (\n",
    "                bool_masked_pos.repeat_interleave(16, 1)\n",
    "                .repeat_interleave(16, 2)\n",
    "                .unsqueeze(1)\n",
    "                .contiguous()\n",
    "            )\n",
    "            reconstruction_loss = torch.nn.functional.l1_loss(img, reconstructed_pixel_values, reduction=\"none\")\n",
    "            masked_im_loss_joint = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / 3\n",
    "\n",
    "        return masked_im_loss_joint\n",
    "    \n",
    "    \n",
    "    def get_joint_mlm_loss(self, joint_reps, sen, masked_sen):\n",
    "        embs = joint_reps[:, self.vit_num_patches+1:, :]\n",
    "        scores = self.mlm_head_joint(torch.nn.functional.leaky_relu(embs))\n",
    "        labels = torch.where(masked_sen == self.mask_token_id, sen, -100)\n",
    "        loss = torch.nn.functional.cross_entropy(scores.view(-1, self.vocab_size), labels.view(-1), ignore_index=-100)\n",
    "        return loss\n",
    "    \n",
    "    def get_mlm_loss(self, scores, sen, masked_sen):\n",
    "        labels = torch.where(masked_sen == self.mask_token_id, sen, -100)\n",
    "        loss = torch.nn.functional.cross_entropy(scores.view(-1, self.vocab_size), labels.view(-1), ignore_index=-100)\n",
    "        \n",
    "        return loss\n",
    "    \n",
    "    \n",
    "    def get_itc_loss(self, img_feats, txt_feats):\n",
    "        # Calculate similarity\n",
    "        with torch.no_grad():\n",
    "            self.tau.clamp_(0.001,0.5)\n",
    "\n",
    "        # pool and flatten text and visual features obtained before fusion\n",
    "        img_feats = self.img_proj(self.pooler(img_feats.transpose(1,2)))\n",
    "        txt_feats = self.txt_proj(self.pooler(txt_feats.transpose(1,2)))\n",
    "        \n",
    "        \n",
    "        sim = (img_feats@txt_feats.T)/self.tau\n",
    "        # sim = torch.clip(sim, max = 1e4, min = 1e-4)\n",
    "        self_mask = torch.eye(sim.shape[0], device=sim.device)\n",
    "        \n",
    "        loss_i2t = -torch.sum(torch.nn.functional.log_softmax(sim, dim = 1)*self_mask, dim = 1).mean()\n",
    "        loss_t2i = -torch.sum(torch.nn.functional.log_softmax(sim.T, dim = 1)*self_mask, dim = 1).mean()\n",
    "\n",
    "        return sim, (loss_i2t+loss_t2i)/2.0\n",
    "    \n",
    "    def get_samples(self, similarities):\n",
    "        probs = torch.nn.functional.softmax(similarities, dim = 1) + 1e-8       # term to make nonnegative\n",
    "        probs = probs.fill_diagonal_(0)         # eliminate full samples\n",
    "        \n",
    "        txt_indices = torch.multinomial(probs, num_samples=1, replacement=True).squeeze(1)\n",
    "        img_indices = torch.multinomial(probs.T, num_samples=1, replacement=True).squeeze(1)\n",
    "        \n",
    "        return txt_indices, img_indices\n",
    "    \n",
    "    def get_vicreg_loss(self, features):\n",
    "        #Source: https://arxiv.org/pdf/2105.04906.pdf\n",
    "              \n",
    "        # Calculate invariance -- not required for now as we're only using a part of vicreg\n",
    "        # invariance_loss = torch.nn.functional.mse_loss(img_txt_feats, img_txt_feats_prime)\n",
    "\n",
    "        # Calculate variance of img_txt_feats using torch.var\n",
    "        features = torch.nn.functional.adaptive_avg_pool1d(features.transpose(1,2), 1).flatten(1)\n",
    "        variance = torch.var(features, dim=0)\n",
    "        std = torch.sqrt(variance + self.eps)\n",
    "        std_loss = torch.mean(torch.relu(1 - std))\n",
    "\n",
    "        # Calculate covariance of img_txt_feats using torch after centering the data\n",
    "        centered_feats = features - torch.mean(features, dim=0)\n",
    "        covariance = torch.mm(centered_feats.T, centered_feats) / (features.shape[0] - 1)        # Get off diagonal elements of covariance\n",
    "        cov_loss = torch.sum(torch.pow(covariance - torch.diag(torch.diag(covariance)),2)) / (features.shape[1])\n",
    "\n",
    "        return  + self.mew * std_loss + self.nu * cov_loss\n",
    "    \n",
    "    def get_vicreg_invariance_loss(self, features, features_prime):\n",
    "        # Calculate invariance\n",
    "        invariance_loss = torch.nn.functional.mse_loss(features, features_prime)\n",
    "        return self.lamda * invariance_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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\")\n",
    "\n",
    "online_network = VicVLM(\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",
    "                    mu = MU,\n",
    "                    nu = NU,\n",
    "                    lamda = LAMBDA\n",
    "                    # cls_token_id=tokenizer.cls_token_id\n",
    "                ).train().to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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-20240423_071811-l4qpp8th</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/madhava20217/MAMO%20-%20Pretrain/runs/l4qpp8th' target=\"_blank\">VicVLM (no MLM and MIM Reg) - 5.0 - 3.0 - 2.0 - ViT-S, BERT-S</a></strong> to <a href='https://wandb.ai/madhava20217/MAMO%20-%20Pretrain' 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-%20Pretrain' target=\"_blank\">https://wandb.ai/madhava20217/MAMO%20-%20Pretrain</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-%20Pretrain/runs/l4qpp8th' target=\"_blank\">https://wandb.ai/madhava20217/MAMO%20-%20Pretrain/runs/l4qpp8th</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-%20Pretrain/runs/l4qpp8th?jupyter=true' style='border:none;width:100%;height:420px;display:none;'></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x7f302a4f92d0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#optimiser\n",
    "optim = torch.optim.AdamW(online_network.parameters(),\n",
    "                          lr = lr,\n",
    "                          weight_decay = decay,\n",
    "                          betas = [0.9, 0.999],\n",
    "                          )\n",
    "\n",
    "epoch_steps = math.ceil(len(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",
    "    )\n",
    "\n",
    "\n",
    "# wandB init\n",
    "wandb.init(\n",
    "    id = id,# id,\n",
    "    resume =  'allow',\n",
    "    project = 'MAMO - Pretrain',\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' : MASKING_RATIO_IMG,\n",
    "        'masking_ratio_itxt' : MASKING_RATIO_TXT,\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",
    "        'mu': MU,\n",
    "        'nu': NU,\n",
    "        'lambda': LAMBDA,         \n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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",
    "    online_network.load_state_dict(chkpt['online_model_state_dict'])\n",
    "    optim.load_state_dict(chkpt['optim_state_dict'])\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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1511 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Train Loss: 19.467766275159732: 100%|██████████| 1511/1511 [12:54<00:00,  1.95it/s]\n",
      "Train Loss: 9.19007332888287: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 8.66620550275559: 100%|██████████| 1511/1511 [12:48<00:00,  1.97it/s] \n",
      "Train Loss: 8.071235903835234: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s]\n",
      "Train Loss: 7.581368010537976: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 7.276085405457028: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 7.013295777028637: 100%|██████████| 1511/1511 [12:48<00:00,  1.97it/s] \n",
      "Train Loss: 6.829367954473792: 100%|██████████| 1511/1511 [12:48<00:00,  1.97it/s] \n",
      "Train Loss: 6.677253091990119: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 6.566205223848455: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 6.490323475699484: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 6.447409676842781: 100%|██████████| 1511/1511 [12:48<00:00,  1.97it/s] \n",
      "Train Loss: 6.428000217553876: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 6.414773126850847: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s] \n",
      "Train Loss: 6.40138408318803: 100%|██████████| 1511/1511 [12:49<00:00,  1.96it/s]  \n"
     ]
    }
   ],
   "source": [
    "scaler = torch.cuda.amp.grad_scaler.GradScaler()\n",
    "itm_loss_fn = torch.nn.BCEWithLogitsLoss()\n",
    "online_network = online_network.to(DEVICE).train()\n",
    "\n",
    "for epoch in range(CHKPT+1, EPOCHS + warmup_epochs):\n",
    "    num_samples = 0\n",
    "    pretrain_loss = 0\n",
    "    # net all losses\n",
    "    net_mim_unimodal_loss = 0\n",
    "    net_mim_joint_loss = 0\n",
    "    net_mlm_loss = 0\n",
    "    # net_mlm_joint_loss = 0\n",
    "    net_vicreg_var_loss = 0\n",
    "    net_vicreg_invariance_loss = 0\n",
    "    net_itc_loss = 0\n",
    "    net_itm_loss = 0\n",
    "    for idx, data in (pbar := tqdm(enumerate(pretrain_dataloader), total = len(pretrain_dataloader))):\n",
    "        img, img_mask, txt, attn_mask, masked_toks, masked_attn_mask, mask_indices = data\n",
    "        \n",
    "        # vision\n",
    "        img = img.to(DEVICE)\n",
    "        img_mask = img_mask.to(DEVICE)\n",
    "        \n",
    "        # language\n",
    "        txt = txt.to(DEVICE)\n",
    "        attn_mask = attn_mask.to(DEVICE)\n",
    "        masked_toks = masked_toks.to(DEVICE)\n",
    "        masked_attn_mask = masked_attn_mask.to(DEVICE)\n",
    "\n",
    "        # indices for masked text: will be used for masked modeling\n",
    "        mask_indices = mask_indices.float().to(DEVICE)\n",
    "        \n",
    "        # # masked image\n",
    "        # masked_image = create_masked_image(img, img_mask)\n",
    "        flattened_img_mask = img_mask.float().flatten(1)\n",
    "        \n",
    "        # create masks for joint representation modeling\n",
    "        img_rep_masks = torch.cat([flattened_img_mask, torch.zeros_like(mask_indices)], axis = 1).unsqueeze(-1)\n",
    "        txt_rep_masks = torch.cat([torch.zeros_like(flattened_img_mask), mask_indices], axis = 1).unsqueeze(-1)\n",
    "        \n",
    "        \n",
    "        \n",
    "        # masked modeling pretraining\n",
    "        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):                           # casting to f16\n",
    "            # forward step for online network\n",
    "            c_img_m_txt, m_img_c_txt, img_txt_joint, mask_img_rep, txt_prediction, img_rep, txt_rep = online_network(img,\n",
    "                                                                                                      txt,\n",
    "                                                                                                      attn_mask,\n",
    "                                                                                                      image_text_matching = False,\n",
    "                                                                                                      masked_pos = flattened_img_mask,\n",
    "                                                                                                      masked_text = masked_toks)\n",
    "\n",
    "\n",
    "\n",
    "            # MIM loss\n",
    "            mim_loss_unimodal = online_network.get_mim_loss(img_rep, img, flattened_img_mask)\n",
    "            mim_loss_joint = online_network.get_joint_mim_loss(m_img_c_txt, img, flattened_img_mask)\n",
    "            mim_loss = mim_loss_unimodal + mim_loss_joint\n",
    "            \n",
    "            # MLM loss\n",
    "            mlm_loss = online_network.get_mlm_loss(txt_prediction, txt, masked_toks)\n",
    "            # mlm_loss_joint = online_network.get_joint_mlm_loss(c_img_m_txt, txt, masked_toks)\n",
    "            \n",
    "            # VICReg loss\n",
    "            vicreg_joint_clean = online_network.get_vicreg_loss(img_txt_joint)\n",
    "            vicreg_joint_c_img_m_txt = online_network.get_vicreg_loss(c_img_m_txt)\n",
    "            vicreg_joint_m_img_c_txt = online_network.get_vicreg_loss(m_img_c_txt)\n",
    "            \n",
    "            vicreg_invariance_1 = online_network.get_vicreg_invariance_loss(img_txt_joint, c_img_m_txt)\n",
    "            vicreg_invariance_2 = online_network.get_vicreg_invariance_loss(img_txt_joint, m_img_c_txt)\n",
    "            \n",
    "            vicreg_var_loss = (vicreg_joint_clean + vicreg_joint_c_img_m_txt + vicreg_joint_m_img_c_txt)/3.\n",
    "            vicreg_invariance_loss = (vicreg_invariance_1 + vicreg_invariance_2)/2.\n",
    "            \n",
    "            vicreg_loss = vicreg_var_loss + vicreg_invariance_loss\n",
    "\n",
    "            \n",
    "            # ITC loss\n",
    "            sim, itc_loss = online_network.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 = online_network.get_samples(sim)\n",
    "            \n",
    "            itm_labels = torch.cat([torch.ones(len(img)),torch.zeros(2*len(img))],\n",
    "                               dim=0).unsqueeze(1).float().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",
    "            joint_rep_negs = online_network.fusion(itm_img_feats, itm_txt_feats, itm_txt_attn)['last_hidden_state']\n",
    "            combined_reps = torch.vstack([img_txt_joint, joint_rep_negs])\n",
    "            \n",
    "            itm_outputs = online_network.itm__head(combined_reps[:, 0, :])\n",
    "            \n",
    "            \n",
    "            # softmax probabilities\n",
    "            itm_loss = itm_loss_fn(itm_outputs, itm_labels)\n",
    "            \n",
    "            # TOTAL LOSS\n",
    "            net_loss = (mim_loss) + \\\n",
    "                       (mlm_loss) + \\\n",
    "                       (vicreg_loss) + \\\n",
    "                       (itc_loss) + (itm_loss)\n",
    "            \n",
    "        scaler.scale(net_loss).backward()\n",
    "        scaler.unscale_(optim)\n",
    "        torch.nn.utils.clip_grad_norm_(online_network.parameters(), 1.)\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_mim_unimodal_loss+= mim_loss_unimodal.item()\n",
    "        net_mim_joint_loss+= mim_loss_joint.item()\n",
    "        \n",
    "        net_mlm_loss+= mlm_loss.item()\n",
    "        \n",
    "        net_vicreg_invariance_loss+= vicreg_invariance_loss.item()\n",
    "        net_vicreg_var_loss+= vicreg_var_loss.item()\n",
    "        \n",
    "        net_itc_loss+= itc_loss.item()\n",
    "        net_itm_loss+= itm_loss.item()\n",
    "        pretrain_loss+= net_loss.item()\n",
    "        pbar.set_description(f\"Train Loss: {pretrain_loss/num_samples}\")\n",
    "        \n",
    "\n",
    "    wandb.log({\n",
    "            'epoch': epoch,\n",
    "            'pretrain_loss': pretrain_loss/num_samples,\n",
    "            'mim_loss': net_mim_unimodal_loss/num_samples,\n",
    "            'mim_joint_loss': net_mim_joint_loss/num_samples,\n",
    "            'mlm_loss': net_mlm_loss/num_samples,\n",
    "            # 'mlm_joint_loss': net_mlm_joint_loss/num_samples,\n",
    "\n",
    "            'vicreg_variance_loss': net_vicreg_var_loss/num_samples,\n",
    "            'vicreg_invariance_loss': net_vicreg_invariance_loss/num_samples,\n",
    "\n",
    "            'itc_loss': net_itc_loss/num_samples,\n",
    "            'itm_loss': net_itm_loss/num_samples,\n",
    "    })\n",
    "    save_path = '{}_{}.pth'.format(MODEL_SAVE_PATH, epoch)\n",
    "    torch.save(\n",
    "            {\n",
    "            'epoch': epoch,\n",
    "            'online_model_state_dict': online_network.state_dict(),\n",
    "            'optim_state_dict': optim.state_dict()\n",
    "            },\n",
    "        save_path\n",
    "        )\n",
    "    # if (epoch-warmup_epochs+1) % 10 == 0:\n",
    "    #     wandb.save(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/ml/ML Projects/Denoising MAMO/wandb/run-20240423_071811-l4qpp8th/files/Models/ViT-S BERT-S (fixed everything)/VicVLM (no MLM and MIM Reg) - 5.0 - 3.0 - 2.0/checkpoint_final.pth']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "save_path = '{}_{}.pth'.format(MODEL_SAVE_PATH, 'final')\n",
    "torch.save(online_network.state_dict(), save_path)\n",
    "wandb.save(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>▁▁▂▃▃▃▄▅▅▅▆▇▇▇█</td></tr><tr><td>itc_loss</td><td>█▂▂▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>itm_loss</td><td>███▆▅▄▃▃▂▂▁▁▁▁▁</td></tr><tr><td>mim_joint_loss</td><td>█▄▄▃▃▂▂▂▂▁▁▁▁▁▁</td></tr><tr><td>mim_loss</td><td>█▃▂▂▂▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>mlm_loss</td><td>█▄▄▃▃▂▂▂▁▁▁▁▁▁▁</td></tr><tr><td>pretrain_loss</td><td>█▂▂▂▂▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>vicreg_invariance_loss</td><td>█▃▃▂▂▂▁▁▁▁▁▁▁▁▁</td></tr><tr><td>vicreg_variance_loss</td><td>█▄▃▂▂▂▁▁▁▁▁▁▁▁▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>epoch</td><td>14</td></tr><tr><td>itc_loss</td><td>0.05338</td></tr><tr><td>itm_loss</td><td>0.15449</td></tr><tr><td>mim_joint_loss</td><td>0.30698</td></tr><tr><td>mim_loss</td><td>0.16178</td></tr><tr><td>mlm_loss</td><td>2.19843</td></tr><tr><td>pretrain_loss</td><td>6.40138</td></tr><tr><td>vicreg_invariance_loss</td><td>0.09534</td></tr><tr><td>vicreg_variance_loss</td><td>3.43099</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">VicVLM (no MLM and MIM Reg) - 5.0 - 3.0 - 2.0 - ViT-S, BERT-S</strong> at: <a href='https://wandb.ai/madhava20217/MAMO%20-%20Pretrain/runs/l4qpp8th' target=\"_blank\">https://wandb.ai/madhava20217/MAMO%20-%20Pretrain/runs/l4qpp8th</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 1 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20240423_071811-l4qpp8th/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torchnew",
   "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": 2
}