train_google_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Downloading the Dataset\n",
    "Make sure you are connected to an environment with an Nvidia GPU supporting CUDA. If you are training/evaluating the model, download the Vimeo-90k septuplet dataset from the official source:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget http://data.csail.mit.edu/tofu/dataset/vimeo_septuplet.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unzip the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!unzip vimeo_septuplet.zip"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execution on Google Colab\n",
    "\n",
    "First, clone the project repository and navigate to the root directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/starboi-63/ArTEMIS.git\n",
    "\n",
    "%cd ArTEMIS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install Python dependencies from `requirements.txt`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the CUDA driver installed on the current compute instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvcc --version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that CUDA is working as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "print(\"PyTorch version: \", torch.__version__)\n",
    "print(\"PyTorch CUDA version: \", torch.version.cuda)\n",
    "print(\"CUDA available: \", torch.cuda.is_available())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mounting to Drive\n",
    "\n",
    "Mount to Google Drive to save logs and checkpoints in real-time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logging, Output, and Checkpoints\n",
    "Define which directories to save logs, outputs, and model training checkpoints will be saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "logging_dir = '/content/drive/My Drive/Deep Learning/ArTEMIS/training/tensorboard_logs'\n",
    "checkpoint_dir = '/content/drive/My Drive/Deep Learning/ArTEMIS/training/tensorboard_logs/' # fill in the rest of the path\n",
    "output_dir = '/content/drive/My Drive/Deep Learning/ArTEMIS/training/output'\n",
    "data_dir = '/content/vimeo_septuplet'\n",
    "\n",
    "if not os.path.exists(logging_dir):\n",
    "    os.makedirs(logging_dir)\n",
    "    print(f\"Created directory at {logging_dir}.\")\n",
    "else:\n",
    "    print(f\"Directory {logging_dir} already exists.\\n\")\n",
    "\n",
    "if not os.path.exists(checkpoint_dir):\n",
    "    os.makedirs(checkpoint_dir)\n",
    "    print(f\"Created directory at {checkpoint_dir}.\")\n",
    "else:\n",
    "    print(f\"Directory {checkpoint_dir} already exists.\")\n",
    "\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "    print(f\"Created directory at {output_dir}.\")\n",
    "else:\n",
    "    print(f\"Directory {output_dir} already exists.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensor Board:\n",
    "\n",
    "Run the following cells to launch tensorboard, which helps visualize training progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --logdir \"{logging_dir}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the Model!\n",
    "\n",
    "### Command-Line Arguments\n",
    "\n",
    "#### Key Arguments for Training/Testing\n",
    "- `--model`: Model to train/test.\n",
    "- `--mode`: Mode to run the model in. Options are `train`, `test`, and `interpolate`.\n",
    "- `--dataset`: Dataset to train on.\n",
    "- `--data_dir`: Path to the dataset.\n",
    "- `--checkpoint_dir`: Directory saving intermediate model states.\n",
    "- `--use_checkpoint`: Whether to resume training from a checkpoint at `checkpoint_dir`.\n",
    "- `--log_dir`: Directory saving TensorBoard logs.\n",
    "- `--output_dir`: Directory in which predictions and ground-truth images are saved.\n",
    "- `--batch_size`: Batch size for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python main.py --model ArTEMIS --mode train --dataset vimeo90K_septuplet --data_dir \"{data_dir}\" --log_dir \"{logging_dir}\" --output_dir \"{output_dir}\" --checkpoint_dir \"{checkpoint_dir}\" --use_checkpoint --batch_size 12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Keep Session Alive When Idling\n",
    "\n",
    "To prevent the session from disconnecting due to inactivity, you can run the following script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from tqdm import tqdm\n",
    "\n",
    "# Duration to keep the session alive in hours\n",
    "hours_to_keep_alive = 8\n",
    "\n",
    "# Convert hours to seconds\n",
    "seconds_to_keep_alive = hours_to_keep_alive * 3600\n",
    "\n",
    "# Initialize the tqdm progress bar\n",
    "with tqdm(total=seconds_to_keep_alive, desc=\"Keeping session alive\") as pbar:\n",
    "    start_time = time.time()\n",
    "\n",
    "    while (time.time() - start_time) < seconds_to_keep_alive:\n",
    "        pbar.update(time.time() - (start_time + pbar.n))\n",
    "        time.sleep(1)  # Wait for a second\n",
    "\n",
    "    pbar.n = seconds_to_keep_alive  # Ensure the progress bar fills completely\n",
    "    pbar.refresh()  # Update the progress bar one last time\n",
    "\n",
    "print(\"Session kept alive successfully!\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}