train.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The dojo"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup training device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git init\n",
    "# !git remote add origin https://github.com/andreasalstrup/chrome-vision.git\n",
    "# !git pull origin <branch>\n",
    "# !git reset --hard FETCH_HEAD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import device, nn\n",
    "import torchvision\n",
    "import torchvision.io as io\n",
    "\n",
    "#torch.manual_seed(42)\n",
    "\n",
    "print(f'PyTorch version: {torch.__version__}\\ntorchvision version: {torchvision.__version__}')\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f'Using device: {device}')\n",
    "\n",
    "ok = io.read_image(\"data/E45Vejle_1011.jpg\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.models as models\n",
    "\n",
    "# Data loading\n",
    "BATCH_SIZE = 256\n",
    "IMAGE_RESIZE = 64\n",
    "\n",
    "# Model\n",
    "OUT_FEATURES = 128\n",
    "QUEUE_SIZE = 65536\n",
    "MOMENTUM = 0.9\n",
    "SOFTMAX_TEMP = 0.07\n",
    "\n",
    "# Encoder\n",
    "ENCODER = models.resnet50\n",
    "\n",
    "# Optimizer\n",
    "OPTIMIZER = torch.optim.Adam\n",
    "LEARNING_RATE = 0.001\n",
    "ADJUST_LEARNING_RATE = False\n",
    "BETAS = (0.9, 0.999)\n",
    "EPS = 1e-08\n",
    "WEIGHT_DECAY = 1e-5\n",
    "\n",
    "# Loss function\n",
    "LOSS_FN = nn.CrossEntropyLoss()\n",
    "\n",
    "# Training loop\n",
    "EPOCHS = 400"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data loading"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine train data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chrome_utils.merge_dir import MergeDir\n",
    "\n",
    "src_dir = 'data/leftImg8bit/train'\n",
    "dst_dir = 'data/leftImg8bit/train_combined'\n",
    "dst_index = 'data/leftImg8bit/indices/trainIndex/combined.csv'\n",
    "\n",
    "MergeDir(src_dir, dst_dir, dst_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chrome_utils.merge_dir import MergeDir\n",
    "\n",
    "src_dir = 'data/leftImg8bit/test'\n",
    "dst_dir = 'data/leftImg8bit/test_combined'\n",
    "dst_index = 'data/leftImg8bit/indices/testIndex/combined.csv'\n",
    "\n",
    "MergeDir(src_dir, dst_dir, dst_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using chromecut to prepare datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chrome_cut import ChromeCut\n",
    "cutter = ChromeCut()\n",
    "cutter.CutImagesInFolder(annotations_file=\"data/leftImg8bit/indices/trainIndex/combined.csv\",\n",
    "                         img_dir=\"data/leftImg8bit/train_combined/\", \n",
    "                         name=\"trainCombinedCut\", \n",
    "                         new_img_dir=\"data/leftImg8bit/train/cut/trainCombinedCut\",\n",
    "                         new_annotations_file_location=\"data/leftImg8bit/indices/trainIndex\")\n",
    "\n",
    "cutter.CutImagesInFolder(annotations_file=\"data/leftImg8bit/indices/testIndex/combined.csv\",\n",
    "                         img_dir=\"data/leftImg8bit/test_combined/\", \n",
    "                         name=\"testCombinedCut\",\n",
    "                         new_img_dir= \"data/leftImg8bit/test/cut/testCombinedCut\",\n",
    "                         new_annotations_file_location= \"data/leftImg8bit/indices/testIndex\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "from chrome_utils.transforms import ContrastiveTransform\n",
    "\n",
    "transform_MoCoV1 = ContrastiveTransform(\n",
    "                        transforms.Compose([\n",
    "                            transforms.ToPILImage(),\n",
    "                            transforms.Resize((IMAGE_RESIZE, IMAGE_RESIZE)),\n",
    "                            transforms.RandomResizedCrop(IMAGE_RESIZE, scale=(0.2, 1.0)), # 224 -> 64 \n",
    "                            transforms.RandomGrayscale(p=0.2),\n",
    "                            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),\n",
    "                            transforms.RandomHorizontalFlip(),\n",
    "                            transforms.ToTensor(),\n",
    "                            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "                            ])\n",
    "                        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from custom_image_dataset import CustomImageDataset\n",
    "\n",
    "TRAIN_DATA = CustomImageDataset(\"data/leftImg8bit/indices/trainIndex/trainCombinedCut.csv\",\"data/leftImg8bit/train/cut/trainCombinedCut\", transform=transform_MoCoV1)\n",
    "TEST_DATA = CustomImageDataset(\"data/leftImg8bit/indices/testIndex/testCombinedCut.csv\",\"data/leftImg8bit/test/cut/testCombinedCut\", transform=transform_MoCoV1)\n",
    "\n",
    "TRAIN_DATALOADER = DataLoader(TRAIN_DATA,\n",
    "                              batch_size=BATCH_SIZE,\n",
    "                              shuffle=True,\n",
    "                              drop_last=True)\n",
    "\n",
    "TEST_DATALOADER = DataLoader(TEST_DATA,\n",
    "                             batch_size=BATCH_SIZE,\n",
    "                             shuffle=False,\n",
    "                             drop_last=True)\n",
    "\n",
    "print(f'Len of train dataloader: {len(TRAIN_DATALOADER)} batches of {BATCH_SIZE}')\n",
    "print(f'Len of test dataloader: {len(TEST_DATALOADER)} batches of {BATCH_SIZE}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Check out what's inside the training dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_features_batch = next(iter(TRAIN_DATALOADER))\n",
    "\n",
    "print(F\"Train features query_image shape: {train_features_batch[0].shape}\")\n",
    "print(F\"Train features key_image shape: {train_features_batch[1].shape}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ChromeMoCo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.chrome_moco import ChromeMoCo\n",
    "\n",
    "model = ChromeMoCo(base_encoder=ENCODER,\n",
    "                  feature_dim=OUT_FEATURES,\n",
    "                  queue_size=QUEUE_SIZE,\n",
    "                  momentum=MOMENTUM,\n",
    "                  softmax_temp=SOFTMAX_TEMP).to(device)\n",
    "\n",
    "optimizer = OPTIMIZER(params=model.parameters(), lr=LEARNING_RATE, betas=BETAS, eps=EPS, weight_decay=WEIGHT_DECAY)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ChromeMoCoV2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.chrome_moco_v2 import ChromeMoCoV2\n",
    "\n",
    "OUT_FEATURES = (1000, 10)\n",
    "\n",
    "model = ChromeMoCoV2(base_encoder=ENCODER,\n",
    "                  feature_dim=OUT_FEATURES,\n",
    "                  queue_size=QUEUE_SIZE,\n",
    "                  momentum=MOMENTUM,\n",
    "                  softmax_temp=SOFTMAX_TEMP).to(device)\n",
    "\n",
    "optimizer = OPTIMIZER(params=model.parameters(), lr=LEARNING_RATE, betas=BETAS, eps=EPS, weight_decay=WEIGHT_DECAY)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop - MoCo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, os.path as path\n",
    "from timeit import default_timer as timer\n",
    "from tqdm.auto import tqdm\n",
    "from model.evaluation import train_step, test_step, adjust_learning_rate # use torchmetrics.utilities.data.select_topk\n",
    "from chrome_utils.model_utils import print_train_time, accuracy_top_k, saveModel, saveCheckpoint, loadCheckpoint\n",
    "from chrome_utils.show_progress import ShowProgress\n",
    "\n",
    "# Setup progress curves\n",
    "show_training = ShowProgress('Train')\n",
    "show_testing = ShowProgress('Test')\n",
    "time_start = timer()\n",
    "\n",
    "# Setup/Load checkpoint\n",
    "checkpoint_epoch = -1\n",
    "CHECKPOINT = 'model/models/checkpoint.pt'\n",
    "CHECKPOINT_BACKUP = 'model/models/checkpoint_backup.pt'\n",
    "\n",
    "if path.exists(CHECKPOINT):\n",
    "    (model, optimizer, checkpoint_epoch, show_training, show_testing, TRAIN_DATALOADER, TEST_DATALOADER) = loadCheckpoint(CHECKPOINT, model, optimizer)\n",
    "\n",
    "# Train loop\n",
    "for epoch in tqdm(range(EPOCHS)):\n",
    "\n",
    "    print(f'\\n\\tEpoch: {epoch}\\n')\n",
    "\n",
    "    # Skip epochs already trained\n",
    "    if epoch <= checkpoint_epoch:\n",
    "        continue\n",
    "    \n",
    "    saveCheckpoint(CHECKPOINT_BACKUP,\n",
    "                   model,\n",
    "                   optimizer,\n",
    "                   epoch - 1,\n",
    "                   show_training,\n",
    "                   show_testing,\n",
    "                   TRAIN_DATALOADER,\n",
    "                   TEST_DATALOADER)\n",
    "    \n",
    "    if ADJUST_LEARNING_RATE:\n",
    "        current_learning_rate = optimizer.param_groups[0]['lr']\n",
    "        optimizer = adjust_learning_rate(optimizer, epoch, EPOCHS, current_learning_rate)\n",
    "        print(optimizer.param_groups[0]['lr'])\n",
    "    \n",
    "    # Train\n",
    "    (train_loss, train_top1, train_top5) = train_step(model=model,\n",
    "                                                     data_loader=TRAIN_DATALOADER,\n",
    "                                                     loss_fn=LOSS_FN,\n",
    "                                                     optimizer=optimizer,\n",
    "                                                     accuracy_fn=accuracy_top_k,\n",
    "                                                     device=device)\n",
    "    # Draw train curve\n",
    "    print(f'Train loss: {train_loss:.5f} | Train acc1: {train_top1:.2f}% | Train acc5: {train_top5:.2f}%')\n",
    "    show_training.appendData(train_loss, train_top1, train_top5)\n",
    "    show_training.draw_curve(epoch)\n",
    "    \n",
    "    # Test\n",
    "    if epoch % 10 == 0:\n",
    "        (test_loss, test_top1, test_top5) = test_step(model=model,\n",
    "                                                    data_loader=TEST_DATALOADER,\n",
    "                                                    loss_fn=LOSS_FN,\n",
    "                                                    accuracy_fn=accuracy_top_k,\n",
    "                                                    device=device)\n",
    "        \n",
    "        # Draw test curve\n",
    "        print(f'Test loss: {test_loss:.5f} | Test acc1: {test_top1:.2f}% | Test acc5: {test_top5:.2f}%')\n",
    "        show_testing.appendData(test_loss, test_top1, test_top5)\n",
    "        show_testing.draw_curve(epoch)\n",
    "    \n",
    "    saveCheckpoint(CHECKPOINT,\n",
    "                   model,\n",
    "                   optimizer,\n",
    "                   epoch,\n",
    "                   show_training,\n",
    "                   show_testing,\n",
    "                   TRAIN_DATALOADER,\n",
    "                   TEST_DATALOADER)\n",
    "\n",
    "# Save model and curves\n",
    "NAME = f\"{model.__class__.__name__}_BatchSize{BATCH_SIZE}_OutFeat{OUT_FEATURES}_LR{LEARNING_RATE}_Adj{ADJUST_LEARNING_RATE}_ImageSize{IMAGE_RESIZE}_Epochs{EPOCHS}\"\n",
    "saveModel(\"model/models\", f\"{NAME}.pt\", model)\n",
    "show_training.saveFig(f\"model/models/{NAME}_train.png\")\n",
    "show_testing.saveFig(f\"model/models/{NAME}_test.png\")\n",
    "\n",
    "# Remove checkpoints\n",
    "if path.exists(CHECKPOINT):\n",
    "    os.remove(CHECKPOINT)\n",
    "    os.remove(CHECKPOINT_BACKUP)\n",
    "\n",
    "# Print time taken\n",
    "time_end = timer()\n",
    "total_train = print_train_time(time_start, time_end, str(next(model.parameters()).device))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chrome_utils.merge_dir import MergeDir\n",
    "\n",
    "src_dir = 'data/leftImg8bit/val'\n",
    "dst_dir = 'data/leftImg8bit/val_combined'\n",
    "dst_index = 'data/leftImg8bit/indices/valIndex/combined.csv'\n",
    "\n",
    "MergeDir(src_dir, dst_dir, dst_index)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using chromecut to prepare datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chrome_cut import ChromeCut\n",
    "cutter = ChromeCut()\n",
    "cutter.CutImagesInFolder(annotations_file=\"data/leftImg8bit/indices/valIndex/combined.csv\",\n",
    "                         img_dir=\"data/leftImg8bit/val_combined/\", \n",
    "                         name=\"valCombinedCut\", \n",
    "                         new_img_dir=\"data/leftImg8bit/val/cut/valCombinedCut\",\n",
    "                         new_annotations_file_location=\"data/leftImg8bit/indices/valIndex\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from custom_image_dataset import CustomImageDataset\n",
    "\n",
    "VAL_DATA = CustomImageDataset(\"data/leftImg8bit/indices/valIndex/valCombinedCut.csv\",\"data/leftImg8bit/val/cut/valCombinedCut\", transform=transform_MoCoV1)\n",
    "\n",
    "VAL_DATALOADER = DataLoader(VAL_DATA,\n",
    "                              batch_size=BATCH_SIZE,\n",
    "                              shuffle=False,\n",
    "                              drop_last=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "from model.chrome_moco import ChromeMoCo\n",
    "from model.evaluation import evaluate_models\n",
    "from chrome_utils.model_utils import accuracy_top_k\n",
    "\n",
    "model = ChromeMoCo(base_encoder=ENCODER,\n",
    "                  feature_dim=OUT_FEATURES,\n",
    "                  queue_size=QUEUE_SIZE,\n",
    "                  momentum=MOMENTUM,\n",
    "                  softmax_temp=SOFTMAX_TEMP).to(device)\n",
    "\n",
    "# Evaluate models in directory\n",
    "DIR = \"model/models/\"\n",
    "NAME = \"eval.png\"\n",
    "df = evaluate_models(model=model,\n",
    "                     models_dir=DIR,\n",
    "                     data_loader=VAL_DATALOADER,\n",
    "                     loss_fn=LOSS_FN,\n",
    "                     accuracy_fn=accuracy_top_k,\n",
    "                     device=device)\n",
    "\n",
    "# Plot results\n",
    "df.set_index('model_name', inplace=True)\n",
    "ax = df.plot(kind='barh', figsize=(10, 8))\n",
    "ax.set_xlabel('Accuracy / Loss')\n",
    "\n",
    "for i, (acc1, acc5, loss) in enumerate(zip(df['model_acc1'], df['model_acc5'], df['model_loss'])):\n",
    "    acc1_width = ax.get_xlim()[1] * (acc1 / 100)\n",
    "    acc5_width = ax.get_xlim()[1] * (acc5 / 100)\n",
    "    loss_width = ax.get_xlim()[1] * (loss / 100)\n",
    "    acc1 = round(acc1, 2)\n",
    "    ax.text(acc1_width - 5, i, f'{round(acc1, 2)}%', color='white', fontweight='bold')\n",
    "    ax.text(acc5_width - 5, i + 0.15, f'{round(acc5, 2)}%', color='white', fontweight='bold')\n",
    "    ax.text(loss_width + 1 , i - 0.2, f'{round(loss, 2)}', color='black', fontweight='bold')\n",
    "\n",
    "plt.savefig(os.path.join(DIR, NAME), bbox_inches='tight')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example of how to run ChromeVision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromevision as cVision\n",
    "import torch\n",
    "import cv2\n",
    "\n",
    "#The model used for object detection\n",
    "model.load_state_dict(torch.load(\"model/models/ChromeMoCo_BatchSize256_OutFeat10_LR0.001_AdjFalse_ImageSize64_Epochs200.pt\", map_location=torch.device(device)))\n",
    "chromeModel = cVision.Chromevision(model, ChromeCut())                \n",
    "identifiedImage = chromeModel.identify(\"data/leftImg8bit/val/munster/munster_000039_000019_leftImg8bit.png\")\n",
    "\n",
    "cv2.imshow('Result', identifiedImage)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cv",
   "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.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "f3fba41177e14bd00b55b5b38da464dfe75c2a11f18b855391927c023f3b0184"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}