mvitac.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e872cd85492721e",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## MVitAC Pretraining\n",
    "This notebook contains the code for the pretraining of the MVitAC model. The model is trained using the Calandra and Touch-and-Go datasets, which contains both vision and tactile modalities. The model is trained using contrastive learning, which is a self-supervised learning method. The model is trained to maximize the agreement between differently augmented views of the same modality and minimize the agreement between differently augmented views of different modalities. The model is trained for 800 epochs, with a batch size of 256, and a learning rate of 0.0001. The model is trained on a single NVIDIA GeForce RTX 4090 Ti GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc76e7b148afca41",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## **1. Library and Module Imports**\n",
    "\n",
    "This section imports necessary libraries and modules for the rest of the notebook.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbedf8ea2e0fa574",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import random_split\n",
    "from data_aug.contrastive_learning_dataset import ContrastiveLearningDataset\n",
    "import wandb\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import logging\n",
    "from utils import accuracy, save_checkpoint\n",
    "from model import MultiModalMoCo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7559717e5330731a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## **2. Configuration Setup**\n",
    "\n",
    "The `config` dictionary outlines several parameters:\n",
    "\n",
    "- `train_dataset_name`: Specifies the name of the training dataset.\n",
    "- `num_channels`: Details the number of channels (6 for 'calandra_label' and 3 for 'tag').\n",
    "- `epochs`: Determines the total number of training epochs.\n",
    "- Various hyperparameters are set, including the learning rate, weight decay, momentum, and temperature.\n",
    "- Model specifications include the neural network model type (`nn_model`) and dimensions (`intra_dim` and `inter_dim`).\n",
    "- Weight parameters for intra-vision, intra-tactile, and inter-modal contrastive losses are defined.\n",
    "- Flags like `pretrained_encoder` and `use_wandb` indicate whether a pretrained encoder should be used and if Weights & Biases logging is active.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82ecc33f1b9e9cd8",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "config={\n",
    "    # \"train_dataset_name\": 'calandra_label_train',\n",
    "    \"train_dataset_name\": 'tag_train',\n",
    "    # \"data_folder\": \"objects_split_object_wise/\",\n",
    "    \"data_folder\": \"TAG_dataset\",\n",
    "    \"model_name\": \"TAG\",\n",
    "    # \"num_channels\": 6, # should be 6 for calandra_label and 3 for tag\n",
    "    \"num_channels\": 3, # should be 6 for calandra_label and 3 for tag\n",
    "    \"epochs\": 240,\n",
    "    \"log_every_n_epochs\": 10,\n",
    "    \"batch_size\": 256,\n",
    "    \"num_workers\": 16,\n",
    "    \"momentum\": 0.99,\n",
    "    \"temperature\": 0.07,\n",
    "    \"lr\": 1e-3,\n",
    "    \"weight_decay\": 1e-6,\n",
    "    \"nn_model\": 'resnet18',\n",
    "    \"intra_dim\": 128,\n",
    "    \"inter_dim\": 128,\n",
    "    \"weight_inter_tv\": 1,\n",
    "    \"weight_inter_vt\": 1,\n",
    "    \"weight_intra_vision\": 1,\n",
    "    \"weight_intra_tactile\": 1,\n",
    "    \"pretrained_encoder\": True,\n",
    "    \"use_wandb\": True\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16a13371d2b33e85",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## **3. Data Preparation**\n",
    "\n",
    "- The `ContrastiveLearningDataset` is instantiated with a specified root folder.\n",
    "- The training dataset is retrieved using the specified name and splits it into training and validation subsets.\n",
    "- A DataLoader (`train_loader`) is established for the efficient fetching of training samples.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dc9470762c4db43",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dataset = ContrastiveLearningDataset(root_folder=config['data_folder'])\n",
    "train_dataset = dataset.get_dataset(config['train_dataset_name'], 2)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True,\n",
    "                                           num_workers=config['num_workers'], drop_last=False, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45d7e6b91f216006",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## **4. Model Initialization**\n",
    "\n",
    "- The `MultiModalMoCo` model, designed for multi-modal contrastive learning accommodating both vision and tactile modalities, is initialized using parameters from the `config` dictionary.\n",
    "- Different contrastive losses are given various weight parameters.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b18c80ebb8e04165",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# load model\n",
    "model = MultiModalMoCo(n_channels=config['num_channels'], m=config['momentum'], T=config['temperature'],\n",
    "                       intra_dim=config['intra_dim'], inter_dim=config['inter_dim'], nn_model=config['nn_model'],\n",
    "                       weight_inter_tv=config['weight_inter_tv'], weight_inter_vt=config['weight_inter_vt'],\n",
    "                       weight_intra_vision=config['weight_intra_vision'], weight_intra_tactile=config['weight_intra_tactile'],\n",
    "                       pretrained_encoder=config['pretrained_encoder'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae12ad8ad2631d02",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## **5. Training Setup and Execution**\n",
    "\n",
    "- The training device (either GPU or CPU) is determined and the model is shifted to this device.\n",
    "- The optimizer and loss function are delineated.\n",
    "- An extensive training loop ensures the model is trained for the specified number of epochs.\n",
    "- Model performance metrics, such as accuracy, are computed.\n",
    "- Model checkpoints, based on performance, are saved.\n",
    "- Weights & Biases logs the training metrics.\n",
    "- At the end of training, the final model checkpoint is saved and log messages are generated.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a426de0c22ced2",
   "metadata": {
    "collapsed": false,
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "# define loss function (criterion) and optimizer\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Training with gpu: {device}.\")\n",
    "model = model.to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], weight_decay=config['weight_decay'])\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.9)\n",
    "writer = SummaryWriter()\n",
    "logging.basicConfig(filename=os.path.join(writer.log_dir, 'training.log'), level=logging.DEBUG)\n",
    "criterion = torch.nn.CrossEntropyLoss().to(device)\n",
    "# Set number of training epochs\n",
    "logging.info(f\"Start MViTaC training for {config['epochs']} epochs.\")\n",
    "logging.info(f\"Training with gpu: {device}.\")\n",
    "best_acc = 0\n",
    "if config['use_wandb']:\n",
    "    wandb.init(project=\"mvitac_pretraining\", config=config)\n",
    "    # name the model\n",
    "    wandb.run.name = f\"{config['nn_model']}_lr_{config['lr']}_batch_{config['batch_size']}_epochs_{config['epochs']}\"\n",
    "    \n",
    "for epoch in range(config['epochs']):\n",
    "    loss_epoch, vis_loss_intra_epoch, tac_loss_intra_epoch, vis_tac_inter_epoch, tac_vis_inter_epoch = 0, 0, 0, 0, 0\n",
    "    pbar = tqdm(train_loader)  # Wrap train_loader with tqdm\n",
    "    for idx, values in enumerate(pbar):  # Use enumerate to get idx\n",
    "        x_vision_q, x_vision_k, x_tactile_q, x_tactile_k, label = values\n",
    "        model.train()\n",
    "        # send to device\n",
    "        x_vision_q = x_vision_q.to(device, non_blocking=True)\n",
    "        x_vision_k = x_vision_k.to(device, non_blocking=True)\n",
    "\n",
    "        x_tactile_q = x_tactile_q.to(device, non_blocking=True)\n",
    "        x_tactile_k = x_tactile_k.to(device, non_blocking=True)\n",
    "\n",
    "        # Forward pass to get the loss\n",
    "        loss, vis_loss_intra, tac_loss_intra, vis_tac_inter, tac_vis_inter, logits, labels = model(x_vision_q, x_vision_k, x_tactile_q, x_tactile_k)\n",
    "        loss_epoch += loss.item()\n",
    "        vis_loss_intra_epoch += vis_loss_intra.item()\n",
    "        tac_loss_intra_epoch += tac_loss_intra.item()\n",
    "        vis_tac_inter_epoch += vis_tac_inter.item()\n",
    "        tac_vis_inter_epoch += tac_vis_inter.item()\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # update pbar message\n",
    "        pbar.set_description(f\"Epoch {epoch}:\\tTrain loss: {loss_epoch/ (idx + 1):.2f}\")\n",
    "    \n",
    "    if epoch % config['log_every_n_epochs'] == 0:\n",
    "        top1, top5 = accuracy(logits, labels, topk=(1, 5))\n",
    "        writer.add_scalar('loss', loss_epoch / len(train_loader), global_step=epoch)\n",
    "        writer.add_scalar('loss/vis_loss_intra', vis_loss_intra_epoch / len(train_loader), global_step=epoch)\n",
    "        writer.add_scalar('loss/tac_loss_intra', tac_loss_intra_epoch / len(train_loader), global_step=epoch)\n",
    "        writer.add_scalar('loss/vis_tac_inter', vis_tac_inter_epoch / len(train_loader), global_step=epoch)\n",
    "        writer.add_scalar('loss/tac_vis_inter', tac_vis_inter_epoch / len(train_loader), global_step=epoch)\n",
    "        writer.add_scalar('acc/top1', top1[0], global_step=epoch)\n",
    "        writer.add_scalar('acc/top5', top5[0], global_step=epoch)\n",
    "        writer.add_scalar('learning_rate', scheduler.get_last_lr()[0], global_step=epoch)\n",
    "        if top1[0] > best_acc:\n",
    "            best_acc = top1[0]\n",
    "            # save both the vision and tactile models\n",
    "\n",
    "            save_checkpoint({\n",
    "                'epoch': epoch,\n",
    "                'arch': 'resnet18',\n",
    "                'state_dict_vis': model.vision_base_q.state_dict(),\n",
    "                'state_dict_tac': model.tactile_base_q.state_dict(),\n",
    "                'optimizer': optimizer.state_dict(),\n",
    "            }, filename=os.path.join(writer.log_dir, 'model_{}_best_{}.pth'.format(epoch, config[\"model_name\"])))\n",
    "        # torch.save(state, f'models/calandra/model_{args.task}_{epoch}_{args.batch_size}_best_object_wise_05_t05.pth')\n",
    "        if config['use_wandb']:\n",
    "            wandb.log({\"epoch\": epoch, \"loss\": loss_epoch / len(train_loader), \"vis_loss_intra\": vis_loss_intra_epoch / len(train_loader),\n",
    "                       \"tac_loss_intra\": tac_loss_intra_epoch / len(train_loader), \"vis_tac_inter\": vis_tac_inter_epoch / len(train_loader),\n",
    "                       \"tac_vis_inter\": tac_vis_inter_epoch / len(train_loader), \"top1\": top1[0], \"top5\": top5[0],\n",
    "                       \"learning_rate\": scheduler.get_last_lr()[0]})\n",
    "            wandb.save('models/{}/model_{}_best.pth'.format(config[\"model_name\"], epoch))\n",
    "        \n",
    "    # warmup for the first 10 epochs\n",
    "    if epoch >= 10:\n",
    "        scheduler.step()\n",
    "    logging.debug(f\"Epoch: {epoch}\\tLoss: {loss_epoch / len(train_loader)}\\tTop1: {top1[0]}\\tTop5: {top5[0]}\")\n",
    "\n",
    "    logging.info(\"Training has finished.\")\n",
    "    # save model checkpoints\n",
    "    checkpoint_name = 'checkpoint_{:04d}.pth.tar'.format(config['epochs'])\n",
    "    save_checkpoint({\n",
    "        'epoch': config['epochs'],\n",
    "        'arch': config['nn_model'],\n",
    "        'state_dict_vis': model.vision_base_q.state_dict(),\n",
    "        'state_dict_tac': model.tactile_base_q.state_dict(),   \n",
    "        'optimizer': optimizer.state_dict(),\n",
    "    }, filename=os.path.join(writer.log_dir, checkpoint_name))\n",
    "    logging.info(f\"Model checkpoint and metadata has been saved at {writer.log_dir}.\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}