clip_nca.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "### Install CLIP"
      ],
      "metadata": {
        "id": "V4vGc1nIYfgo"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y8UlAJoO6F5X"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/openai/CLIP.git"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Check for GPU"
      ],
      "metadata": {
        "id": "qLW-uLWdYiTx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi -L"
      ],
      "metadata": {
        "id": "PXvZkWnnHTEd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download Code"
      ],
      "metadata": {
        "id": "bWewKSh8Ym-P"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!rm -rf /content/project-omega/ *.py\n",
        "!git clone https://github.com/Mainakdeb/project-omega.git\n",
        "!cp /content/project-omega/*.py /content/"
      ],
      "metadata": {
        "id": "DbZQ433Z8_oD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports"
      ],
      "metadata": {
        "id": "Ugd2Xz-iYxTw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fj9Q--hlPEKi"
      },
      "outputs": [],
      "source": [
        "from language_model import clip_encode_images, clip_encode_text, get_clip_loss\n",
        "from nca import ca_model, to_rgb\n",
        "from video_utils import create_inference_video, show_video, create_inference_gif\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "import torchvision.models as models\n",
        "import torchvision.transforms.functional as F\n",
        "from torchvision import transforms\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "from tqdm import tqdm_notebook, tnrange\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.pyplot as plt\n",
        "plt.rcParams[\"figure.figsize\"] = (10, 10)\n",
        "import requests\n",
        "import io\n",
        "import os\n",
        "from IPython.display import clear_output\n",
        "\n",
        "os.environ['FFMPEG_BINARY'] = 'ffmpeg'\n",
        "from google.colab import files"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Set params"
      ],
      "metadata": {
        "id": "qvBcM67tY2cK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nOlGhxB_PIEf"
      },
      "outputs": [],
      "source": [
        "# for CA model\n",
        "total_channels = 12\n",
        "hidden_filters = 128\n",
        "\n",
        "# training\n",
        "text_prompt = \"Leopard Skin\"\n",
        "learning_rate = 2e-3\n",
        "batch_size = 4\n",
        "lr_decay_checkpoints = [300, 500, 850]\n",
        "decay_factor = 0.4\n",
        "train_iterations = 1000\n",
        "\n",
        "# training pool\n",
        "pool_size = 256\n",
        "image_size = 128\n",
        "\n",
        "# NCA model\n",
        "ca = ca_model(chn=total_channels, hidden_n=hidden_filters).to(device)\n",
        "\n",
        "# optimizer\n",
        "opt = torch.optim.Adam(ca.parameters(), learning_rate)\n",
        "\n",
        "# learning rate scheduler\n",
        "lr_sched = torch.optim.lr_scheduler.MultiStepLR(opt, lr_decay_checkpoints, decay_factor)\n",
        "\n",
        "# array to log loss\n",
        "loss_log = []\n",
        "\n",
        "#initiate pool\n",
        "pool = ca.seed(n=pool_size, sz=image_size).to(device)\n",
        "\n",
        "# experimental pool transforms\n",
        "transform = transforms.Compose([# transforms.RandomHorizontalFlip(p=0.5),\n",
        "                                #transforms.RandomVerticalFlip(p=0.5),\n",
        "                                # tansforms.RandomRotation(degrees=30),\n",
        "                                # A.GridDistortion(.9)\n",
        "                                ])"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Train"
      ],
      "metadata": {
        "id": "ZZYnXoLeY6FE"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dLf17ePJPLQP"
      },
      "outputs": [],
      "source": [
        "for i in range(train_iterations):\n",
        "  with torch.no_grad():\n",
        "\n",
        "    # get ixs to sample \n",
        "    batch_idx = np.random.choice(len(pool), batch_size, replace=False)\n",
        "\n",
        "    #get sample from pool\n",
        "    x = pool[batch_idx]\n",
        "\n",
        "  # periodically replace element in pool with noise tensor \n",
        "  if i%8 == 0:\n",
        "    x[:1] = ca.seed(1, image_size).to(device)\n",
        "\n",
        "  # step_n: number of times the batch is passed through the model -> random integer\n",
        "  step_n = np.random.randint(64, 96)\n",
        "\n",
        "  # 1 forward pass\n",
        "  x = torch.utils.checkpoint.checkpoint_sequential([ca]*step_n, 16, x)\n",
        "\n",
        "  # convert images to RGB for loss computation (consider only first 3 chaannels  of 12 channels)\n",
        "  imgs = to_rgb(x)\n",
        "\n",
        "  # compute overflow loss  \n",
        "  overflow_loss = (x-x.clamp(-1.0, 1.0)).abs().sum() \n",
        "\n",
        "  # 2 compute loss\n",
        "  loss = get_clip_loss(text_prompt, imgs) + overflow_loss\n",
        "\n",
        "  # 3 clean gradients\n",
        "  opt.zero_grad()\n",
        "\n",
        "  # 4 accumulate partial derivatives\n",
        "  loss.backward()\n",
        "\n",
        "  # normalize gradients \n",
        "  with torch.no_grad():\n",
        "    for p in ca.parameters():\n",
        "      p.grad /= (p.grad.norm()+1e-8)   \n",
        "\n",
        "  # 5 step in opposite direction of gradient\n",
        "  opt.step()\n",
        "  \n",
        "  # step learning rate scheduler\n",
        "  lr_sched.step()\n",
        "\n",
        "  # update pool (replace sampled batch)\n",
        "  pool[batch_idx] = x                \n",
        "    \n",
        "  #log loss for plot  \n",
        "  loss_log.append(loss.item())\n",
        "\n",
        "  # periodically visualize loss and current batch\n",
        "  if i%10==0:\n",
        "    clear_output(True)\n",
        "    imgs = to_rgb(x).permute([0, 3, 2, 1]).cpu()\n",
        "    f, axarr = plt.subplots(2,1)\n",
        "    img_grid = torchvision.utils.make_grid(imgs.permute(0, 3, 2, 1,), nrow=4).cpu().detach().permute(1,2,0).numpy()\n",
        "    axarr[0].imshow(img_grid)\n",
        "    axarr[0].set_title('batch')\n",
        "    axarr[0].axis('off')\n",
        "    axarr[1].plot(loss_log[-20:], alpha=0.8)\n",
        "    axarr[1].set_xlabel('step:'+str(len(loss_log))+' loss:'+str(loss.item())+' overflow loss: '+str(overflow_loss.item()))\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vid_path = create_inference_video(ca_model=ca,\n",
        "                                  size=128,\n",
        "                                  num_frames=300, \n",
        "                                  steps_per_frame=5, \n",
        "                                  filename=\"test_vid.mov\")\n",
        "show_video(vid_path)"
      ],
      "metadata": {
        "id": "hqLQ1Uz4yHA8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gif_path = create_inference_gif(ca_model=ca,\n",
        "                                  size=128,\n",
        "                                  num_frames=120, \n",
        "                                  steps_per_frame=10, \n",
        "                                  fps=20,\n",
        "                                  filename=\"test_gif.gif\")\n",
        "\n",
        "files.download(gif_path) "
      ],
      "metadata": {
        "id": "QvP_XCzP8ahT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Bf_Mtg4Mabj"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "clip_nca_mobile.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}