tutorials/02_conditional_flow-based_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "TFGENZOO_Tutorial (2)",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HDzO_ynar37g",
        "colab_type": "text"
      },
      "source": [
        "# TFGENZOO_Tutorial (2) How to add conditional Input into the invertible layer\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/MokkeMeguru/TFGENZOO/blob/master/tutorials/01_What_is_the_invertible_layer.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HgkAgC9psKWp",
        "colab_type": "text"
      },
      "source": [
        "Some novels of Flow-based Model requires Conditional Input to control some attribution e.g. image style."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NSUO6XYXtbTF",
        "colab_type": "text"
      },
      "source": [
        "Below image is the example from the paper,\n",
        "\"Guided Image Generation \n",
        "with Conditional Invertible Neural Networks (cINN)\" \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eR3uqoJouyIw",
        "colab_type": "text"
      },
      "source": [
        "![cinn.png](https://raw.githubusercontent.com/MokkeMeguru/TFGENZOO/master/tutorials/img/cinn.png)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ENLSjAeu6kl",
        "colab_type": "text"
      },
      "source": [
        "# Conditional Input into the AffineCoupling Layer\n",
        "\n",
        "Conditional Input usually concated in the AffineCoupling Layer."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GDbcR8ZCu5h-",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 785
        },
        "outputId": "b9429e5f-cc18-4f2b-dd4c-636aa5df87a5"
      },
      "source": [
        "# install TFGENZOO\n",
        "!pip install TFGENZOO==1.2.4.post7"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: TFGENZOO==1.2.4.post7 in /usr/local/lib/python3.6/dist-packages (1.2.4.post7)\n",
            "Requirement already satisfied: tensorflow>=2.1.0 in /usr/local/lib/python3.6/dist-packages (from TFGENZOO==1.2.4.post7) (2.2.0)\n",
            "Requirement already satisfied: tensorflow-probability>=0.9.0 in /usr/local/lib/python3.6/dist-packages (from TFGENZOO==1.2.4.post7) (0.10.0)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.6/dist-packages (from TFGENZOO==1.2.4.post7) (1.0.5)\n",
            "Requirement already satisfied: numpy<2.0,>=1.16.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.18.5)\n",
            "Requirement already satisfied: keras-preprocessing>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.1.2)\n",
            "Requirement already satisfied: protobuf>=3.8.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.10.0)\n",
            "Requirement already satisfied: grpcio>=1.8.6 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.30.0)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.12.0)\n",
            "Requirement already satisfied: tensorflow-estimator<2.3.0,>=2.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2.2.0)\n",
            "Requirement already satisfied: wheel>=0.26; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.34.2)\n",
            "Requirement already satisfied: wrapt>=1.11.1 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.12.1)\n",
            "Requirement already satisfied: tensorboard<2.3.0,>=2.2.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2.2.2)\n",
            "Requirement already satisfied: scipy==1.4.1; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.4.1)\n",
            "Requirement already satisfied: astunparse==1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.6.3)\n",
            "Requirement already satisfied: gast==0.3.3 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.3.3)\n",
            "Requirement already satisfied: h5py<2.11.0,>=2.10.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2.10.0)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.2.1)\n",
            "Requirement already satisfied: absl-py>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.9.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.1.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.8 in /usr/local/lib/python3.6/dist-packages (from tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.2.0)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.6/dist-packages (from tensorflow-probability>=0.9.0->TFGENZOO==1.2.4.post7) (4.4.2)\n",
            "Requirement already satisfied: cloudpickle>=1.2.2 in /usr/local/lib/python3.6/dist-packages (from tensorflow-probability>=0.9.0->TFGENZOO==1.2.4.post7) (1.3.0)\n",
            "Requirement already satisfied: python-dateutil>=2.6.1 in /usr/local/lib/python3.6/dist-packages (from pandas->TFGENZOO==1.2.4.post7) (2.8.1)\n",
            "Requirement already satisfied: pytz>=2017.2 in /usr/local/lib/python3.6/dist-packages (from pandas->TFGENZOO==1.2.4.post7) (2018.9)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from protobuf>=3.8.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (47.3.1)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.0.1)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.4.1)\n",
            "Requirement already satisfied: google-auth<2,>=1.6.3 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.17.2)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.6.0.post3)\n",
            "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2.23.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.6/dist-packages (from tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.2.2)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.6/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.3.0)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (4.1.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.2.8)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4; python_version >= \"3\" in /usr/local/lib/python3.6/dist-packages (from google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (4.6)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2020.6.20)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (2.9)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3,>=2.21.0->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.24.3)\n",
            "Requirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (1.6.1)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.6/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.1.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.6/dist-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (0.4.8)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata; python_version < \"3.8\"->markdown>=2.6.8->tensorboard<2.3.0,>=2.2.0->tensorflow>=2.1.0->TFGENZOO==1.2.4.post7) (3.1.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fFDRwGMwyJA4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "from TFGENZOO.flows.cond_affine_coupling import ConditionalAffineCoupling"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EwSHfbTb19JD",
        "colab_type": "text"
      },
      "source": [
        "## Conditional Affine Coupling forward formulation:\n",
        "\n",
        "```\n",
        "[x_1, x_2] = split(x)\n",
        "log_scale, shift <- NN([x_1, cond])\n",
        "z_1 = x_1\n",
        "z_2 = (x_2 + shift) * scale\n",
        "z = concat([z_1, z_2])\n",
        "\n",
        "log_det_jacobian = sum(log(scale))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B0Z4QGS62K7q",
        "colab_type": "text"
      },
      "source": [
        "## NN layer to deal with Conditional Input\n",
        "Here is the basic NN Layer proposed in cINN.\n",
        "\n",
        "It is same as TFGENZOO.layers.resnet.ShallowResNet."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4_QlQms-2ILq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from TFGENZOO.flows.utils import Conv2D, Conv2DZeros\n",
        "\n",
        "def ShallowResNet(\n",
        "    inputs: tf.keras.Input,\n",
        "    cond: tf.keras.Input = None,\n",
        "    width: int = 512,\n",
        "    out_scale: int = 2,\n",
        "):\n",
        "    \"\"\"ResNet of OpenAI's Glow\n",
        "    Args:\n",
        "        inputs (tf.Tensor): input tensor rank == 4\n",
        "        cond   (tf.Tensor): input tensor rank == 4 (optional)\n",
        "        width        (int): hidden width\n",
        "        out_scale    (int): output channel width scale\n",
        "    Returns:\n",
        "        model: tf.keras.Model\n",
        "    \"\"\"\n",
        "    _inputs = inputs if cond is None else tf.concat([inputs, cond], axis=-1)\n",
        "\n",
        "    conv1 = Conv2D(width=width)\n",
        "    conv2 = Conv2D(width=width)\n",
        "    conv_out = Conv2DZeros(width=int(inputs.shape[-1] * out_scale))\n",
        "\n",
        "    outputs = _inputs\n",
        "    outputs = tf.keras.layers.ReLU()(conv1(outputs))\n",
        "    outputs = tf.keras.layers.ReLU()(conv2(outputs))\n",
        "    outputs = conv_out(outputs)\n",
        "    return tf.keras.Model(inputs if cond is None else [inputs, cond], outputs)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gAeLVVlN3uDW",
        "colab_type": "text"
      },
      "source": [
        "# Conditional AffineCouplig Layer \n",
        "\n",
        "To construct Conditional AffineCoupling Layer, we can take two type.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OAPUDErD6Yes",
        "colab_type": "text"
      },
      "source": [
        "## scale_shift_net_template with tf.keras.layer.build (recommend)\n",
        "\n",
        "This solution uses the template function which generates tf.keras.Model\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0AIiRD4U3tXX",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "f16b3685-44c6-4e93-b41b-3822cede8070"
      },
      "source": [
        "from TFGENZOO.layers.resnet import ShallowResNet\n",
        "from TFGENZOO.flows.cond_affine_coupling import ConditionalAffineCoupling\n",
        "\n",
        "# You need to define the shape of condional input\n",
        "cond = tf.keras.Input([16, 16, 32])\n",
        "\n",
        "condAffineCoupling = ConditionalAffineCoupling(\n",
        "    scale_shift_net_template=lambda x: ShallowResNet(x, cond))\n",
        "\n",
        "# To build the template, we need to build your model.\n",
        "condAffineCoupling.build([None, 16, 16, 12])\n",
        "\n",
        "x = tf.random.normal([32, 16, 16, 12])\n",
        "c = tf.random.normal([32, 16, 16, 32])\n",
        "\n",
        "# here is the forward function\n",
        "z, ldj = condAffineCoupling(x, cond=c)\n",
        "\n",
        "# here is the inverse function\n",
        "rev_x, ildj = condAffineCoupling(z, cond=c, inverse=True)\n",
        "\n",
        "print(f'reconstruction loss is {tf.reduce_mean((rev_x - x)**2)}, and also, reconstruction effect is  {tf.reduce_mean(ldj + ildj)}')"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "initialization at actnorm_activation\n",
            "initialization at actnorm_activation_1\n",
            "reconstruction loss is 0.0, and also, reconstruction effect is  0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ms_yq1Ym7bjG",
        "colab_type": "text"
      },
      "source": [
        "## scale_shift_net with tf.keras.Model\n",
        "\n",
        "This solution uses the tf.keras.Model directly. We need to find whole layer's dimention.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pSR5_9kk7ztu",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 0
        },
        "outputId": "45036658-5b4d-48d7-dd5e-b8df387ff793"
      },
      "source": [
        "# You need to define the shape of condional input\n",
        "cond = tf.keras.Input([16, 16, 32])\n",
        "x_temp = tf.keras.Input([16, 16, 12 // 2])\n",
        "\n",
        "condAffineCoupling = ConditionalAffineCoupling(\n",
        "    scale_shift_net=ShallowResNet(x_temp, cond))\n",
        "\n",
        "# To build the template, we need to build your model.\n",
        "condAffineCoupling.build([None, 16, 16, 12])\n",
        "\n",
        "x = tf.random.normal([32, 16, 16, 12])\n",
        "c = tf.random.normal([32, 16, 16, 32])\n",
        "\n",
        "# here is the forward function\n",
        "z, ldj = condAffineCoupling(x, cond=c)\n",
        "\n",
        "# here is the inverse function\n",
        "rev_x, ildj = condAffineCoupling(z, cond=c, inverse=True)\n",
        "\n",
        "print(f'reconstruction loss is {tf.reduce_mean((rev_x - x)**2)}, and also, reconstruction effect is  {tf.reduce_mean(ldj + ildj)}')"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "initialization at actnorm_activation_2\n",
            "initialization at actnorm_activation_3\n",
            "reconstruction loss is 0.0, and also, reconstruction effect is  0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jvUhVu_K8JqH",
        "colab_type": "text"
      },
      "source": [
        "# Conditional Flow Module\n",
        "\n",
        "In previous tutorial, we introduce the TFGENZOO.flows.FlowModule to construt a step of flow.\n",
        "\n",
        "However,  in now, we need to input the condtional input into FlowModule.\n",
        "\n",
        "To tackle this problem, we serve TFGENZOO.flows.flowbase.ConditionalFlowModule."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oeO5mRGcBXIS",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "ddee9e10-374e-4a60-e10a-84f463ae0abc"
      },
      "source": [
        "from TFGENZOO.flows.flowbase import ConditionalFlowModule\n",
        "from TFGENZOO.flows import Inv1x1Conv, Actnorm\n",
        "\n",
        "channel = 12\n",
        "assert channel % 2 == 0\n",
        "\n",
        "cond = tf.keras.Input([16, 16, 32])\n",
        "\n",
        "fm = []\n",
        "fm.append(Actnorm())\n",
        "fm.append(Inv1x1Conv())\n",
        "cac = ConditionalAffineCoupling(scale_shift_net_template=lambda x: ShallowResNet(x, cond=cond))\n",
        "fm.append(cac)\n",
        "\n",
        "stepOfFlow= ConditionalFlowModule(fm)\n",
        "\n",
        "stepOfFlow.build([None, 16, 16, channel])\n",
        "\n",
        "\n",
        "x = tf.random.normal([32, 16, 16, channel])\n",
        "c = tf.random.normal([32, 16, 16, 32])\n",
        "\n",
        "# here is the forward function\n",
        "z, ldj = stepOfFlow(x, cond=c)\n",
        "\n",
        "# here is the inverse function\n",
        "rev_x, ildj = stepOfFlow(z, cond=c, inverse=True)\n",
        "\n",
        "print(f'reconstruction loss is {tf.reduce_mean((rev_x - x)**2)}, and also, reconstruction effect is  {tf.reduce_mean(ldj + ildj)}')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "initialization at actnorm\n",
            "initialization at actnorm_activation_4\n",
            "initialization at actnorm_activation_5\n",
            "reconstruction loss is 3.656933629070207e-14, and also, reconstruction effect is  0.0\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fovQ5h0qNqc2",
        "colab_type": "text"
      },
      "source": [
        "# Conclusion\n",
        "- To add the Condtional Input into Flow-based Model, we usually modify Coupling Layer.\n",
        "- TFGENZOO supports Conditional Coupling by ConditionalAffineCoupling and ConditionalFlowModule"
      ]
    }
  ]
}