pretraining_image_fitting.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "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/emmanuel-nwogu/prefit_image_fitting/blob/master/pretraining_image_fitting.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount(\"/content/gdrive\", force_remount=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hDnTK12cTOH3",
        "outputId": "153321b3-c155-479d-d0c0-2ceccae3d4a3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/gdrive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ENhqXVSnRIce"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import sys\n",
        "import time\n",
        "\n",
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "# These most likely have to be dragged and dropped into colab from https://github.com/emmanuel-nwogu/prefit_image_fitting\n",
        "from cifair import TriplesCiFAIR10\n",
        "from modules import Siren\n",
        "from utils import Params, FittingMode, get_mgrid, fit_one_image, save_run_info\n",
        "\n",
        "SEED_FOR_DATASET = 42  # For data loading order sake."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "sr-ae-x-WBqB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "cifair10_save_folder = r\"data/cifair10\"\n",
        "output_dir = r\"/content/gdrive/MyDrive/ImageFitting/run_output\""
      ],
      "metadata": {
        "id": "nXI8tQ5jlA26"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(f\"device: {device}\")\n",
        "\n",
        "hls, start_hl_index = [2, 3, 4, 5, 6, 7, 8], 0\n",
        "hfs, start_hf_index = [32, 64, 128, 256, 512, 1024], 0\n",
        "assert start_hl_index < len(hls)\n",
        "assert start_hf_index < len(hfs)\n",
        "start_run_index = 0\n",
        "for hl_index, hidden_layers in enumerate(hls):\n",
        "    for hf_index, hidden_features in enumerate(hfs):\n",
        "        experiment_name = f\"{hidden_layers}hls_{hidden_features}hfs\"\n",
        "        current_exp_save_folder = f\"{output_dir}/{experiment_name}\"\n",
        "        prev_exps = len(hfs) * hl_index + hf_index\n",
        "        num_exps_to_skip = len(hfs) * start_hl_index + start_hf_index\n",
        "        if prev_exps < num_exps_to_skip:\n",
        "            print(f\"Skipping {experiment_name}...\")\n",
        "            continue\n",
        "        start_time = time.time()\n",
        "        params = Params(hidden_layers=hidden_layers, hidden_features=hidden_features, learning_rate=1e-4,\n",
        "                        batch_size=1, fit_epochs=2000, num_triples=20)\n",
        "        params.save_json(save_path=current_exp_save_folder)\n",
        "\n",
        "        transform = transforms.Compose(\n",
        "            [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
        "        triples_dataset = TriplesCiFAIR10(determinism_seed=SEED_FOR_DATASET, root=cifair10_save_folder, train=True,\n",
        "                                          download=True, transform=transform, num_triples=params.num_triples)\n",
        "\n",
        "        triples_dataloader = DataLoader(triples_dataset, shuffle=True, batch_size=params.batch_size,\n",
        "                                        pin_memory=True,\n",
        "                                        num_workers=0)\n",
        "\n",
        "        # Three images in `images` from dataloader\n",
        "        # 0 -> random image of class C0, A.\n",
        "        # 1 -> random image of class C1, B where C1 != C0 and A != B. Last conditional is redundant but still :)\n",
        "        # 2 -> another random image of class C0, C where C != A.\n",
        "\n",
        "        # coords is always the same, so it seems memory-inefficient to store it in every batch.\n",
        "        coords = get_mgrid(sidelen=32, dim=2)\n",
        "        is_first_run_in_exp = True\n",
        "        run_seeds = random.sample(range(1, sys.maxsize), params.num_triples)\n",
        "        print(f\"Starting {experiment_name}...\")\n",
        "        for run_index, (images, labels) in enumerate(triples_dataloader):\n",
        "            if run_index < start_run_index:\n",
        "                print(f\"Skipping {experiment_name}...{run_index}\")\n",
        "                continue\n",
        "            else:\n",
        "                start_run_index = -1\n",
        "            torch.manual_seed(\n",
        "                run_seeds[run_index])  # The random init for each triple is generated using a unique rand seed.\n",
        "            labels = labels[0]\n",
        "            # Save time and avg it? maybe leave this to tqdm? then run a lot of these (>100) on colab asap!\n",
        "            siren_mlp = Siren(in_features=params.in_features, hidden_features=params.hidden_features,\n",
        "                              hidden_layers=params.hidden_layers, out_features=params.out_features)\n",
        "            if is_first_run_in_exp:\n",
        "                # We don't change is_first_run to True yet since we need it for save_run_info.\n",
        "                print(siren_mlp)\n",
        "            siren_mlp = siren_mlp.to(device)\n",
        "            initial_state_dict = siren_mlp.state_dict()\n",
        "            for fitting_mode in list(FittingMode):\n",
        "                if fitting_mode == FittingMode.FIT_FROM_RANDOM_INIT:\n",
        "                    pass  # no op\n",
        "                elif fitting_mode == FittingMode.FIT_FROM_MODEL_FIT_TO_IMAGE_FROM_SAME_CLASS:\n",
        "                    # fit to C\n",
        "                    siren_mlp.load_state_dict(initial_state_dict)\n",
        "                    fit_one_image(model=siren_mlp, coords=coords, image_to_fit=images[2], params=params,\n",
        "                                  device=device,\n",
        "                                  plot_losses=False, current_fitting_mode=fitting_mode, pretrain=True)\n",
        "                else:\n",
        "                    # fit to B\n",
        "                    siren_mlp.load_state_dict(initial_state_dict)\n",
        "                    fit_one_image(model=siren_mlp, coords=coords, image_to_fit=images[1], params=params,\n",
        "                                  device=device,\n",
        "                                  plot_losses=False, current_fitting_mode=fitting_mode, pretrain=True)\n",
        "\n",
        "                tqdm_desc = f\"{experiment_name}: triple {run_index}/{params.num_triples}, labels {labels.tolist()}, \" \\\n",
        "                            f\"fit mode {fitting_mode.name}\"\n",
        "                # fit to A\n",
        "                fitting_metrics = fit_one_image(model=siren_mlp, coords=coords, image_to_fit=images[0],\n",
        "                                                params=params,\n",
        "                                                device=device, tqdm_description=tqdm_desc, plot_losses=False,\n",
        "                                                current_fitting_mode=fitting_mode, pretrain=False,\n",
        "                                                calculate_psnr=True)\n",
        "                run_info = {\"labels\": labels.tolist(), \"fitting_mode\": fitting_mode.value,\n",
        "                            \"metrics\": fitting_metrics}\n",
        "                save_run_info(run_info, current_exp_save_folder, clear_existing_file=is_first_run_in_exp)\n",
        "\n",
        "                if is_first_run_in_exp:\n",
        "                    is_first_run_in_exp = False\n",
        "\n",
        "        exp_duration = time.strftime('%H:%M:%S', time.gmtime(time.time() - start_time))\n",
        "        print(f\"DONE: {experiment_name}. hl_index: {hl_index}, hf_index: {hf_index}. Duration: {exp_duration}\")\n"
      ],
      "metadata": {
        "id": "EPxLzbpBRq7P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title So I know when to disconnect the runtime lol\n",
        "from google.colab import output\n",
        "output.eval_js('new Audio(\"https://upload.wikimedia.org/wikipedia/commons/0/05/Beep-09.ogg\").play()')"
      ],
      "metadata": {
        "id": "YjC7k_YRfDK8"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}