12_gan/denoising_diffusion_probabilistic_model.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/machine-perception-robotics-group/MPRGDeepLearningLectureNotebook/blob/master/12_gan/denoising_diffusion_probabilistic_model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJFbw-pAbOiE",
        "tags": []
      },
      "source": [
        "# Denoising Diffusion Probabilistic Model (DDPM)\n",
        "\n",
        "---\n",
        "## 目的\n",
        "Pytorchを用いてDenoising Diffusion Probabilistic Model (DDPM) を構築し，画像の生成を行う．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_z832d5MXUMl",
        "tags": []
      },
      "source": [
        "## モジュールのインポート\n",
        "\n",
        "はじめに必要となるモジュールをインポートします．\n",
        "\n",
        "### GPUの確認\n",
        "GPUを使用した計算が可能かどうかを確認します．\n",
        "\n",
        "`GPU availability: True`と表示されれば，GPUを使用した計算を行うことが可能です．\n",
        "Falseとなっている場合は，上部のメニューバーの「ランタイム」→「ランタイムのタイプを変更」からハードウェアアクセラレータをGPUにしてください．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u8rszdnIbOiF",
        "tags": []
      },
      "outputs": [],
      "source": [
        "import math\n",
        "import numpy as np\n",
        "from inspect import isfunction\n",
        "from functools import partial\n",
        "from time import time\n",
        "\n",
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "from IPython.display import HTML\n",
        "\n",
        "from tqdm.auto import tqdm\n",
        "from einops import rearrange\n",
        "from einops.layers.torch import Rearrange\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch import nn, einsum\n",
        "from torch.optim import Adam\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "# # GPUの確認\n",
        "use_cuda = torch.cuda.is_available()\n",
        "print('Use CUDA:', use_cuda)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pNGKB-mjbOiI",
        "tags": []
      },
      "source": [
        "## ネットワークの構築\n",
        "DDPMは医療用画像のセマンティックセグメンテーションを行うモデルとして提案されたU-Netをベースにネットワークを構築しています．\n",
        "\n",
        "また，DDPMのU-Netは以下の点でオリジナルのU-Netとは異なります．\n",
        "* Position Embeddingsの追加\n",
        "* Wide ResNetの採用\n",
        "* Attentionの追加\n",
        "* Group Normalizationの追加\n",
        "\n",
        "これらを考慮してDDPMのネットワークを構築しますが，本ノートブックでは実装の簡略化のためWide ResNetではなく，通常のResNetを使用します．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cEPn5DMA8qG1",
        "tags": []
      },
      "source": [
        "## Position Embeddings\n",
        "DDPMでは各時刻$t$のノイズを推定する時，ネットワークのパラメータは共通です．時刻$t$ごとにネットワークを構築するのではなく，どの時刻$t$かを表す情報をネットワークに与えることで各時刻$t$のノイズを推定することが可能となります．\n",
        "\n",
        "Position Embeddingsでは以下に示す式によって時刻の特徴量を求め，U-Net層の各Residual blockに追加されます．\n",
        "$$\n",
        "PE_{(pos, 2i)} = \\sin (\\frac{pos}{10000^{(2i/d)}})\n",
        "$$\n",
        "$$\n",
        "PE_{(pos, 2i+1)} = \\cos (\\frac{pos}{10000^{(2i/d)}})\n",
        "$$\n",
        "ここで，$pos$は時刻，$i$は時刻特徴量の次元のインデックス，$d$は時刻特徴量の次元数を表します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "0k0lLZAAbOiI",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class PositionEmbeddings(nn.Module):\n",
        "    def __init__(self, dim):\n",
        "        super().__init__()\n",
        "        self.dim = dim\n",
        "\n",
        "    def forward(self, time):\n",
        "        device = time.device\n",
        "        half_dim = self.dim // 2\n",
        "        embeddings = math.log(10000) / (half_dim - 1)\n",
        "        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)\n",
        "        embeddings = time[:, None] * embeddings[None, :]\n",
        "        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)\n",
        "        return embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RqgRx0aa8qG2",
        "tags": []
      },
      "source": [
        "## ResNet Block\n",
        "ResNetは，通常のネットワークのように，何かしらの処理ブロックによる変換$F(x)$を単純に次の層に渡していくのではなく，残差接続構造によりその処理ブロックへの入力$x$をショートカットし， $H(x) = F(x)+x$を次の層に渡すようにしています．残差接続構造により，誤差逆伝播時に勾配が消失しても，層をまたいで値を伝播することができます．このショートカットを含めた処理単位をResidual blockと呼びます．\n",
        "\n",
        "DDPMのResNet Blockでは，Position Embeddingsで求めた時刻特徴量を画像特徴量に追加や残差接続を行います．ここで効率的な学習のために活性化関数にはRectified Linear Unit (ReLU)関数ではなく，Sigmoid-weighted Linear Unit (SiLU) 関数を用います．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "R8HPN7598qG2",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# 補助関数\n",
        "def exists(x):\n",
        "    return x is not None\n",
        "\n",
        "# 補助関数\n",
        "def default(val, d):\n",
        "    if exists(val):\n",
        "        return val\n",
        "    return d() if isfunction(d) else d\n",
        "\n",
        "# 残差接続\n",
        "class Residual(nn.Module):\n",
        "    def __init__(self, fn):\n",
        "        super().__init__()\n",
        "        self.fn = fn\n",
        "\n",
        "    def forward(self, x, *args, **kwargs):\n",
        "        return self.fn(x, *args, **kwargs) + x\n",
        "\n",
        "# アップサンプル\n",
        "def Upsample(dim, dim_out=None):\n",
        "    return nn.Sequential(\n",
        "        nn.Upsample(scale_factor=2, mode=\"nearest\"),\n",
        "        nn.Conv2d(dim, default(dim_out, dim), 3, padding=1),\n",
        "    )\n",
        "\n",
        "# ダウンサンプル\n",
        "def Downsample(dim, dim_out=None):\n",
        "    return nn.Sequential(\n",
        "        Rearrange(\"b c (h p1) (w p2) -> b (c p1 p2) h w\", p1=2, p2=2),\n",
        "        nn.Conv2d(dim * 4, default(dim_out, dim), 1),\n",
        "    )\n",
        "\n",
        "# 畳み込みブロック\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, dim, dim_out, groups=8):\n",
        "        super().__init__()\n",
        "        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)\n",
        "        self.norm = nn.GroupNorm(groups, dim_out)\n",
        "        self.act = nn.SiLU()\n",
        "\n",
        "    def forward(self, x, scale_shift=None):\n",
        "        x = self.proj(x)\n",
        "        x = self.norm(x)\n",
        "        if exists(scale_shift):\n",
        "            scale, shift = scale_shift\n",
        "            x = x * (scale + 1) + shift\n",
        "\n",
        "        x = self.act(x)\n",
        "        return x\n",
        "\n",
        "# 残差接続ブロック\n",
        "class ResnetBlock(nn.Module):\n",
        "    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):\n",
        "        super().__init__()\n",
        "        self.mlp = (\n",
        "            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out * 2))\n",
        "            if exists(time_emb_dim)\n",
        "            else None\n",
        "        )\n",
        "        self.block1 = Block(dim, dim_out, groups=groups)\n",
        "        self.block2 = Block(dim_out, dim_out, groups=groups)\n",
        "        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()\n",
        "\n",
        "    def forward(self, x, time_emb=None):\n",
        "        scale_shift = None\n",
        "        if exists(self.mlp) and exists(time_emb):\n",
        "            time_emb = self.mlp(time_emb)\n",
        "            time_emb = rearrange(time_emb, \"b c -> b c 1 1\")\n",
        "            scale_shift = time_emb.chunk(2, dim=1)\n",
        "\n",
        "        h = self.block1(x, scale_shift=scale_shift)\n",
        "        h = self.block2(h)\n",
        "        return h + self.res_conv(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KYS0rK8U8qG3",
        "tags": []
      },
      "source": [
        "## Attention\n",
        "DDPMでは，Dot-Product AttentionとLinear Attentionの2種類のAttentionをU-Netに追加しています．\n",
        "\n",
        "#### Scaled Dot-Product Attention\n",
        "Dot-Product Attentioとは，ある要素が他の要素にどれだけ注目するべきかを計算する手法であり，以下に示す式によって計算されます．\n",
        "$$\n",
        "\\rm{Attention}(Q, K, V)= \\rm{softmax} ( \\frac{QK^{T}}{\\sqrt{d_k} } ) V\n",
        "$$\n",
        "ここで，$Q, K, V$はそれぞれQuery，Key，Valueです．また$d_k$はQueryの次元数を表します．DDPMでは1つのQuery，Key，Valueを持たせるのではなく，小さいQuery，Key，Valueに分割して，分割した特徴表現を計算するMulti-Head Attentionを使用します．これにより，モデルが異なる特徴表現の異なる情報についてAttention weightを計算することが可能となります．\n",
        "\n",
        "#### Linear Attention\n",
        "Linear Attentionとは，シーケンス長$n$に対する計算コストを$O(n^2)$から$O(n)$へと削減した手法であり，以下に示す式によって計算されます\n",
        "$$\n",
        "\\rm{LinearAttention} (Q, K, V) = \\frac{\\phi (Q) ( \\phi (K)^T V)}{\\phi (Q) \\phi (K)^T}\n",
        "$$\n",
        "$$\n",
        "\\phi (x) = \\rm{elu}(x) + 1\n",
        "$$\n",
        "ここで，$\\rm{elu}(\\cdot)$はExponential Linear Unit (ELU)関数を表します．Linear Attentionでは$K$と$V$の積を最初に計算し，次に$KV$と$Q$の積を計算します．これにより，内積計算が線形に抑えられ，計算コストはシーケンスの長さ$n$に対して$O(n)$となります．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "qRm4M-bp8qG3",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class Attention(nn.Module):\n",
        "    def __init__(self, dim, heads=4, dim_head=32):\n",
        "        super().__init__()\n",
        "        self.scale = dim_head**-0.5\n",
        "        self.heads = heads\n",
        "        hidden_dim = dim_head * heads\n",
        "        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)\n",
        "        self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        b, c, h, w = x.shape\n",
        "        qkv = self.to_qkv(x).chunk(3, dim=1)\n",
        "        q, k, v = map(lambda t: rearrange(t, \"b (h c) x y -> b h c (x y)\", h=self.heads), qkv)\n",
        "        q = q * self.scale\n",
        "\n",
        "        sim = einsum(\"b h d i, b h d j -> b h i j\", q, k)\n",
        "        sim = sim - sim.amax(dim=-1, keepdim=True).detach()\n",
        "        attn = sim.softmax(dim=-1)\n",
        "\n",
        "        out = einsum(\"b h i j, b h d j -> b h i d\", attn, v)\n",
        "        out = rearrange(out, \"b h (x y) d -> b (h d) x y\", x=h, y=w)\n",
        "        return self.to_out(out)\n",
        "\n",
        "class LinearAttention(nn.Module):\n",
        "    def __init__(self, dim, heads=4, dim_head=32):\n",
        "        super().__init__()\n",
        "        self.scale = dim_head**-0.5\n",
        "        self.heads = heads\n",
        "        hidden_dim = dim_head * heads\n",
        "        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)\n",
        "        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1),\n",
        "                                    nn.GroupNorm(1, dim))\n",
        "\n",
        "    def forward(self, x):\n",
        "        b, c, h, w = x.shape\n",
        "        qkv = self.to_qkv(x).chunk(3, dim=1)\n",
        "        q, k, v = map(lambda t: rearrange(t, \"b (h c) x y -> b h c (x y)\", h=self.heads), qkv)\n",
        "\n",
        "        q = q.softmax(dim=-2)\n",
        "        k = k.softmax(dim=-1)\n",
        "\n",
        "        q = q * self.scale\n",
        "        context = torch.einsum(\"b h d n, b h e n -> b h d e\", k, v)\n",
        "\n",
        "        out = torch.einsum(\"b h d e, b h d n -> b h e n\", context, q)\n",
        "        out = rearrange(out, \"b h c (x y) -> b (h c) x y\", h=self.heads, x=h, y=w)\n",
        "        return self.to_out(out)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b5boTNiA8qG3",
        "tags": []
      },
      "source": [
        "## Group Normalization\n",
        "Group Normalizationとは，チャネルを$G$個のグループに分割し，グループごとに１つの平均と分散を計算する正規化手法です．以下に正規化手法のイメージ図を示します．図からわかるように，Group NormalizationはLayer NormazliationとInstance Normalizationの中間のような正規化です．Layer Normazliationでは各チャネルの特徴を捉えられず，Instance Normalizationではチャネル間の依存関係を捉えられない場合でも，Group Normalizationではグループの特徴を捉えつつ，グループ内の依存関係もとらえることが可能です．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFhibkeRkMBw"
      },
      "source": [
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3916683/ccf1ff69-8efc-c696-6798-982e32a8f554.png\" width=100%>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "uJajXcb_8qG3",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class PreNorm(nn.Module):\n",
        "    def __init__(self, dim, fn):\n",
        "        super().__init__()\n",
        "        self.fn = fn\n",
        "        self.norm = nn.GroupNorm(1, dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.norm(x)\n",
        "        return self.fn(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mp-AW2GM8qG4",
        "tags": []
      },
      "source": [
        "## Unet\n",
        "DDPMのU-Netでは，ノイズが付与された画像のバッチとそれぞれのノイズレベル（時刻$t$）を入力として受け取り，画像に追加されたノイズを推定しています．そのため，ネットワークの流れは以下のようになります．\n",
        "* 最初に，ノイズが付与された画像に畳み込み処理と時刻$t$のPosition Embeddingsが計算されます．\n",
        "* 次に，2つのResNet Block → Attention → 畳み込みの順番で処理されるダウンサンプリングが複数回適用されます．\n",
        "* ネットワークの中央では，ResNet Block → Attention → ResNet Blockの順番で処理をします．\n",
        "* アップサンプリングでは，2つのResNet Block → Attention → 畳み込みの順番を複数回処理をします．\n",
        "* 最後に，ResNet Blockの後に畳み込み処理をします．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_-pHMUaxkMBy"
      },
      "source": [
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3916683/23f56197-6977-9503-08fa-a39beda8012d.png\" width=100%>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "TiWcUVEt8qG4",
        "tags": []
      },
      "outputs": [],
      "source": [
        "class Unet(nn.Module):\n",
        "    def __init__(self, dim, init_dim=None, out_dim=None, dim_mults=(1, 2, 4, 8), channels=3, resnet_block_groups=4):\n",
        "        super().__init__()\n",
        "        self.channels = channels\n",
        "        input_channels = channels\n",
        "\n",
        "        init_dim = default(init_dim, dim)\n",
        "        self.init_conv = nn.Conv2d(input_channels, init_dim, 1, padding=0)\n",
        "\n",
        "        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]\n",
        "        in_out = list(zip(dims[:-1], dims[1:]))\n",
        "\n",
        "        block_klass = partial(ResnetBlock, groups=resnet_block_groups)\n",
        "\n",
        "        time_dim = dim * 4\n",
        "\n",
        "        self.time_mlp = nn.Sequential(\n",
        "            PositionEmbeddings(dim),\n",
        "            nn.Linear(dim, time_dim),\n",
        "            nn.GELU(),\n",
        "            nn.Linear(time_dim, time_dim),\n",
        "        )\n",
        "\n",
        "        self.downs = nn.ModuleList([])\n",
        "        self.ups = nn.ModuleList([])\n",
        "        num_resolutions = len(in_out)\n",
        "        # ダウンサンプリング\n",
        "        for ind, (dim_in, dim_out) in enumerate(in_out):\n",
        "            is_last = ind >= (num_resolutions - 1)\n",
        "            self.downs.append(\n",
        "                nn.ModuleList(\n",
        "                    [\n",
        "                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),\n",
        "                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),\n",
        "                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),\n",
        "                        Downsample(dim_in, dim_out)\n",
        "                        if not is_last else nn.Conv2d(dim_in, dim_out, 3, padding=1),\n",
        "                    ]\n",
        "                )\n",
        "            )\n",
        "        # 中間ブロック\n",
        "        mid_dim = dims[-1]\n",
        "        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)\n",
        "        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))\n",
        "        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)\n",
        "        # アップサンプリング\n",
        "        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):\n",
        "            is_last = ind == (len(in_out) - 1)\n",
        "            self.ups.append(\n",
        "                nn.ModuleList(\n",
        "                    [\n",
        "                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),\n",
        "                        block_klass(dim_out + dim_in, dim_out, time_emb_dim=time_dim),\n",
        "                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),\n",
        "                        Upsample(dim_out, dim_in)\n",
        "                        if not is_last else nn.Conv2d(dim_out, dim_in, 3, padding=1),\n",
        "                    ]\n",
        "                )\n",
        "            )\n",
        "\n",
        "        self.out_dim = default(out_dim, channels)\n",
        "        self.final_res_block = block_klass(dim * 2, dim, time_emb_dim=time_dim)\n",
        "        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)\n",
        "\n",
        "    def forward(self, x, time):\n",
        "        x = self.init_conv(x)\n",
        "        r = x.clone()\n",
        "        t = self.time_mlp(time)\n",
        "        h = []\n",
        "        # ダウンサンプリング\n",
        "        for block1, block2, attn, downsample in self.downs:\n",
        "            x = block1(x, t)\n",
        "            h.append(x)\n",
        "            x = block2(x, t)\n",
        "            x = attn(x)\n",
        "            h.append(x)\n",
        "            x = downsample(x)\n",
        "        # 中間ブロック\n",
        "        x = self.mid_block1(x, t)\n",
        "        x = self.mid_attn(x)\n",
        "        x = self.mid_block2(x, t)\n",
        "        # アップサンプリング\n",
        "        for block1, block2, attn, upsample in self.ups:\n",
        "            x = torch.cat((x, h.pop()), dim=1)\n",
        "            x = block1(x, t)\n",
        "            x = torch.cat((x, h.pop()), dim=1)\n",
        "            x = block2(x, t)\n",
        "            x = attn(x)\n",
        "            x = upsample(x)\n",
        "\n",
        "        x = torch.cat((x, r), dim=1)\n",
        "        x = self.final_res_block(x, t)\n",
        "        return self.final_conv(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ghwiV388qG5",
        "tags": []
      },
      "source": [
        "## Forward Process\n",
        "Forward Processは，入力画像$\\mathbf{x}_0$に対してノイズを付与し最終的には完全なノイズ$\\mathbf{x}_T$へと変換する確率過程であり，以下に示す式のように正規分布に従うマルコフ過程で定義されます．\n",
        "$$\n",
        "q(\\mathbf{x}_{1:T} | \\mathbf{x}_0) = \\prod_{t=1}^T q(\\mathbf{x}_t | \\mathbf{x}_{t-1})\n",
        "$$\n",
        "$$\n",
        "q(\\mathbf{x}_t |  \\mathbf{x}_{t-1}) = N(\\mathbf{x}_t ; \\sqrt{1 - \\beta_t} \\mathbf{x}_{t-1} \\mathbf{I})\n",
        "$$\n",
        "ここで，$\\beta_t$は変化量を表すパラメータを表します．Forward Processでは学習を行わず，単純に連続的な微小変化によって解析可能な分布に変換することが目的です．ノイズ量を調整するスケジューラの値は$0.0001 \\sim 0.02$，ノイズを付与するステップは$300$回とします．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oWuxpviTkMB0"
      },
      "source": [
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3916683/44f8adec-7e88-557e-125d-3c5148616cc8.png\" width=100%>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "vDd-r-tQ8qG5",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# スケジューラ\n",
        "def linear_beta_schedule(timesteps):\n",
        "    beta_start = 0.0001\n",
        "    beta_end = 0.02\n",
        "    return torch.linspace(beta_start, beta_end, timesteps)\n",
        "\n",
        "# ステップ数\n",
        "timesteps = 300\n",
        "\n",
        "betas = linear_beta_schedule(timesteps=timesteps)\n",
        "alphas = 1. - betas\n",
        "alphas_cumprod = torch.cumprod(alphas, axis=0)\n",
        "alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)\n",
        "sqrt_recip_alphas = torch.sqrt(1.0 / alphas)\n",
        "\n",
        "sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)\n",
        "sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)\n",
        "posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)\n",
        "\n",
        "def extract(a, t, x_shape):\n",
        "    batch_size = t.shape[0]\n",
        "    out = a.gather(-1, t.cpu())\n",
        "    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)\n",
        "\n",
        "def q_sample(x_start, t, noise=None):\n",
        "    if noise is None:\n",
        "        noise = torch.randn_like(x_start)\n",
        "    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)\n",
        "    sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x_start.shape)\n",
        "    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pehiRsVK8qG6",
        "tags": []
      },
      "source": [
        "## Reverse Process\n",
        "変化量を表すパラメータ$\\beta_t$が十分に小さい連続変換（Forward Process）の場合，その逆変換（Reverse Process）は同じ関数系で表現することが可能であり，ガウスノイズの除去として考えることができます．そのため，Reverse Processは以下に示す式のように定義します．\n",
        "$$\n",
        "p_\\theta (\\mathbf{x}_{0:T}) = p_\\theta (\\mathbf{x}_T) \\prod_{t=1}^T p_\\theta (\\mathbf{x}_{t-1} | \\mathbf{x}_t)\n",
        "$$\n",
        "$$\n",
        "p_\\theta (\\mathbf{x}_{t-1} | \\mathbf{x}_t) = N( \\mathbf{x}_{t-1}; \\mathbf{\\mu}_\\theta (\\mathbf{x}_t, t), \\Sigma_\\theta(\\mathbf{x}_t, t))\n",
        "$$\n",
        "上記の式では，平均$\\mathbf{\\mu}_\\theta$と共分散$\\Sigma_\\theta(\\mathbf{x}_t, t)$をニューラルネットワークで学習することになっていますが，論文では共分散$\\Sigma_\\theta(\\mathbf{x}_t, t)$をあらかじめ固定し学習しません．そのため，本ノートブックでは共分散を固定して実装していますが，ステップ数$t$を減らした場合においては平均と共分散の両方を学習した方が良いです．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T7FymmUxkMB5"
      },
      "source": [
        "<img src=\"https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3916683/68d24842-f3bc-0905-9588-9b279365484a.png\" width=100%>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "lYUmw7IE8qG6",
        "tags": []
      },
      "outputs": [],
      "source": [
        "@torch.no_grad()\n",
        "def p_sample(model, x, t, t_index):\n",
        "    betas_t = extract(betas, t, x.shape)\n",
        "    sqrt_one_minus_alphas_cumprod_t = extract(sqrt_one_minus_alphas_cumprod, t, x.shape)\n",
        "    sqrt_recip_alphas_t = extract(sqrt_recip_alphas, t, x.shape)\n",
        "    model_mean = sqrt_recip_alphas_t * (x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t)\n",
        "\n",
        "    if t_index == 0:\n",
        "        return model_mean\n",
        "    else:\n",
        "        posterior_variance_t = extract(posterior_variance, t, x.shape)\n",
        "        noise = torch.randn_like(x)\n",
        "        return model_mean + torch.sqrt(posterior_variance_t) * noise\n",
        "\n",
        "@torch.no_grad()\n",
        "def p_sample_loop(model, shape, img_input=None):\n",
        "    device = next(model.parameters()).device\n",
        "    b = shape[0]\n",
        "    if img_input == None:\n",
        "        print('noise')\n",
        "        img = torch.randn(shape, device=device)\n",
        "    else:\n",
        "        print('image')\n",
        "        img = img_input\n",
        "    imgs = []\n",
        "\n",
        "    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=timesteps):\n",
        "        img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), i)\n",
        "        imgs.append(img.cpu().numpy())\n",
        "    return imgs\n",
        "\n",
        "@torch.no_grad()\n",
        "def sample(model, image_size, batch_size=16, channels=3, img_input=None):\n",
        "    return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size), img_input=img_input)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n-Vv7kv4xoNu",
        "tags": []
      },
      "source": [
        "## 損失関数\n",
        "DDPMでは，Variational Autoencoderと同様にデータ$\\mathbf{x}_0$の対数尤度の変分下限を最大化するように学習します．DDPMの損失関数は以下に示す式のようになります．\n",
        "$$\n",
        "\\mathcal{L}_{\\rm{simple}} = \\mathbb{E}_{t, \\mathbf{x}_0, \\epsilon} [\\parallel \\epsilon - \\epsilon_\\theta(\\sqrt{\\bar{a}_t} \\mathbf{x}_0 + \\sqrt{1 - \\bar{a}_t}\\epsilon, t) \\parallel^2]\n",
        "$$\n",
        "ここで，$\\mathbf{x}_t = \\sqrt{\\bar{a}_t} \\mathbf{x}_0 + \\sqrt{1 - \\bar{a}_t}\\epsilon$と表せるため，損失関数は以下のようになります．\n",
        "$$\n",
        "\\mathcal{L}_{\\rm{simple}} = \\mathbb{E}_{t, \\mathbf{x}_0, \\epsilon} [\\parallel \\epsilon - \\epsilon_\\theta(\\mathbf{x}_t, t) \\parallel^2]\n",
        "$$\n",
        "従ってDDPMでは，ノイズが付与された画像$\\mathbf{x}_t$と時刻$t$が入力され，付与されたノイズ$\\epsilon$を推定するニューラルネットワーク$\\epsilon_\\theta$を学習します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "mLDvF4_MxoNu",
        "tags": []
      },
      "outputs": [],
      "source": [
        "def p_losses(denoise_model, x_start, t, noise=None):\n",
        "    if noise is None:\n",
        "        noise = torch.randn_like(x_start)\n",
        "    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)\n",
        "    predicted_noise = denoise_model(x_noisy, t)\n",
        "    loss = F.smooth_l1_loss(noise, predicted_noise)\n",
        "    return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZSistJSbOiK",
        "tags": []
      },
      "source": [
        "## データセット，最適化関数などの設定\n",
        "データセットはMNISTを用いて学習をします．\n",
        "最適化関数にはAdam Optimizerを使用します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J4zatgwzxoNu",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# データセットの設定\n",
        "transform_train = transforms.Compose([transforms.ToTensor()])\n",
        "mnist_data = datasets.MNIST(root='./data', train=True, transform=transform_train, download=True)\n",
        "train_loader = DataLoader(dataset=mnist_data, batch_size=128, shuffle=True)\n",
        "\n",
        "mnist_testdata = datasets.MNIST(root='./data', train=False, transform=transform_train)\n",
        "test_loader =DataLoader(dataset=mnist_testdata, batch_size=1, shuffle=True)\n",
        "\n",
        "\n",
        "image_size = 28\n",
        "channels = 1\n",
        "epochs = 6\n",
        "\n",
        "# ネットワークモデル・最適化手法の設定\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "model = Unet(\n",
        "    dim=image_size,\n",
        "    channels=channels,\n",
        "    dim_mults=(1, 2, 4,)\n",
        ")\n",
        "model.to(device)\n",
        "optimizer = Adam(model.parameters(), lr=1e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D0f83-lrbOiP"
      },
      "source": [
        "## ネットワークの学習"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IeMNuhEsbOiP",
        "tags": []
      },
      "outputs": [],
      "source": [
        "model.train()\n",
        "\n",
        "start = time()\n",
        "for epoch in range(epochs):\n",
        "    sum_loss = 0.0\n",
        "    for x, _ in train_loader:\n",
        "        optimizer.zero_grad()\n",
        "        batch_size = x.shape[0]\n",
        "        x = x.to(device)\n",
        "        t = torch.randint(0, timesteps, (batch_size,), device=device).long()\n",
        "        loss = p_losses(model, x, t)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        sum_loss += loss.item()\n",
        "\n",
        "    print(\"epoch:{}, Loss:{}, elapsed time: {}\".format(epoch, sum_loss / len(train_loader), time() - start))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QiBN_Q6ybOiS"
      },
      "source": [
        "## 学習済みモデルを用いたノイズからの画像生成\n",
        "\n",
        "先ほど学習した重みパラメータを用いて，ノイズから画像の生成をします．mnistの画像サイズと同じランダムなノイズを作成し，その値をモデルに入力した結果を確認します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3fetGCOxxoNv",
        "tags": []
      },
      "outputs": [],
      "source": [
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    # ノイズからのサンプリング\n",
        "    samples = sample(model, image_size=image_size, batch_size=1, channels=channels)\n",
        "\n",
        "plt.imshow(samples[0].reshape(image_size, image_size, channels), cmap=\"gray\")\n",
        "plt.title(\"Input\")\n",
        "plt.show()\n",
        "plt.imshow(samples[-1].reshape(image_size, image_size, channels), cmap=\"gray\")\n",
        "plt.title(\"Output\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JRI6CgCXxoNv"
      },
      "source": [
        "## デノイジングの可視化\n",
        "\n",
        "先ほど生成した画像のデノイジング過程 (Reverse Process) を可視化して確認します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RFO5N1IhxoNv",
        "tags": []
      },
      "outputs": [],
      "source": [
        "fig = plt.figure()\n",
        "ims = []\n",
        "for i in range(timesteps):\n",
        "    if i % 50 == 0 or i == 299:\n",
        "        im = plt.imshow(samples[i].reshape(image_size, image_size, channels), cmap=\"gray\", animated=True)\n",
        "        ims.append([im])\n",
        "\n",
        "animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000)\n",
        "HTML(animate.to_jshtml())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dpawTBOBkMCB"
      },
      "source": [
        "# ノイズからの画像生成\n",
        "5つのノイズから画像を5枚生成し，どのような画像が生成されるか確認します．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIfD-e04xoNw"
      },
      "outputs": [],
      "source": [
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    samples = sample(model, image_size=image_size, batch_size=5, channels=channels)\n",
        "\n",
        "input_imgs = np.concatenate([samples[0][i].reshape(image_size, image_size) for i in range(5)], axis=1)\n",
        "output_imgs = np.concatenate([samples[-1][i].reshape(image_size, image_size) for i in range(5)], axis=1)\n",
        "\n",
        "plt.imshow(input_imgs, cmap=\"gray\")\n",
        "plt.title(\"Input\")\n",
        "plt.show()\n",
        "plt.imshow(output_imgs, cmap=\"gray\")\n",
        "plt.title(\"Output\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dUEfiAj3kMCD",
        "tags": []
      },
      "source": [
        "# 課題\n",
        "\n",
        "1. 拡散過程を適用していない画像に対してデノイジング過程を適用し，どのような画像が生成されるか確認してみましょう．"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XTN7N0FskMCD",
        "tags": []
      },
      "outputs": [],
      "source": [
        "# ノイズが付与されていない画像に対してデノイジング過程を適用するには以下のようにします．\n",
        "\n",
        "model.eval()\n",
        "with torch.no_grad():\n",
        "    x, y = next(iter(test_loader))\n",
        "    samples = sample(model, image_size=image_size, batch_size=1, channels=channels, img_input=x.cuda())\n",
        "\n",
        "plt.imshow(x.reshape(image_size, image_size, channels), cmap=\"gray\")\n",
        "plt.title(\"Input\")\n",
        "plt.show()\n",
        "plt.imshow(samples[-1].reshape(image_size, image_size, channels), cmap=\"gray\")\n",
        "plt.title(\"Output\")\n",
        "plt.show()\n",
        "\n",
        "fig = plt.figure()\n",
        "ims = []\n",
        "for i in range(timesteps):\n",
        "    if i % 50 == 0 or i == 299:\n",
        "        im = plt.imshow(samples[i].reshape(image_size, image_size, channels), cmap=\"gray\", animated=True)\n",
        "        ims.append([im])\n",
        "\n",
        "animate = animation.ArtistAnimation(fig, ims, interval=50, blit=True, repeat_delay=1000)\n",
        "HTML(animate.to_jshtml())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsJKYWVBkMCE"
      },
      "source": [
        "2. 拡散過程のステップ数を変更して学習した際にどのような傾向が現れるか確認してみましょう．\n",
        "\n",
        "ヒント：拡散過程のステップ数を変更したい場合，Forward Processのセルで定義されているtimestepsを変更する．"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3_G1PaFn05t"
      },
      "source": [
        "# 参考文献\n",
        "[1] Jonathan Ho, Ajay Jain and Pieter Abbeel, Denoising Diffusion Probabilistic Models, NeurIPS, 2020."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "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.8.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}