MIMO_U_Net_NYUv2_depth.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "V100",
      "authorship_tag": "ABX9TyN6lFQFNEpIP3zuwQTpznu3",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/antonbaumann/MIMO-Unet/blob/main/MIMO_U_Net_NYUv2_depth.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Download data and dependencies"
      ],
      "metadata": {
        "id": "nQEexLS70FwV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard"
      ],
      "metadata": {
        "id": "U3ohfXOm3iPC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "\n",
        "# clone MIMO U-Net repository\n",
        "!rm -r MIMO-Unet; git clone https://github.com/antonbaumann/MIMO-Unet.git\n",
        "\n",
        "# add repository to PATH\n",
        "sys.path.append('/content/MIMO-Unet/')\n",
        "\n",
        "# install MIMO U-Net dependencies\n",
        "!pip install -r MIMO-Unet/requirements.txt"
      ],
      "metadata": {
        "id": "axO11WLdfyDk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# download NYU Depth Dataset V2\n",
        "!wget -c https://www.dropbox.com/s/qtab28cauzalqi7/depth_data.tar.gz?dl=1 -O depth_data.tar.gz\n",
        "!mkdir data && tar -xzvf depth_data.tar.gz -C data"
      ],
      "metadata": {
        "id": "qQ75Rb1FKlvQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from datetime import datetime\n",
        "\n",
        "import lightning.pytorch as pl\n",
        "from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint\n",
        "from lightning.pytorch.loggers import WandbLogger, TensorBoardLogger\n",
        "\n",
        "from mimo.models.mimo_unet import MimoUnetModel\n",
        "from mimo.tasks.depth.nyuv2_datamodule import NYUv2DepthDataModule\n",
        "from mimo.tasks.depth.callbacks import OutputMonitor, WandbMetricsDefiner"
      ],
      "metadata": {
        "id": "8KJBJ7iOSSJZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def default_callbacks(validation: bool = True) -> List[pl.Callback]:\n",
        "    callbacks = [\n",
        "        OutputMonitor(),\n",
        "        ModelCheckpoint(save_last=True),\n",
        "    ]\n",
        "    if validation:\n",
        "        callbacks_validation = [\n",
        "            ModelCheckpoint(\n",
        "                monitor=\"val_loss\",\n",
        "                save_top_k=1,\n",
        "                filename=\"epoch-{epoch}-step-{step}-valloss-{val_loss:.8f}-mae-{metric_val/mae_epoch:.8f}\",\n",
        "                auto_insert_metric_name=False,\n",
        "            ),\n",
        "        ]\n",
        "        callbacks += callbacks_validation\n",
        "    return callbacks"
      ],
      "metadata": {
        "id": "_7-Tu2PJU_f4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example: Pytorch Lightning\n",
        "Initialize `datamodule`, `model`, `logger` and `trainer`"
      ],
      "metadata": {
        "id": "LNSX9NxG1pOe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pl.seed_everything(1)\n",
        "\n",
        "dm = NYUv2DepthDataModule(\n",
        "    dataset_dir='data',\n",
        "    batch_size=32,\n",
        "    num_workers=3,\n",
        "    pin_memory=True,\n",
        "    normalize=True,\n",
        ")\n",
        "\n",
        "model = MimoUnetModel(\n",
        "    in_channels=3,\n",
        "    out_channels=2,\n",
        "    num_subnetworks=2,\n",
        "    filter_base_count=21,\n",
        "    center_dropout_rate=0.0,\n",
        "    final_dropout_rate=0.0,\n",
        "    encoder_dropout_rate=0.0,\n",
        "    core_dropout_rate=0.0,\n",
        "    decoder_dropout_rate=0.0,\n",
        "    loss_buffer_size=10,\n",
        "    loss_buffer_temperature=0.3,\n",
        "    input_repetition_probability=0.0,\n",
        "    batch_repetitions=1,\n",
        "    loss='laplace_nll',\n",
        "    weight_decay=0.0,\n",
        "    learning_rate=1e-3,\n",
        "    seed=1,\n",
        ")\n",
        "\n",
        "tensorboard_logger = TensorBoardLogger(\n",
        "    save_dir='/content/logs',\n",
        ")\n",
        "\n",
        "trainer = pl.Trainer(\n",
        "    callbacks=default_callbacks(),\n",
        "    accelerator='gpu',\n",
        "    devices=1,\n",
        "    precision=\"16-mixed\",\n",
        "    max_epochs=15,\n",
        "    default_root_dir='/content/runs',\n",
        "    log_every_n_steps=200,\n",
        "    logger=tensorboard_logger,\n",
        ")"
      ],
      "metadata": {
        "id": "t9SoBe8bVUxh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorboard --logdir /content/logs"
      ],
      "metadata": {
        "id": "Glu27Z6lFu0n"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trainer.started_at = str(datetime.now().isoformat(timespec=\"seconds\"))\n",
        "trainer.fit(model, dm)"
      ],
      "metadata": {
        "id": "zpuDgLIMeN63"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example: raw pytorch example\n",
        "You can build your own training loop as shown below if you do not want to use `lightning`"
      ],
      "metadata": {
        "id": "YSzGZjO04TD3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "from mimo.models.mimo_components.model import MimoUNet\n",
        "from mimo.models.mimo_components.loss_buffer import LossBuffer\n",
        "from mimo.models.utils import apply_input_transform\n",
        "from mimo.datasets.nyuv2 import NYUv2DepthDataset\n",
        "from mimo.losses import LaplaceNLL"
      ],
      "metadata": {
        "id": "wbSVX7Qq4yj2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "out_channels = 1\n",
        "num_subnetworks = 2\n",
        "input_repetition_probability = 0\n",
        "batch_repetitions = 1\n",
        "device = torch.device(\"cuda\")\n",
        "\n",
        "model = MimoUNet(\n",
        "    in_channels=3,\n",
        "    out_channels=out_channels * 2,\n",
        "    num_subnetworks=num_subnetworks,\n",
        "    filter_base_count=21,\n",
        ")\n",
        "\n",
        "model = model.to(device)\n",
        "\n",
        "loss_buffer = LossBuffer(\n",
        "    subnetworks=num_subnetworks,\n",
        "    temperature=0.3,\n",
        "    buffer_size=10,\n",
        ")\n",
        "\n",
        "criterion = LaplaceNLL()\n",
        "\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
        "\n",
        "data_train = NYUv2DepthDataset(\n",
        "    dataset_path='data/depth_train.h5',\n",
        "    normalize=True,\n",
        "    shuffle_on_load=True,\n",
        ")\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "    data_train,\n",
        "    batch_size=32,\n",
        "    shuffle=True,\n",
        ")\n",
        "\n",
        "data_test = NYUv2DepthDataset(\n",
        "    dataset_path='data/depth_test.h5',\n",
        "    normalize=True,\n",
        "    shuffle_on_load=False,\n",
        ")\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "    data_test,\n",
        "    batch_size=32,\n",
        "    shuffle=False,\n",
        ")"
      ],
      "metadata": {
        "id": "kmAGW0bb2TaR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# train for one epoch\n",
        "for batch in train_loader:\n",
        "  image, label = batch[\"image\"].to(device), batch[\"label\"].to(device)\n",
        "\n",
        "  image_transformed, label_transformed, mask_transformed = apply_input_transform(\n",
        "    image,\n",
        "    label,\n",
        "    mask=None,\n",
        "    num_subnetworks=num_subnetworks,\n",
        "    input_repetition_probability=input_repetition_probability,\n",
        "    batch_repetitions=batch_repetitions,\n",
        "  )\n",
        "\n",
        "  out = model(image_transformed)\n",
        "  p1 = out[:, :, :out_channels, ...]\n",
        "  p2 = out[:, :, out_channels:, ...]\n",
        "\n",
        "  y_pred = criterion.mode(p1, p2)\n",
        "\n",
        "  raw_loss = criterion.forward(p1, p2, label_transformed, reduce_mean=False, mask=None)\n",
        "  loss = raw_loss.mean(dim=(0, 2, 3, 4))\n",
        "  weights = loss_buffer.get_weights().to(loss.device)\n",
        "  loss_buffer.add(loss.detach())\n",
        "\n",
        "  weighted_loss = (loss * weights).mean()\n",
        "\n",
        "  weighted_loss.backward()\n",
        "  optimizer.step()\n",
        "\n",
        "  print(f'loss={loss.mean().cpu().detach()}')"
      ],
      "metadata": {
        "id": "SClwCfl06ntQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}