Flow Matching Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e2k5H1y4aZys"
   },
   "source": [
    "### Connecting to Gdrive for imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4729,
     "status": "ok",
     "timestamp": 1701944233264,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "KmF1M7d4Yr4M",
    "outputId": "8f61d300-29cb-4a27-f885-edc0552089bc"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive/')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zmAGOeWLaTwf"
   },
   "source": [
    "### Importing necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kdTDWWerZEIO"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/content/drive/MyDrive/Colab_Notebooks/CS_682/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 20156,
     "status": "ok",
     "timestamp": 1701944253415,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "uHJtKZRtugPP",
    "outputId": "07062aad-ec0e-4ca9-88ec-eabbe8e71f88"
   },
   "outputs": [],
   "source": [
    "!pip install torchdiffeq\n",
    "!pip install torchmetrics\n",
    "!pip install torchviz\n",
    "!pip install torch-fidelity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 5399,
     "status": "ok",
     "timestamp": 1701944258796,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "60hiPDXUV4WE",
    "outputId": "03b720a4-a998-4754-a529-414105dcdbca"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset\n",
    "from torchmetrics.functional.multimodal import clip_score\n",
    "from torchmetrics.image.fid import FrechetInceptionDistance\n",
    "import pickle\n",
    "from functools import partial\n",
    "import torchviz\n",
    "from torchviz import make_dot\n",
    "import PIL\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import tqdm\n",
    "from torchdiffeq import odeint\n",
    "import os\n",
    "\n",
    "from caption_generation import CIFAR10WithCaptions\n",
    "from unet_attn import UNet\n",
    "from text_encoding import reshape_text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lv9JsUZNfPfN"
   },
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4jsTtwIyfRRo"
   },
   "outputs": [],
   "source": [
    "def convert(loss_vals):\n",
    "  new_loss= []\n",
    "  for i in loss_vals:\n",
    "    new_loss.append(i)\n",
    "  return new_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lm91DGfSbtzd"
   },
   "source": [
    "### Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gcISdkGUcMqS"
   },
   "outputs": [],
   "source": [
    "def uncond_dataset():\n",
    "  transform = transforms.Compose(\n",
    "    [transforms.ToTensor()])\n",
    "\n",
    "  train_dataset = torchvision.datasets.CIFAR10(root='./drive/MyDrive/CS 682/CS682 Project/uncond_dataset/train/', train=True,\n",
    "                                          download=True, transform=transform)\n",
    "  test_dataset = torchvision.datasets.CIFAR10(root='./drive/MyDrive/CS 682/CS682 Project/uncond_dataset/train/', train=False,\n",
    "                                          download=True, transform=transform)\n",
    "\n",
    "  class Custom_CIFAR_train(torch.utils.data.Dataset):\n",
    "    def __init__(self, train_dataset):\n",
    "      self.target_imgs = train_dataset\n",
    "    def __getitem__(self, idx):\n",
    "      return self.target_imgs[idx][0]\n",
    "    def __len__(self):\n",
    "      return len(train_dataset)\n",
    "\n",
    "  flow_train_dataset = Custom_CIFAR_train(train_dataset)\n",
    "  flow_test_dataset = Custom_CIFAR_train(test_dataset)\n",
    "\n",
    "  return flow_train_dataset, flow_test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C4KaXOe5cpuY"
   },
   "outputs": [],
   "source": [
    "def cond_dataset():\n",
    "  with open('/content/drive/MyDrive/Colab_Notebooks/CS_682/Data/CAPTIONED_CIFAR_TRAIN.pkl', 'rb') as file:\n",
    "    flow_train_dataset = pickle.load(file)\n",
    "  with open('/content/drive/MyDrive/Colab_Notebooks/CS_682/Data/CAPTIONED_CIFAR_TEST.pkl', 'rb') as file:\n",
    "    flow_test_dataset = pickle.load(file)\n",
    "  return flow_train_dataset, flow_test_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tnABxZ0nbvf0"
   },
   "outputs": [],
   "source": [
    "def load_dataset(conditional_gen=False):\n",
    "  if conditional_gen:\n",
    "    return cond_dataset()\n",
    "  else:\n",
    "    return uncond_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zOeAspzUc6fe"
   },
   "source": [
    "### Sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "on4mXxoTdBN9"
   },
   "outputs": [],
   "source": [
    "def sample_from_dataset(dataset, conditional_dataset=False):\n",
    "  idx = torch.randint(0, len(dataset))\n",
    "  if conditional_dataset:\n",
    "    c_img, label, caption = dataset[idx]\n",
    "  else:\n",
    "    g_img, c_img = dataset[idx]\n",
    "  plt.imshow(g_img.permute(1,2,0))\n",
    "  plt.imshow(c_img.permute(1,2,0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BGvGinveepan"
   },
   "source": [
    "### Train-Val Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OBftiTxSeojR"
   },
   "outputs": [],
   "source": [
    "def split_train_dataset(flow_train_dataset, train_frac):\n",
    "  flow_train_dataset, flow_val_dataset = torch.utils.data.random_split(flow_train_dataset, [train_frac, 1-train_frac])\n",
    "  return flow_train_dataset, flow_val_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W0xZra5Ygkg6"
   },
   "source": [
    "### Create DataLoaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4HQ-pqcMgnuB"
   },
   "outputs": [],
   "source": [
    "def gen_loaders(dataset, batch_size):\n",
    "  return torch.utils.data.DataLoader(dataset, batch_size=batch_size, pin_memory=True,num_workers = 6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jkZcS_ZcfXVH"
   },
   "source": [
    "### Model Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RmZhJ03ufU0M"
   },
   "outputs": [],
   "source": [
    "def model_init(convnet = True, conditional_gen=False):\n",
    "  if conditional_gen:\n",
    "    return UNet(conditional_gen = conditional_gen)\n",
    "  else:\n",
    "    return UNet(conditional_gen = conditional_gen)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hasADKm_ChOA"
   },
   "source": [
    "### Masking Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6uzMsyiGCi-O"
   },
   "outputs": [],
   "source": [
    "def masking_tokens(tokens, pad_token=0):\n",
    "  p = torch.rand(1)\n",
    "  mask = torch.zeros(tokens.size())\n",
    "  if p<0.1:\n",
    "    tokens = tokens * mask\n",
    "  return tokens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F4FzziRkfzvS"
   },
   "source": [
    "### Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Fy1kmPH8fywc"
   },
   "outputs": [],
   "source": [
    "def loss(vf_flow, x_1, t, reshape_text, tokens = None, conditional_gen = False):\n",
    "\n",
    "  x_0 = torch.rand(x_1.shape).to(\"cuda\")\n",
    "\n",
    "  xt = t[:, None, None, None]*x_1 + (1-t[:,None, None, None])*x_0\n",
    "  xt = xt.cuda()\n",
    "\n",
    "  true_flow = x_1 - x_0\n",
    "  if conditional_gen:\n",
    "    tokens = reshape_text(tokens)\n",
    "    tokens = masking_tokens(tokens).to('cuda')\n",
    "  if conditional_gen:\n",
    "     predicted_flow, _ = vf_flow(t, (xt, tokens))\n",
    "  else:\n",
    "    predicted_flow  = vf_flow(t, xt)\n",
    "  flow_objective_loss = torch.sum((predicted_flow - true_flow)**2, axis=(1,2,3))\n",
    "  avg_obj_loss = torch.mean(flow_objective_loss)\n",
    "  return avg_obj_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cGuNC-QMgCRX"
   },
   "source": [
    "### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eRK5bffpgIJI"
   },
   "outputs": [],
   "source": [
    "def train_one_epoch(min_delta, patience, epochs,\n",
    "                    flow_train_dataset,\n",
    "                    lr, batch_size, reshape_text,\n",
    "                    convnet = True, conditional_gen=False,\n",
    "                    epoch_print = 1, infer_num = 0, epoch_save = 250, final_infer_num = 3):\n",
    "\n",
    "    loss_vals = []\n",
    "    infer_imgs = []\n",
    "    num_iter = -1\n",
    "    flow_train_loader = gen_loaders(flow_train_dataset, batch_size)\n",
    "    vf_flow = model_init(convnet, conditional_gen)\n",
    "    vf_flow.to('cuda')\n",
    "    optimizer = torch.optim.Adam(vf_flow.parameters(), lr = lr)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "      if (epoch % epoch_save == 0 and epoch != 0):\n",
    "        plot_loss(loss_vals)\n",
    "        torch.save(vf_flow, f'/content/drive/MyDrive/Colab_Notebooks/CS_682/Models/Conditional/Intermediate/conditional_unet_honey_ham_epoch_{epoch}.pth')\n",
    "      if (epoch % epoch_print == 0 and epoch!=0):\n",
    "        print(f'Epoch: {epoch} \\nLoss: {loss_vals[-1]}')\n",
    "        print('-------------------')\n",
    "      for i, data in enumerate(flow_train_loader):\n",
    "        num_iter+=1\n",
    "        stop = True\n",
    "        t = torch.rand(len(data[0])).to('cuda')\n",
    "\n",
    "        if conditional_gen:\n",
    "          x_1, labels, captions, x_0 = data\n",
    "          loss_val = loss(vf_flow, x_1.to('cuda'), t, reshape_text.to('cuda'), captions, conditional_gen=True)\n",
    "        else:\n",
    "          x_0, x_1 = data\n",
    "          loss_val = loss(vf_flow, x_1.to('cuda'), t, reshape_text.to('cuda'))\n",
    "\n",
    "        loss_vals.append(loss_val.item())\n",
    "\n",
    "        # Zero your gradients for every batch!\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        loss_val.backward()\n",
    "\n",
    "        # Adjust learning weights\n",
    "        optimizer.step()\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    final_images = []\n",
    "    torch.save(vf_flow, '/content/drive/MyDrive/Colab_Notebooks/CS_682/Models/Conditional/Final/conditional_unet_honey_ham_final.pth')\n",
    "    for infers in range(final_infer_num):\n",
    "          idx = torch.randint(0, len(flow_train_dataset), (1,)).item()\n",
    "          if conditional_gen:\n",
    "            img = inference(vf_flow, caption = flow_train_dataset[idx][2], reshape_text=reshape_text, conditional_gen=True).to('cpu').permute(1,2,0)\n",
    "          else:\n",
    "            img = inference(vf_flow, conditional_gen=False).to('cpu').permute(1,2,0)\n",
    "          plt.imshow(img)\n",
    "          plt.show()\n",
    "          final_images.append(img)\n",
    "    plot_loss(loss_vals)\n",
    "    return loss_vals, infer_imgs, final_images, vf_flow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wvbSkXJqh_xp"
   },
   "source": [
    "### Model Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_DLP5rqjiBz8"
   },
   "outputs": [],
   "source": [
    "def model_viz_print(model):\n",
    "  print(model)\n",
    "\n",
    "def model_viz_graph(model, conditional_gen=False):\n",
    "  if conditional_gen:\n",
    "    x = torch.rand(2,4,32,32)\n",
    "    t = torch.FloatTensor([0.0,1.0])\n",
    "  else:\n",
    "    x = torch.rand(2,3,32,32)\n",
    "    t = torch.FloatTensor([0.0,1.0])\n",
    "  make_dot(model(t,x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JNh9HQw_jcY_"
   },
   "source": [
    "### ODE Solver for Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "27v39aZ0jfRP"
   },
   "outputs": [],
   "source": [
    "def inference(model, caption=None, reshape_text=None, conditional_gen=False):\n",
    "  x_0 = torch.rand(1,3,32,32).to('cuda')\n",
    "  if conditional_gen:\n",
    "    tokens = reshape_text([caption]).to('cuda')\n",
    "\n",
    " # t = torch.linspace(0.0,1.0,10) # To observe the change in the image\n",
    "  t = torch.tensor([0.0, 1.0]).to('cuda')\n",
    "\n",
    "  with torch.no_grad():\n",
    "      if conditional_gen:\n",
    "        x_1, _ = odeint(model, (x_0, tokens), t, method='dopri5', atol=1e-5, rtol=1e-5)\n",
    "      else:\n",
    "        x_1 = odeint(model, x_0, t, method='dopri5', atol=1e-5, rtol=1e-5)\n",
    "  return x_1[-1,0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XQERkwbLk1Qi"
   },
   "source": [
    "### Loss Curve Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kcO74jNYk4CM"
   },
   "outputs": [],
   "source": [
    "def plot_loss(loss_vals):\n",
    "  plt.xlabel('#iterations')\n",
    "  plt.ylabel('Loss')\n",
    "  plt.plot(list(range(len(loss_vals))), convert(loss_vals))\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oaolYI0YM2fD"
   },
   "source": [
    "## Unconditional Flow Matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ntFCn9oCusXU"
   },
   "outputs": [],
   "source": [
    "flow_train_dataset, flow_test_dataset = load_dataset(False)\n",
    "flow_train_dataset, flow_val_dataset = split_train_dataset(flow_train_dataset, 0.98)\n",
    "ds_size = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AKElLufR2uUK"
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "import math\n",
    "imgs = []\n",
    "\n",
    "for img in range(ds_size):\n",
    "  imgs.append(flow_train_dataset[img][1].permute(1,2,0))\n",
    "\n",
    "  ncols = 10\n",
    "  nrows = math.ceil(ds_size / ncols)\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(10., 4.))\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(nrows, ncols),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, imgs):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_delta = 0\n",
    "patience = 10\n",
    "epochs = 2000\n",
    "lr = 1e-4\n",
    "batch_size = 128\n",
    "epoch_print = 1\n",
    "epoch_save = 6\n",
    "infer_num = 0\n",
    "final_infer_num = 10\n",
    "conditional_gen = False\n",
    "convnet = True\n",
    "\n",
    "\n",
    "print(f'train for <{epochs}> epochs')\n",
    "print(f'learning rate is <{lr}>')\n",
    "print(f'batch size is <{batch_size}>')\n",
    "print(f'conditional generation is set to <{conditional_gen}>')\n",
    "print(f'print loss every <{epoch_print}> epochs')\n",
    "print(f'generate <{infer_num}> images every <{epoch_print}> epochs')\n",
    "print(f'save intermediate models every <{epoch_save}> epochs')\n",
    "print(f'generate <{final_infer_num}> images after training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2vdN-IxUPFYM"
   },
   "outputs": [],
   "source": [
    "loss_vals, infer_imgs,final_infers, model = train_one_epoch(min_delta, patience, epochs,\n",
    "                                                            flow_train_dataset,\n",
    "                                                            lr, batch_size, reshape_text,\n",
    "                                                            convnet=convnet, conditional_gen=conditional_gen,\n",
    "                                                            epoch_print=epoch_print, epoch_save=epoch_save, infer_num = infer_num,\n",
    "                                                            final_infer_num=final_infer_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CH_DAIo_J-R4"
   },
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import set_loglevel\n",
    "import math\n",
    "\n",
    "plot_loss(loss_vals)\n",
    "\n",
    "total_infers = len(infer_imgs)\n",
    "grid_cols = 10\n",
    "grid_rows = math.ceil(total_infers / grid_cols)\n",
    "\n",
    "plt.set_loglevel(\"critical\")\n",
    "fig = plt.figure(figsize=(grid_cols, grid_rows))\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(grid_rows, grid_cols),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, infer_imgs):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "final_infer_num = 3\n",
    "grid_rows = math.ceil(final_infer_num / grid_cols)\n",
    "fig = plt.figure(figsize=(grid_cols, grid_rows))\n",
    "\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(grid_rows, grid_cols),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, final_infers):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()\n",
    "plt.set_loglevel(\"warning\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jGw0zJSvMkzd"
   },
   "source": [
    "## Conditional Flow Matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YGh1JFYD1iwg"
   },
   "outputs": [],
   "source": [
    "flow_train_dataset, flow_test_dataset = load_dataset(True)\n",
    "flow_train_dataset, flow_val_dataset = split_train_dataset(flow_train_dataset, 0.98)\n",
    "ds_size = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 290
    },
    "executionInfo": {
     "elapsed": 5630,
     "status": "ok",
     "timestamp": 1701944265518,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "mbY0XiuXmC3m",
    "outputId": "4566ae96-e32d-41bd-8b7e-0f3f5b59a0d8"
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "imgs = []\n",
    "\n",
    "for img in range(ds_size):\n",
    "  imgs.append(flow_train_dataset[img][0].permute(1,2,0))\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(10., 4.))\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(3, 10),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, imgs):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 25,
     "status": "ok",
     "timestamp": 1701944265519,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "DTpnml3TnvuC",
    "outputId": "16349396-b087-420f-eb54-75d3a72ce92c"
   },
   "outputs": [],
   "source": [
    "for i in range(ds_size):\n",
    "  print(flow_train_dataset[i][2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 16,
     "status": "ok",
     "timestamp": 1701944265519,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "u7kWxBY34z2w",
    "outputId": "2676c046-4a76-4328-d7e0-d96fb037c6bd"
   },
   "outputs": [],
   "source": [
    "min_delta = 0\n",
    "patience = 10\n",
    "epochs = 2500\n",
    "lr = 3e-4\n",
    "batch_size = 128\n",
    "epoch_print = 1\n",
    "epoch_save = 6\n",
    "infer_num = 0\n",
    "final_infer_num = 10\n",
    "conditional_gen = True\n",
    "convnet = True\n",
    "\n",
    "\n",
    "print(f'train for <{epochs}> epochs')\n",
    "print(f'learning rate is <{lr}>')\n",
    "print(f'batch size is <{batch_size}>')\n",
    "print(f'conditional generation is set to <{conditional_gen}>')\n",
    "print(f'print loss every <{epoch_print}> epochs')\n",
    "print(f'generate <{infer_num}> images every <{epoch_print}> epochs')\n",
    "print(f'save intermediate models every <{epoch_save}> epochs')\n",
    "print(f'generate <{final_infer_num}> images after training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1088214,
     "status": "error",
     "timestamp": 1701972662281,
     "user": {
      "displayName": "Ashvath Balgovind",
      "userId": "14860850667161268200"
     },
     "user_tz": 300
    },
    "id": "Fp4G1Pmpe76x",
    "outputId": "dd25ab3d-c83e-4897-d902-32c68c200869"
   },
   "outputs": [],
   "source": [
    "loss_vals, infer_imgs,final_infers, model = train_one_epoch(min_delta, patience, epochs,\n",
    "                                                            flow_train_dataset,\n",
    "                                                            lr, batch_size, reshape_text,\n",
    "                                                            convnet=convnet, conditional_gen=conditional_gen,\n",
    "                                                            epoch_print=epoch_print, epoch_save=epoch_save, infer_num = infer_num,\n",
    "                                                            final_infer_num=final_infer_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import set_loglevel\n",
    "import math\n",
    "\n",
    "plot_loss(loss_vals)\n",
    "\n",
    "total_infers = len(infer_imgs)\n",
    "grid_cols = 10\n",
    "grid_rows = math.ceil(total_infers / grid_cols)\n",
    "\n",
    "plt.set_loglevel(\"critical\")\n",
    "fig = plt.figure(figsize=(grid_cols, grid_rows))\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(grid_rows, grid_cols),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, infer_imgs):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()\n",
    "\n",
    "\n",
    "final_infer_num = 3\n",
    "grid_rows = math.ceil(final_infer_num / grid_cols)\n",
    "fig = plt.figure(figsize=(grid_cols, grid_rows))\n",
    "\n",
    "grid = ImageGrid(fig, 111,  # similar to subplot(111)\n",
    "                 nrows_ncols=(grid_rows, grid_cols),  # creates 2x2 grid of axes\n",
    "                 axes_pad=0.1,  # pad between axes in inch.\n",
    "                 )\n",
    "\n",
    "for ax, im in zip(grid, final_infers):\n",
    "    # Iterating over the grid returns the Axes.\n",
    "    ax.imshow(im)\n",
    "\n",
    "plt.show()\n",
    "plt.set_loglevel(\"warning\")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "wvbSkXJqh_xp",
    "JNh9HQw_jcY_",
    "XQERkwbLk1Qi",
    "hSjVWzdWlMzP",
    "oaolYI0YM2fD"
   ],
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}