invertible-mnist-logdet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "import multiprocessing as mp\n",
    "mp.set_start_method('forkserver')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch\n",
    "torch.cuda.set_device(1)\n",
    "import torch.nn as nn\n",
    "from torch.nn.utils import weight_norm, spectral_norm\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "import numpy as np \n",
    "from matplotlib import pyplot as plt\n",
    "from torch.optim import *\n",
    "from InvertibleResnet import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "transform = transforms.Compose(\n",
    "    [   transforms.Pad(2),\n",
    "        transforms.ToTensor(),     \n",
    "     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "\n",
    "trainset = torchvision.datasets.MNIST(root='./data', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,\n",
    "                                          shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testset = torchvision.datasets.MNIST(root='./data', train=False,\n",
    "                                       download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=256,\n",
    "                                         shuffle=True, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testnet = InvertibleResnetConv(1,32, list_num_blocks=(32,32,32)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = testnet(next(iter(trainloader))[0].cuda(), return_logdet=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.distributions import *\n",
    "std = 1\n",
    "ndim = 32*32\n",
    "height,width = 32,32\n",
    "n_bins = 5\n",
    "pzes = []\n",
    "logdets = []\n",
    "losses = []\n",
    "latent_numpys =[]\n",
    "gaussian = Normal(0.0,std)\n",
    "normal_np = gaussian.sample((20000,2)).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = InvertibleResnetConv(1,32, list_num_blocks=(32,32,32)).cuda()\n",
    "optim = Adam(net.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent = net(testset[0][0].unsqueeze(0).cuda(), return_logdet=False)\n",
    "fixed_z = [torch.randn_like(l[:1]).cuda() for l in latent]\n",
    "print ([l.shape for l in fixed_z])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "fig,(ax,ax2) = plt.subplots(1,2,figsize=(18,8))\n",
    "axim = None\n",
    "axscatter = None\n",
    "fig.canvas.draw()\n",
    "\n",
    "use_normal = True\n",
    "\n",
    "def torchSigWMW( target,pred, beta=8):    \n",
    "    mask = target.type(torch.cuda.ByteTensor)\n",
    "    x = torch.masked_select(pred, 1+(-1)*mask).view(1,-1)\n",
    "    y = torch.masked_select(pred, mask).view(-1,1)\n",
    "    xn = x.expand(y.size(0), x.size(1))\n",
    "    yn = y.expand(y.size(0), x.size(1))\n",
    "    ur = torch.sigmoid(-beta*(xn-yn))\n",
    "    return torch.sum(ur) / torch.sum(target) / torch.sum(1-target)\n",
    "\n",
    "def criterion(pred, target):\n",
    "    # return ((pred - (target * dist * 2 - dist)) ** 2).mean() \n",
    "    return torchSigWMW(target, pred)\n",
    "\n",
    "\n",
    "for epoch in tqdm(range(0,200)):  \n",
    "    itercount = 0\n",
    "    for inp, _ in trainloader:\n",
    "        \n",
    "        net.train()\n",
    "        optim.zero_grad()\n",
    "        inp = inp.cuda()\n",
    "        latent, logdet = net(inp + (torch.rand_like(inp).cuda() - 0.5) / ( n_bins)  , return_logdet=True, num_logdet_iter=1) \n",
    "        logdet /= inp.shape[2] * inp.shape[3]\n",
    "        flat_latent = torch.cat([l.view(inp.shape[0], -1) for l in latent], dim=1)\n",
    "        latentstd = torch.mean(flat_latent.std(dim=0))\n",
    "        latentstdmax = torch.max(flat_latent.std(dim=0))\n",
    "        latentstdmin = torch.min(flat_latent.std(dim=0))\n",
    "\n",
    "        pz = -(flat_latent ** 2).sum(dim=1) / 2 / std /std            \n",
    "        pz = pz.mean()\n",
    "        pz /= inp.shape[2] * inp.shape[3]\n",
    "        pz -= 0.5 * np.log(2 * 3.1415 * std * std)\n",
    "        \n",
    "        np_loss = ( np.log(n_bins) -logdet.item()-pz.item()) / np.log(2)\n",
    "        print ('loss:{:.4f}, logdet:{:.4f} pz:{:.4f} latentstd:{:.4f} latentstdmax:{:.4f} latentstdmin:{:.4f}'.format(\n",
    "                np_loss, logdet.item(),pz.item(),latentstd.item(),latentstdmax.item(),latentstdmin.item()))\n",
    "        ( - logdet - pz).backward()\n",
    "        if torch.isnan(logdet):\n",
    "            raise Exception()\n",
    "        pzes.append(pz.item())\n",
    "        logdets.append(logdet.item())\n",
    "        losses.append(np_loss)\n",
    "        torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)\n",
    "        optim.step()\n",
    "\n",
    "        if itercount % 10 == 0:\n",
    "            net.eval()\n",
    "            latents = []\n",
    "            with torch.no_grad():\n",
    "                latent_numpy = []\n",
    "                for test_batch, _ in testloader:\n",
    "                    test_batch = test_batch.cuda()\n",
    "                    latent = net(test_batch + (torch.rand_like(test_batch).cuda() - 0.5) / ( n_bins), return_logdet=False)\n",
    "                    _latent_numpy = torch.cat([l.contiguous().view(test_batch.shape[0], -1) for l in latent], dim=1).detach().cpu().numpy()\n",
    "                    latent_numpy.append(_latent_numpy)\n",
    "                    break\n",
    "                latent_numpy = np.concatenate(latent_numpy)\n",
    "\n",
    "                pred = net([l[:1] for l in latent], reverse=True, reverse_iterations=200)                \n",
    "                pred_sample = net(fixed_z, reverse=True, reverse_iterations=200)\n",
    "            \n",
    "            \n",
    "\n",
    "            recon_loss = F.mse_loss(test_batch[:1], pred)\n",
    "            if itercount >0:\n",
    "                print ('loss:{:.4f}, logdet:{:.4f} pz:{:.4f} latentstd:{:.4f} latentmean:{:.4f} recon_loss:{:.4f}'.format(\n",
    "                np_loss, logdet.item(),pz.item(), np.mean(latent_numpy.std(axis=0)),np.mean(latent_numpy.mean(axis=0)),  recon_loss.item()))\n",
    "\n",
    "\n",
    "\n",
    "            pred_np = pred[0].cpu().detach().numpy().reshape((height,width))\n",
    "            inp_np = test_batch[0].cpu().detach().numpy().reshape((height,width))            \n",
    "            pred_sample_np = pred_sample[0].cpu().detach().numpy().reshape((height,width))\n",
    "            disp = np.concatenate([pred_np, inp_np, pred_sample_np], axis=1)\n",
    "            if axim:\n",
    "                axim.set_data(disp)\n",
    "            else:\n",
    "                axim = ax.imshow(disp, vmin=-1.0, vmax=1.0)\n",
    "\n",
    "\n",
    "            ax2.clear()\n",
    "            ax2.scatter(normal_np[:1000,0], normal_np[:1000,1], c='g')            \n",
    "            ax2.scatter(latent_numpy[:1000,-1], latent_numpy[:1000,-2],  c='r')  \n",
    "            fig.canvas.draw()\n",
    "            latent_numpys.append(latent_numpy) \n",
    "        itercount += 1\n",
    "        del pz,  logdet, latent, latentstd\n",
    "    #with torch.no_grad():\n",
    "    #    for name, param in net.named_parameters():\n",
    "    #        if '.weight' in name:\n",
    "    #            param /= torch.sqrt((param ** 2).sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_numpys = np.array(latent_numpys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dists = torch.from_numpy(np.array([[dist,-dist], [dist, dist], [-dist, dist], [-dist,-dist]]).astype(np.float32)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_numpys.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_numpys = latent_numpys.reshape((-1,80,2560,2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_numpy = latent_numpys[0,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_numpy.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.std(latent_numpy[np.logical_and(latent_numpy[:,0] > 0, latent_numpy[:,1] < 0)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.mean(latent_numpy[np.logical_and(latent_numpy[:,0] < 0, latent_numpy[:,1] < 0)], axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "fig,ax =  plt.subplots(1,1,figsize=(8,8))\n",
    "def scroll(event):\n",
    "    global i, latent_numpys\n",
    "    oldi = i\n",
    "    if event.key == 'e':\n",
    "        i -= 1\n",
    "    elif event.key =='d':\n",
    "        i += 1\n",
    "    if i < 0: i = 0\n",
    "    if i >= len(latent_numpys[0]) - 0: i = len(latent_numpys[0]) - 1\n",
    "    if oldi != i:\n",
    "        ax.clear()\n",
    "        latent_numpy = latent_numpys[0][i][-1]\n",
    "        \n",
    "        ax.clear()\n",
    "        _ = ax.hist(latent_numpys[0][i][:,1], bins=100, density=True)\n",
    "        _ = ax.hist(normal_np[:,1], bins=100, density=True)\n",
    "        \n",
    "        fig.canvas.draw()\n",
    "        #ax.set_xlim(-3,3)\n",
    "        #ax.set_ylim(-3,3)\n",
    "\n",
    "\n",
    "cid1 = fig.canvas.mpl_connect('key_press_event', scroll)\n",
    "i = 0\n",
    "latent_numpy = latent_numpys[0][i][-1]\n",
    "ax.clear()\n",
    "_ = ax.hist(latent_numpys[0][i][:,1], bins=100, density=True)\n",
    "_ = ax.hist(normal_np[:,1], bins=100, density=True)\n",
    "\n",
    "\n",
    "fig.canvas.draw()\n",
    "#ax.set_xlim(-3,3)\n",
    "#ax.set_ylim(-3,3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "fig,ax =  plt.subplots(1,1,figsize=(8,8))\n",
    "def scroll(event):\n",
    "    global i, latent_numpys\n",
    "    oldi = i\n",
    "    if event.key == 'e':\n",
    "        i -= 1\n",
    "    elif event.key =='d':\n",
    "        i += 1\n",
    "    if i < 0: i = 0\n",
    "    if i >= len(latent_numpys[0]) - 0: i = len(latent_numpys[0]) - 1\n",
    "    if oldi != i:\n",
    "        ax.clear()\n",
    "        latent_numpy = latent_numpys[0,i]\n",
    "        \n",
    "        ax.clear()\n",
    "        ax.scatter(normal_np[:,0], normal_np[:,1], c='g') \n",
    "        ax.scatter(latent_numpy[:500,0], latent_numpy[:500,1],  c='r')  \n",
    "        fig.canvas.draw()\n",
    "        #ax.set_xlim(-3,3)\n",
    "        #ax.set_ylim(-3,3)\n",
    "\n",
    "\n",
    "cid1 = fig.canvas.mpl_connect('key_press_event', scroll)\n",
    "i = 0\n",
    "latent_numpy = latent_numpys[0,i]\n",
    "ax.clear()\n",
    "ax.scatter(normal_np[:,0], normal_np[:,1], c='g')\n",
    "ax.scatter(latent_numpy[:500,0], latent_numpy[:500,1],  c='r')  \n",
    "fig.canvas.draw()\n",
    "#ax.set_xlim(-3,3)\n",
    "#ax.set_ylim(-3,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.plot(range(len(losses[:])), losses[:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.plot(range(len(logdets[:])), logdets[:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.plot(range(len(pzes[250:])), pzes[250:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}