entry.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Denoising Diffusion Probabilistic Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from data import DiffSet\n",
    "import pytorch_lightning as pl\n",
    "from model import DiffusionModel\n",
    "from torch.utils.data import DataLoader\n",
    "import imageio\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training hyperparameters\n",
    "diffusion_steps = 1000\n",
    "dataset_choice = \"CIFAR\"\n",
    "max_epoch = 10\n",
    "batch_size = 128\n",
    "\n",
    "# Loading parameters\n",
    "load_model = False\n",
    "load_version_num = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load dataset and train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code for optionally loading model\n",
    "pass_version = None\n",
    "last_checkpoint = None\n",
    "\n",
    "if load_model:\n",
    "    pass_version = load_version_num\n",
    "    last_checkpoint = glob.glob(\n",
    "        f\"./lightning_logs/{dataset_choice}/version_{load_version_num}/checkpoints/*.ckpt\"\n",
    "    )[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create datasets and data loaders\n",
    "train_dataset = DiffSet(True, dataset_choice)\n",
    "val_dataset = DiffSet(False, dataset_choice)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=4, shuffle=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, num_workers=4, shuffle=True)\n",
    "\n",
    "# Create model and trainer\n",
    "if load_model:\n",
    "    model = DiffusionModel.load_from_checkpoint(last_checkpoint, in_size=train_dataset.size*train_dataset.size, t_range=diffusion_steps, img_depth=train_dataset.depth)\n",
    "else:\n",
    "    model = DiffusionModel(train_dataset.size*train_dataset.size, diffusion_steps, train_dataset.depth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load Trainer model\n",
    "tb_logger = pl.loggers.TensorBoardLogger(\n",
    "    \"lightning_logs/\",\n",
    "    name=dataset_choice,\n",
    "    version=pass_version,\n",
    ")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    max_epochs=max_epoch, \n",
    "    log_every_n_steps=10, \n",
    "    gpus=1, \n",
    "    auto_select_gpus=True,\n",
    "    resume_from_checkpoint=last_checkpoint, \n",
    "    logger=tb_logger\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "trainer.fit(model, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample from model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gif_shape = [3, 3]\n",
    "sample_batch_size = gif_shape[0] * gif_shape[1]\n",
    "n_hold_final = 10\n",
    "\n",
    "# Generate samples from denoising process\n",
    "gen_samples = []\n",
    "x = torch.randn((sample_batch_size, train_dataset.depth, train_dataset.size, train_dataset.size))\n",
    "sample_steps = torch.arange(model.t_range-1, 0, -1)\n",
    "for t in sample_steps:\n",
    "    x = model.denoise_sample(x, t)\n",
    "    if t % 50 == 0:\n",
    "        gen_samples.append(x)\n",
    "for _ in range(n_hold_final):\n",
    "    gen_samples.append(x)\n",
    "gen_samples = torch.stack(gen_samples, dim=0).moveaxis(2, 4).squeeze(-1)\n",
    "gen_samples = (gen_samples.clamp(-1, 1) + 1) / 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process samples and save as gif\n",
    "gen_samples = (gen_samples * 255).type(torch.uint8)\n",
    "gen_samples = gen_samples.reshape(-1, gif_shape[0], gif_shape[1], train_dataset.size, train_dataset.size, train_dataset.depth)\n",
    "\n",
    "def stack_samples(gen_samples, stack_dim):\n",
    "    gen_samples = list(torch.split(gen_samples, 1, dim=1))\n",
    "    for i in range(len(gen_samples)):\n",
    "        gen_samples[i] = gen_samples[i].squeeze(1)\n",
    "    return torch.cat(gen_samples, dim=stack_dim)\n",
    "\n",
    "gen_samples = stack_samples(gen_samples, 2)\n",
    "gen_samples = stack_samples(gen_samples, 2)\n",
    "\n",
    "imageio.mimsave(\n",
    "    f\"{trainer.logger.log_dir}/pred.gif\",\n",
    "    list(gen_samples),\n",
    "    fps=5,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "78f961cc1adb1d23b163d13487a7bbea047622fac31372ac86f89c77664d2955"
  },
  "kernelspec": {
   "display_name": "Python 3.8.7 64-bit ('3.8.7': pyenv)",
   "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.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}