conditional_flow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "conditional_flow.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1e983ddebd9c4669a41c6ef8b2dd6671": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_0aebc72ab14f4a09a3be98d89d7d02af",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_9b8be4bdd94b456cbea4ba1bacbea050",
              "IPY_MODEL_1010a30da9c24155bd19350683922f1f"
            ]
          }
        },
        "0aebc72ab14f4a09a3be98d89d7d02af": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "9b8be4bdd94b456cbea4ba1bacbea050": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_136dd165f96546d08e292d817c1e7571",
            "_dom_classes": [],
            "description": "",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 170498071,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 170498071,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_d5f020174e2d42c185e96b6825c4a0a5"
          }
        },
        "1010a30da9c24155bd19350683922f1f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_ef0b715721734686a377332bfdc80a07",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 170499072/? [01:12&lt;00:00, 2346468.00it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_36531c1064314fc5bf7ce3f266d19b63"
          }
        },
        "136dd165f96546d08e292d817c1e7571": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "d5f020174e2d42c185e96b6825c4a0a5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "ef0b715721734686a377332bfdc80a07": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "36531c1064314fc5bf7ce3f266d19b63": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hWeFrroXdnEx"
      },
      "source": [
        "# COND-GLOW"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ysHRHn043l8P"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "device=torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "import os\n",
        "from tqdm import tqdm\n",
        "import torch.nn as nn\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z5wqyu5y2PiL"
      },
      "source": [
        "class AverageMeter(object):\n",
        "    \"\"\"Computes and stores the average and current value.\n",
        "\n",
        "    Adapted from: https://github.com/pytorch/examples/blob/master/imagenet/train.py\n",
        "    \"\"\"\n",
        "    def __init__(self):\n",
        "        self.val = 0.\n",
        "        self.avg = 0.\n",
        "        self.sum = 0.\n",
        "        self.count = 0.\n",
        "\n",
        "    def reset(self):\n",
        "        self.val = 0.\n",
        "        self.avg = 0.\n",
        "        self.sum = 0.\n",
        "        self.count = 0.\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.val = val\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QQ0YjC4A3vS9"
      },
      "source": [
        "def mean_dim(tensor, dim=None, keepdims=False):\n",
        "    if dim is None:\n",
        "        return tensor.mean()\n",
        "    else:\n",
        "        if isinstance(dim, int):\n",
        "            dim = [dim]\n",
        "        dim = sorted(dim)\n",
        "        for d in dim:\n",
        "            tensor = tensor.mean(dim=d, keepdim=True)\n",
        "        if not keepdims:\n",
        "            for i, d in enumerate(dim):\n",
        "                tensor.squeeze_(d-i)\n",
        "        return tensor\n",
        "class act_norm(torch.nn.Module):\n",
        "  def __init__(self,n_feats):\n",
        "    super(act_norm,self).__init__()\n",
        "    self.register_buffer('initialized',torch.zeros(1))\n",
        "    self.mean=torch.nn.Parameter(torch.zeros(1,n_feats,1,1))\n",
        "    self.s=torch.nn.Parameter(torch.zeros(1,n_feats,1,1))\n",
        "    self.n_feats=n_feats\n",
        "\n",
        "  def init_params(self,x):\n",
        "    #actnorm is initialized so output of coupling layer had null mean and unit variance\n",
        "    #scale the data  with the negative of it's mean and inverse of sqrt of variance\n",
        "    if self.training == 0:\n",
        "      return None\n",
        "    with torch.no_grad():\n",
        "      #mean=-mean_dim(x.clone(),dim=[0,2,3],keepdims=True)\n",
        "      #var=mean_dim((x.clone()+mean) ** 2, dim=[0, 2, 3], keepdims=True)\n",
        "      mean=-x.mean(dim=[0,2,3],keepdims=True)\n",
        "      var=((x+mean)**2).mean(dim=[0,2,3],keepdims=True)\n",
        "      s=(float(1.)/(var.sqrt()+(1e-6))).log()\n",
        "      self.mean.data.copy_(mean.data)\n",
        "      self.s.data.copy_(s.data)\n",
        "      self.initialized+=1.\n",
        "\n",
        "  def forward(self,x,log_det):\n",
        "    if self.initialized == 0:\n",
        "      self.init_params(x)\n",
        "    x=x+self.mean\n",
        "    x=x*torch.exp(self.s)\n",
        "    #h*w*sum(log(s))\n",
        "    log_det_jac=x.shape[2]*x.shape[3]*torch.sum(self.s)\n",
        "    if log_det is not None:\n",
        "      log_det+=log_det_jac\n",
        "    return (x,log_det)\n",
        "  \n",
        "  def backward(self,x,log_det):\n",
        "    if self.initialized == 0:\n",
        "      self.init_params(x)\n",
        "    x=x*torch.exp(self.s*-1)\n",
        "    #h*w*sum(log(s))\n",
        "    log_det_jac=x.shape[2]*x.shape[3]*torch.sum(self.s)\n",
        "    if log_det is not None:\n",
        "      log_det-=log_det_jac\n",
        "    x=x-self.mean\n",
        "    return (x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7iZypj9-VNS6"
      },
      "source": [
        "class invertible_conv(torch.nn.Module):\n",
        "  def __init__(self,n_channels):\n",
        "    super(invertible_conv,self).__init__()\n",
        "    self.n_channels=n_channels\n",
        "    #weight matrix for 1x1 inertible convolution has to be chann x chann and orthogonal(det=0)\n",
        "    #this is reshaped into chann x chann where each row becomes (channelx1x1==>n*h*w) filter\n",
        "    #There are n_channel such matrices of this diemsnion\n",
        "    # the purpose of the invertible convolution is to permute the channels of the image instead of just splitting and \\\n",
        "    # then applying transformation. This can be done by using a fixed permutation matrix(weight of the 1d conv matrix) but\n",
        "    # GLOW decided to make it learnable instead.\n",
        "    #refer to notes    \n",
        "    weight_matrix=np.random.randn(n_channels,n_channels)\n",
        "    weight_matrix=np.linalg.qr(weight_matrix)[0]\n",
        "    weight_matrix=weight_matrix.astype(\"float32\")\n",
        "    weight_matrix=torch.from_numpy(weight_matrix)\n",
        "    self.filter=torch.nn.Parameter(weight_matrix)\n",
        "  \n",
        "  def forward(self,x,log_det):\n",
        "    #h*w*log|det(W)|\n",
        "    log_det_jacobian=x.shape[2]*x.shape[3]*torch.slogdet(self.filter)[1]\n",
        "    #print(log_det_jacobian)\n",
        "    log_det+=log_det_jacobian\n",
        "    filter=self.filter.view(self.n_channels,self.n_channels,1,1)\n",
        "    #ref_weight=ref_weight.view(self.n_channels,self.n_channels,1,1)\n",
        "    op=torch.nn.functional.conv2d(x,filter)\n",
        "    return (op,log_det)\n",
        "\n",
        "  def backward(self,x,log_det):\n",
        "    log_det_jacobian=x.shape[2]*x.shape[3]*torch.slogdet(self.filter)[1]\n",
        "    log_det-=log_det_jacobian\n",
        "    inv_filter=torch.inverse(self.filter.double()).float()\n",
        "    #ref_weight=torch.inverse(ref_weight.double()).float()\n",
        "    inv_filter=inv_filter.view(self.n_channels,self.n_channels,1,1)\n",
        "    #ref_weight=ref_weight.view(self.n_channels,self.n_channels,1,1)\n",
        "\n",
        "    op=torch.nn.functional.conv2d(x,inv_filter)\n",
        "    return (op,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TsKr20YR-LTG"
      },
      "source": [
        "# innorm,inconv,condconv==>midconv1,midcondconv1==>midnorm,midconv2,midcondconv2==>out_norm,outcon\n",
        "# condition it on label"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vT8zOIsdTDh8"
      },
      "source": [
        "class cond_coupling(torch.nn.Module):\n",
        "  def __init__(self,in_channel,cond_channel,mid_channel,norm_type=\"batch_norm\"):\n",
        "    super(cond_coupling,self).__init__()\n",
        "    self.scale=torch.nn.Parameter(torch.ones(in_channel,1,1))\n",
        "    out_channel=2*in_channel\n",
        "\n",
        "    if norm_type == \"act_norm\":\n",
        "      self.norm1=act_norm(in_channel)\n",
        "    else:\n",
        "      self.norm1=torch.nn.BatchNorm2d(in_channel)\n",
        "    self.conv1=torch.nn.Conv2d(in_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    self.cond_conv1=torch.nn.Conv2d(cond_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    nn.init.normal_(self.conv1.weight,0.,0.05)\n",
        "    nn.init.normal_(self.cond_conv1.weight,0.,0.05)\n",
        "\n",
        "    self.conv2=torch.nn.Conv2d(mid_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    self.cond_conv2=torch.nn.Conv2d(cond_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    nn.init.normal_(self.conv2.weight,0.,0.05)\n",
        "    nn.init.normal_(self.cond_conv2.weight,0.,0.05)\n",
        "    \n",
        "    if norm_type == \"act_norm\":\n",
        "      self.norm2=act_norm(mid_channel)\n",
        "    else:\n",
        "      self.norm2=torch.nn.BatchNorm2d(mid_channel)\n",
        "    self.conv3=torch.nn.Conv2d(mid_channel,mid_channel,kernel_size=(1,1),padding=0,bias=False)\n",
        "    self.cond_conv3=torch.nn.Conv2d(cond_channel,mid_channel,kernel_size=(1,1),padding=0,bias=False)\n",
        "    nn.init.normal_(self.conv3.weight,0.,0.05)\n",
        "    nn.init.normal_(self.cond_conv3.weight,0.,0.05)\n",
        "\n",
        "    if norm_type==\"act_norm\":\n",
        "      self.norm3=act_norm(mid_channel)\n",
        "    else:\n",
        "      self.norm3=torch.nn.BatchNorm2d(mid_channel)\n",
        "      self.conv4=torch.nn.Conv2d(mid_channel,out_channel,kernel_size=(3,3),padding=1,bias=True)\n",
        "      nn.init.zeros_(self.conv4.weight)\n",
        "      nn.init.zeros_(self.conv4.bias)\n",
        "\n",
        "     \n",
        "  \n",
        "  def forward(self,x,x_cond,log_det):\n",
        "    x1,x3=x.chunk(2,dim=1)\n",
        "    #s,t function\n",
        "    x2=self.norm1(x3)\n",
        "    x2=self.conv1(x2) + self.cond_conv1(x_cond)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv2(x2) + self.cond_conv2(x_cond)\n",
        "    x2=self.norm2(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv3(x2) + self.cond_conv3(x_cond)\n",
        "    x2=self.norm3(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv4(x2)\n",
        "\n",
        "    s=x2[:,0::2,:,:]\n",
        "    t=x2[:,1::2,:,:]\n",
        "\n",
        "    s=self.scale*torch.tanh(s)\n",
        "    x1=(x1+t)*torch.exp(s)\n",
        "    \n",
        "\n",
        "    log_jac=torch.sum(s,dim=(1,2,3))\n",
        "    log_det=log_det+log_jac\n",
        "    x=torch.cat((x1,x3),dim=1)\n",
        "\n",
        "    return(x,log_det)\n",
        "  \n",
        "  def backward(self,x,x_cond,log_det):\n",
        "    x1,x3=x.chunk(2,dim=1)\n",
        "    #s,t function\n",
        "    x2=self.norm1(x3)\n",
        "    #print(\"PAY ATTENTION HERE:\",x_cond.shape)\n",
        "    x2=self.conv1(x2) + self.cond_conv1(x_cond)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv2(x2) + self.cond_conv2(x_cond)\n",
        "    x2=self.norm2(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv3(x2) + self.cond_conv3(x_cond)\n",
        "    x2=self.norm3(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv4(x2)\n",
        "\n",
        "    s=x2[:,0::2,:,:]\n",
        "    t=x2[:,1::2,:,:]\n",
        "\n",
        "    s=self.scale*torch.tanh(s)\n",
        "    x1=x1*torch.exp(s*-1) - t\n",
        "    log_jac=torch.sum(s,dim=(1,2,3))\n",
        "    log_det=log_det-log_jac\n",
        "    x=torch.cat((x1,x3),dim=1)\n",
        "\n",
        "    return(x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "riAl1P-dy9D9"
      },
      "source": [
        "class condflow_module(torch.nn.Module):\n",
        "  ## act_norm ==> invertible 1x1 conv ==> coupling layer(affine) ##\n",
        "  def __init__(self,in_channel,cond_channel,mid_channel):\n",
        "    super(condflow_module,self).__init__()\n",
        "\n",
        "    self.norm1=act_norm(in_channel)\n",
        "    self.conv1=invertible_conv(in_channel)\n",
        "    self.cond_coupling1=cond_coupling(in_channel//2,cond_channel,mid_channel)\n",
        "\n",
        "  def forward(self,x,x_cond,log_det):\n",
        "    x,log_det=self.norm1.forward(x,log_det)\n",
        "    x,log_det=self.conv1.forward(x,log_det)\n",
        "    x,log_det=self.cond_coupling1.forward(x,x_cond,log_det)\n",
        "\n",
        "    return (x,log_det)\n",
        "  \n",
        "  def backward(self,x,x_cond,log_det):\n",
        "    x,log_det=self.cond_coupling1.backward(x,x_cond,log_det)\n",
        "    x,log_det=self.conv1.backward(x,log_det)\n",
        "    x,log_det=self.norm1.backward(x,log_det)\n",
        "\n",
        "    return (x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U4nqqa57wwb_"
      },
      "source": [
        "class glow(torch.nn.Module):\n",
        "  def __init__(self,in_channel,cond_channel,mid_channel,L,K):\n",
        "    super(glow,self).__init__()\n",
        "    self.glow_flows1=torch.nn.ModuleList([condflow_module(in_channel,cond_channel,mid_channel) for _ in range(K)])\n",
        "    if L>1:\n",
        "      self.glow_flows2=glow(2*in_channel,4*cond_channel,mid_channel,L-1,K)\n",
        "    else:\n",
        "      self.glow_flows2=None\n",
        "  \n",
        "  def forward(self,x,x_cond,log_det):\n",
        "    for block in self.glow_flows1:\n",
        "      x,log_det=block(x,x_cond,log_det)\n",
        "    if self.glow_flows2 is not None:\n",
        "      x=squeeze(x)\n",
        "      x_cond=squeeze(x_cond)\n",
        "      x,x2=x.chunk(2,dim=1)\n",
        "      x,log_det=self.glow_flows2.forward(x,x_cond,log_det)\n",
        "      x=torch.cat((x,x2),dim=1)\n",
        "      x=unsqueeze(x)\n",
        "      x_cond=unsqueeze(x_cond)\n",
        "    \n",
        "    return (x,log_det)\n",
        "  \n",
        "  def backward(self,x,x_cond,log_det):\n",
        "    if self.glow_flows2 is not None:\n",
        "      x=squeeze(x)\n",
        "      x_cond=squeeze(x_cond)\n",
        "      x,x2=x.chunk(2,dim=1)\n",
        "      x,log_det=self.glow_flows2.backward(x,x_cond,log_det)\n",
        "      x=torch.cat((x,x2),dim=1)\n",
        "      x=unsqueeze(x)\n",
        "      x_cond=unsqueeze(x_cond)\n",
        "    \n",
        "    for block in self.glow_flows1[::-1]:\n",
        "      x,log_det=block.backward(x,x_cond,log_det)\n",
        "    \n",
        "    return (x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IkgcgNtJRwcS"
      },
      "source": [
        "class glow_model(torch.nn.Module):\n",
        "  def __init__(self,prior_dist,n_channels,L,K):\n",
        "    super(glow_model,self).__init__()\n",
        "    self.prior=prior_dist\n",
        "    self.model=glow(in_channel=4*3,cond_channel=4,mid_channel=n_channels,L=L,K=K)#for rgb 4*3, for bw 4*1\n",
        "\n",
        "  def inference(self,x,x_cond):\n",
        "    x,log_det=preprocess(x)\n",
        "    x_cond,log_cond=preprocess(x_cond)\n",
        "    x=squeeze(x)\n",
        "    x_cond=squeeze(x_cond)\n",
        "    x,log_det=self.model.forward(x,x_cond,log_det)\n",
        "    x=unsqueeze(x)\n",
        "    return (x,log_det)\n",
        "  \n",
        "  def sampling(self,x,x_cond):\n",
        "    log_det=torch.zeros(x.shape[0])\n",
        "    log_det=log_det.to(device)\n",
        "    x_cond,_=preprocess(x_cond)\n",
        "    z=squeeze(x)\n",
        "    z_cond=squeeze(x_cond)\n",
        "    z,log_det=self.model.backward(z,z_cond,log_det)\n",
        "    z=unsqueeze(z)\n",
        "    return (z)\n",
        "  \n",
        "  def likelihood(self,x,x_cond):\n",
        "    #log(p(x))=log(ph(f(x)))+log(sii)\n",
        "    x_,log_det=self.inference(x,x_cond)\n",
        "    prior_ll=-0.5*(x_**2 + np.log(2*np.pi))\n",
        "    prior_ll=prior_ll.flatten(1).sum(-1) - np.log(256) * np.prod(x_.size()[1:])\n",
        "    return (prior_ll+log_det)\n",
        "  \n",
        "  def forward(self,x,x_cond):\n",
        "    ll=self.likelihood(x,x_cond)\n",
        "    return (ll)\n",
        "  \n",
        "  @torch.no_grad()\n",
        "  def sample_images(self,number,channel,height,width,cond_img):\n",
        "    z=self.prior.sample((number,channel,height,width))\n",
        "    z=z.to(device)\n",
        "    x=self.sampling(z,cond_img)\n",
        "    return (x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 118,
          "referenced_widgets": [
            "1e983ddebd9c4669a41c6ef8b2dd6671",
            "0aebc72ab14f4a09a3be98d89d7d02af",
            "9b8be4bdd94b456cbea4ba1bacbea050",
            "1010a30da9c24155bd19350683922f1f",
            "136dd165f96546d08e292d817c1e7571",
            "d5f020174e2d42c185e96b6825c4a0a5",
            "ef0b715721734686a377332bfdc80a07",
            "36531c1064314fc5bf7ce3f266d19b63"
          ]
        },
        "id": "uh03y9vKNkBY",
        "outputId": "26aa7891-7da3-4162-c435-f3129651029b"
      },
      "source": [
        "#datalaoder\n",
        "bs=64\n",
        "transform_train=transforms.Compose([transforms.ToTensor(),transforms.RandomHorizontalFlip(),\n",
        "                                    transforms.Resize((32,32))])\n",
        "transform_test=transforms.Compose([transforms.ToTensor(),torchvision.transforms.Resize((32,32))])\n",
        "\n",
        "trainset=torchvision.datasets. CIFAR10(root='./data', train=True,download=True,transform=transform_train)\n",
        "train_loader=torch.utils.data.DataLoader(trainset,batch_size=bs,shuffle=True, num_workers=2)\n",
        "\n",
        "testset=torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform_test)\n",
        "test_loader=torch.utils.data.DataLoader(testset, batch_size=bs,shuffle=False,num_workers=2)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1e983ddebd9c4669a41c6ef8b2dd6671",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=170498071.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Extracting ./data/cifar-10-python.tar.gz to ./data\n",
            "Files already downloaded and verified\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_d5Hj9zKe9D6"
      },
      "source": [
        "def train_model(epoch,Glow_net,train_loader,device,optimizer,scheduler):\n",
        "  Glow_net.train()\n",
        "  loss_meter=AverageMeter()\n",
        "  with tqdm(total=len(train_loader.dataset)) as progress_bar:\n",
        "    for batch_idx,data in enumerate(train_loader):\n",
        "      x=data[0].to(device)\n",
        "      cond_x=x.cpu()\n",
        "      cond_x=[torchvision.transforms.Grayscale()(i) for i in cond_x]\n",
        "      cond_x=torch.stack(cond_x).to(device)\n",
        "      optimizer.zero_grad()\n",
        "      loss=-Glow_net(x,cond_x).mean()\n",
        "      loss_meter.update(loss.item(),x.size(0))\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "      scheduler.step()\n",
        "      progress_bar.set_postfix(nll=loss_meter.avg,bpd=bits_per_dim(x,loss_meter.avg),lr=optimizer.param_groups[0][\"lr\"])\n",
        "      progress_bar.update(x.size(0))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XYC9r28Z_wQc"
      },
      "source": [
        "def test_model(epoch,Glow_net,test_loader,device,optimizer,scheduler):\n",
        "  Glow_net.eval()\n",
        "  loss_meter=AverageMeter()\n",
        "\n",
        "  with tqdm(total=len(test_loader.dataset)) as progress_bar:\n",
        "    for batch_idx,data in enumerate(test_loader):\n",
        "      x=data[0].to(device)\n",
        "      cond_x=x.cpu()\n",
        "      cond_x=[torchvision.transforms.Grayscale()(i) for i in cond_x]\n",
        "      cond_x=torch.stack(cond_x).to(device)\n",
        "      z,det=Glow_net.inference(x,cond_x)\n",
        "      prior=-0.5*(z**2+np.log(2*np.pi))\n",
        "      prior_ll=prior.flatten(1).sum(-1) - np.log(256) * np.prod(z.size()[1:])\n",
        "      loss=-(prior_ll+det).mean()\n",
        "      loss_meter.update(loss.item(),x.size(0))\n",
        "      progress_bar.set_postfix(nll=loss_meter.avg,bpd=bits_per_dim(x,loss_meter.avg))\n",
        "      progress_bar.update(x.size(0))\n",
        "  if epoch%1 == 0:\n",
        "    print('Saving...')\n",
        "    state={'net': Glow_net.state_dict(),'test_loss': loss_meter.avg,'epoch':epoch}\n",
        "    os.makedirs('ckpts',exist_ok=True)\n",
        "    torch.save(state,'ckpts/best.pth.tar')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vxAE6Y50vCMP"
      },
      "source": [
        "def bits_per_dim(x, nll):\n",
        "  dim = np.prod(x.size()[1:])\n",
        "  bpd = nll / (np.log(2) * dim)\n",
        "  return bpd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "st8Lkxt9zN0Q"
      },
      "source": [
        "mean=torch.tensor(0.)\n",
        "variance=torch.tensor(1.)\n",
        "gaussian_dist=torch.distributions.normal.Normal(mean,variance)\n",
        "Glow_net=glow_model(prior_dist=gaussian_dist,n_channels=128,L=3,K=8)\n",
        "Glow_net=Glow_net.to(device)\n",
        "optimizer=torch.optim.Adam(Glow_net.parameters(),lr=1e-3,betas=(0.9,0.999),eps=1e-8)\n",
        "scheduler=optim.lr_scheduler.LambdaLR(optimizer, lambda s: min(1.,s/500000))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PFjQcKFs0i71"
      },
      "source": [
        "def squeeze(x):\n",
        "  b,c,h,w=x.size()\n",
        "  x=x.view(b,c,h//2,2,w//2,2)\n",
        "  x=x.permute(0,1,3,5,2,4).contiguous()\n",
        "  x=x.view(b,c*2*2,h//2,w//2)\n",
        "  return (x)\n",
        "    \n",
        "def unsqueeze(x):\n",
        "  # Unsqueeze\n",
        "  b,c,h,w=x.size()\n",
        "  x=x.view(b,c//4,2,2,h,w)\n",
        "  x=x.permute(0,1,4,2,5,3).contiguous()\n",
        "  x=x.view(b,c//4,h*2,w*2)\n",
        "  return (x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PPXEZ6mL0nvz"
      },
      "source": [
        "def preprocess(x):\n",
        "  noise=torch.distributions.Uniform(0.,1.).sample(x.shape)\n",
        "  noise=noise.to(device)\n",
        "  x=(x*255. + noise)/256.\n",
        "  x*=2.\n",
        "  x-=1.\n",
        "  x*=0.9\n",
        "  x+=1.\n",
        "  x/=2.\n",
        "  logit_x=torch.log(x) - torch.log(1.-x)\n",
        "  pre_logit_scale=torch.tensor(np.log(0.9) - np.log(1.-0.9))\n",
        "  log_det=F.softplus(logit_x) + F.softplus(-logit_x) -F.softplus(-pre_logit_scale)\n",
        "  log_det=torch.sum(log_det,dim=(1,2,3))\n",
        "  x=torch.log(x)-torch.log(1.-x)\n",
        "  return (x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cP8--zxg0tV-"
      },
      "source": [
        "global_step=0\n",
        "start_epoch=0\n",
        "num_epochs=100\n",
        "for epoch in range(start_epoch, start_epoch + num_epochs):\n",
        "  train_model(epoch,Glow_net,train_loader,device,optimizer,scheduler)\n",
        "  with torch.no_grad():\n",
        "    test_model(epoch,Glow_net,test_loader,device,optimizer,scheduler)\n",
        "    origin_img,label=next(iter(test_loader))\n",
        "    cond_x=origin_img.cpu()\n",
        "    cond_x=[torchvision.transforms.Grayscale()(i) for i in cond_x]\n",
        "    cond_x=torch.stack(cond_x).to(device)\n",
        "    gen_img=Glow_net.sample_images(64,3,32,32,cond_x)\n",
        "    gen_imgs=torch.sigmoid(gen_img)\n",
        "  os.makedirs('samples_trial2',exist_ok=True)\n",
        "  torchvision.utils.save_image(gen_imgs,'samples_trial2/epoch_{}.png'.format(epoch))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ew2vapDyUEmX",
        "outputId": "2ce40c7f-a1fe-44de-92e3-cabd549dc72c"
      },
      "source": [
        "checkpoint=torch.load('ckpts/best.pth.tar')\n",
        "Glow_net.load_state_dict(checkpoint['net'])"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<All keys matched successfully>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 56
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d30IBDvBU8V2",
        "outputId": "f28cad33-b6e8-4cec-e783-758b0e8d6cb7"
      },
      "source": [
        "for batch_idx,batch in enumerate(test_loader):\n",
        "  #batch=next(iter(data_loader)\n",
        "  image=batch[0].to(device)\n",
        "  gray_batch=image.cpu()\n",
        "  gray_batch=[torchvision.transforms.Grayscale()(i) for i in gray_batch]\n",
        "  gray_batch=torch.stack(gray_batch).to(device)\n",
        "  gen_img=Glow_net.sample_images(gray_batch.shape[0],3,32,32,gray_batch)\n",
        "  gen_imgs=torch.sigmoid(gen_img)\n",
        "  os.makedirs(\"test_gen_imgs\",exist_ok=True)\n",
        "  os.makedirs(\"test_gray_imgs\",exist_ok=True)\n",
        "  torchvision.utils.save_image(gen_imgs,'test_gen_imgs/epoch_{}.png'.format(batch_idx))\n",
        "  torchvision.utils.save_image(gray_batch,'test_gray_imgs/epoch_{}.png'.format(batch_idx))\n",
        "print(\"Done\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Done\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oru75T5PvQz8",
        "outputId": "bd4d6f85-5f69-4b65-925c-c342b20d6b6a"
      },
      "source": [
        "len(os.listdir(\"test_gray_imgs\"))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "157"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 63
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jl4BhMZ-pGuE"
      },
      "source": [
        "gray_img_path=\"test_gray_imgs\"\n",
        "color_img_path=\"test_gen_imgs\"\n",
        "root=\"/content\"\n",
        "os.makedirs(\"stitched_imgs\",exist_ok=True)\n",
        "import cv2\n",
        "i=0\n",
        "for gray,color in zip(os.listdir(gray_img_path),os.listdir(color_img_path)):\n",
        "  gray_path=os.path.join(root,gray_img_path,gray)\n",
        "  color_path=os.path.join(root,color_img_path,color)\n",
        "  gray_img=cv2.imread(gray_path)\n",
        "  color_img=cv2.imread(color_path)\n",
        "  stitched_img=np.hstack((gray_img,color_img))\n",
        "  cv2.imwrite('stitched_imgs/{}.png'.format(i),stitched_img)\n",
        "  i+=1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1mjQIQBMpvIp"
      },
      "source": [
        "!zip -r /content/stitched_imgs.zip /content/stitched_imgs\n",
        "from google.colab import files\n",
        "files.download(\"/content/stitched_imgs.zip\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KrIkitfLMZfZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "722ae4ff-3280-4959-ab52-94745c22bbaf"
      },
      "source": [
        "#embedding the occupancy probablity into a higher dimension\n",
        "x=torch.tensor((0.,1.))\n",
        "x=x.unsqueeze(0)\n",
        "print(x.shape)\n",
        "layer=torch.nn.Linear(2,5)\n",
        "l=layer(x)\n",
        "print(layer(x))\n",
        "print(layer(x).shape)\n",
        "x6,x7=layer(x).chunk(2,dim=1)\n",
        "print(x6)\n",
        "print(x7)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "torch.Size([1, 2])\n",
            "tensor([[-0.9918, -0.2111, -0.3086,  0.8186, -0.2076]],\n",
            "       grad_fn=<AddmmBackward>)\n",
            "torch.Size([1, 5])\n",
            "tensor([[-0.9918, -0.2111, -0.3086]], grad_fn=<SplitBackward>)\n",
            "tensor([[ 0.8186, -0.2076]], grad_fn=<SplitBackward>)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hJDluPtrcayT"
      },
      "source": [
        "### The output dimesnion of the MLP should be same as the channel dimension of the feature map; also the same label embedding should be input into the every conditonal batch norm. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fIkyLb5eQYCY"
      },
      "source": [
        "class conditional_bn(torch.nn.Module):\n",
        "  def __init__(self,emb_dim,mlp_mid_dim,mlp_out_dim,batch_size,channels,eps=1.0e-5):\n",
        "    super(conditional_bn,self).__init__()\n",
        "    self.emb_dim=emb_dim # size of the lstm emb which is input to MLP\n",
        "    self.mlp_mid_dim=mlp_mid_dim # size of hidden layer of MLP\n",
        "    self.mlp_out_dim=mlp_out_dim # output of the MLP - for each channel\n",
        "    self.batch_size=batch_size\n",
        "    self.n_channels=channels\n",
        "    self.eps=eps\n",
        "\n",
        "    self.betas=torch.nn.Parameter(torch.zeros(self.batch_size,self.n_channels))\n",
        "    self.gammas=torch.nn.Parameter(torch.zeros(self.batch_size,self.n_channels))\n",
        "\n",
        "    self.mlp_gamma=torch.nn.Sequential(torch.nn.Linear(emb_dim,mlp_mid_dim),\n",
        "                                       torch.nn.ReLU(inplace=True),\n",
        "                                       torch.nn.Linear(mlp_mid_dim,mlp_out_dim))\n",
        "    self.mlp_betas=torch.nn.Sequential(torch.nn.Linear(emb_dim,mlp_mid_dim),\n",
        "                                      torch.nn.ReLU(inplace=True),\n",
        "                                      torch.nn.Linear(mlp_mid_dim,mlp_out_dim))\n",
        "    for m in self.modules():\n",
        "      if isinstance(m,torch.nn.Linear):\n",
        "        torch.nn.init.xavier_uniform_(m.weight)\n",
        "        torch.nn.init.constant_(m.bias,0.1)\n",
        "  \n",
        "  def forward(self,label_emb,img_feats):\n",
        "    bs,ch,height,width=img_feats.shape\n",
        "    delta_beta=self.mlp_betas(label_emb)\n",
        "    delta_gamma=self.mlp_gamma(label_emb)\n",
        "\n",
        "    betas_cloned=self.betas.clone()\n",
        "    gammas_cloned=self.gammas.clone()\n",
        "    betas_cloned+=delta_beta\n",
        "    gammas_cloned+=delta_gamma\n",
        "    \n",
        "    mean,var=torch.mean(img_feats),torch.var(img_feats,unbiased=True)\n",
        "\n",
        "    all_betas=torch.stack([betas_cloned]*img_feats.shape[2],dim=2)\n",
        "    all_betas=torch.stack([all_betas]*img_feats.shape[3],dim=3)\n",
        "    all_gammas=torch.stack([gammas_cloned]*img_feats.shape[2],dim=2)\n",
        "    all_gammas=torch.stack([all_gammas]*img_feats.shape[3],dim=3)\n",
        "\n",
        "    feat_norm=(img_feats-mean)/torch.sqrt(var+self.eps)\n",
        "    #print(feat_norm.shape)\n",
        "    #print(all_gammas.shape)\n",
        "    cbn_feat=torch.mul(feat_norm,all_gammas)+all_betas\n",
        "\n",
        "    return (cbn_feat) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ebZcRw1seRY7",
        "outputId": "a19330a9-ed99-4527-ec93-59fb8545531c"
      },
      "source": [
        "img=torch.rand(1,3,32,32)\n",
        "label=torch.tensor([0.,1.,0.,0.,0.,0.,0.,0.,0.,0.])\n",
        "label=label.unsqueeze(0)\n",
        "print(label.shape)\n",
        "sample_conv=torch.nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3,stride=1,bias=False)\n",
        "sample_mlp=torch.nn.Linear(in_features=10,out_features=16)\n",
        "sample_feat=sample_conv(img)\n",
        "label_emb=sample_mlp(label)\n",
        "print(sample_feat.shape)\n",
        "print(label_emb.shape)\n",
        "trial_cbn=conditional_bn(emb_dim=16,mlp_mid_dim=32,mlp_out_dim=32,batch_size=1,channels=32)\n",
        "print(trial_cbn)\n",
        "sample_op=trial_cbn(label_emb,sample_feat)\n",
        "print(sample_op.shape)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "torch.Size([1, 10])\n",
            "torch.Size([1, 32, 30, 30])\n",
            "torch.Size([1, 16])\n",
            "conditional_bn(\n",
            "  (mlp_gamma): Sequential(\n",
            "    (0): Linear(in_features=16, out_features=32, bias=True)\n",
            "    (1): ReLU(inplace=True)\n",
            "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
            "  )\n",
            "  (mlp_betas): Sequential(\n",
            "    (0): Linear(in_features=16, out_features=32, bias=True)\n",
            "    (1): ReLU(inplace=True)\n",
            "    (2): Linear(in_features=32, out_features=32, bias=True)\n",
            "  )\n",
            ")\n",
            "torch.Size([1, 32, 30, 30])\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SfSFsh5Ojc1j"
      },
      "source": [
        "class cbn_coupling(torch.nn.Module):\n",
        "  def _init__(self,in_channel,mid_channel,emb_channel,bs):\n",
        "    super(cbn_coupling,self).__init__()\n",
        "    self.scale=torch.nn.Parameter(torch.ones(in_channel,1,1))\n",
        "    out_channel=2*in_channel\n",
        "\n",
        "    self.label_mlp=torch.nn.Linear(in_features=10,out_features=emb_channel)\n",
        "    self.in_norm=torch.nn.BatchNorm2d(in_channel)\n",
        "    self.conv1=torch.nn.Conv2d(in_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    nn.init.normal_(self.conv1.weight,0.,0.05)\n",
        "\n",
        "    self.conv2=torch.nn.Conv2d(mid_channel,mid_channel,kernel_size=(3,3),padding=1,bias=False)\n",
        "    nn.init.normal_(self.conv2.weight,0.,0.05)\n",
        "\n",
        "    self.norm2=conditional_bn(emb_channel,mlp_mid_dim=32,mlp_out_dim=mid_channel,batch_size=bs,channels=mid_channel)\n",
        "    self.conv3=torch.nn.Conv2d(mid_channel,mid_channel,kernel_size=(1,1),padding=0,bias=False)\n",
        "    nn.init.normal_(self.conv3.weight,0.,0.05)\n",
        "\n",
        "    self.norm3=conditional_bn(emb_channel,mlp_mid_dim=32,mlp_out_dim=mid_channel,batch_size=bs,channels=mid_channel)\n",
        "    self.conv4=torch.nn.Conv2d(mid_channel,out_channel,kernel_size=(3,3),padding=1,bias=True)\n",
        "    nn.init.zeros_(self.conv4.weight)\n",
        "    nn.init.zeros_(self.conv4.bias)\n",
        "    ##1finish function ; encoder ; 2:1 paper rangenet++ ; 3: lectrure till 30minutes \n",
        "  \n",
        "  def forward(self,x,label,log_det):\n",
        "    label_emb=self.label_mlp(label)\n",
        "    x1,x3=x.chunk(2,dim=1)\n",
        "    #s,t function\n",
        "    x2=self.in_norm(x3)\n",
        "    x2=self.conv1(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv2(x2)\n",
        "    x2=self.norm2(label_emb,x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv3(x2)\n",
        "    x2=self.norm3(label_emb,x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv4(x2)\n",
        "\n",
        "    s=x2[:,0::2,:,:]\n",
        "    t=x2[:,1::2,:,:]\n",
        "\n",
        "    s=self.scale*torch.tanh(s)\n",
        "    x1=(x1+t)*torch.exp(s)\n",
        "    log_jac=torch.sum(s,dim=(1,2,3))\n",
        "    log_det=log_det+log_jac\n",
        "    x=torch.cat((x1,x3),dim=1)\n",
        "\n",
        "    return(x,log_det)\n",
        "  \n",
        "  def backward(self,x,label,log_det):\n",
        "    label_emb=self.label_mlp(label)\n",
        "    x1,x3=x.chunk(2,dim=1)\n",
        "    #s,t function\n",
        "    x2=self.in_norm(x3)\n",
        "    x2=self.conv1(x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv2(x2)\n",
        "    x2=self.norm2(label_emb,x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv3(x2)\n",
        "    x2=self.norm3(label_emb,x2)\n",
        "    x2=F.relu(x2)\n",
        "\n",
        "    x2=self.conv4(x2)\n",
        "\n",
        "    s=x2[:,0::2,:,:]\n",
        "    t=x2[:,1::2,:,:]\n",
        "\n",
        "    s=self.scale*torch.tanh(s)\n",
        "    x1=x1*torch.exp(s*-1) - t\n",
        "    log_jac=torch.sum(s,dim=(1,2,3))\n",
        "    log_det=log_det-log_jac\n",
        "    x=torch.cat((x1,x3),dim=1)\n",
        "\n",
        "    return(x,log_det)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iRMUeVZHiug2"
      },
      "source": [
        "img=torch.rand(1,3,32,32)\n",
        "label=torch.tensor([0.,1.,0.,0.,0.,0.,0.,0.,0.,0.])\n",
        "label=label.unsqueeze(0)\n",
        "print(label.shape)\n",
        "sample_conv=torch.nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3,stride=1,bias=False)\n",
        "sample_mlp=torch.nn.Linear(in_features=10,out_features=16)\n",
        "sample_feat=sample_conv(img)\n",
        "label_emb=sample_mlp(label)\n",
        "samp_coup=cbn_coupling(in_channel=32,mid_channel=32,emb_channel=16,bs=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FfynT8BNRxPw"
      },
      "source": [
        "#shapenet code (broyden implicit poof)(till game) + (encoder decoder) + git clone"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X6fuK8Np86Ig"
      },
      "source": [
        "#1st half==> 1deep equilibrium models talk and understand,#\n",
        "#2 2-3shapenet code(tomorrow) + till 2 1mdeq code recreate  \n",
        "#2nd half==>3nf1 code,4mg2mesh code,5tokyo appn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dfKPxl0NG51Z"
      },
      "source": [
        "#use conditional batch norm for the segmentation labels\n",
        "#use all layers condtion with all layers"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}