notebooks/celeba/Continue_Adversarial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import cm\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'png'\n",
    "matplotlib.rcParams['figure.figsize'] = (12.0, 4.0)\n",
    "matplotlib.rcParams['font.size'] = 7\n",
    "\n",
    "import matplotlib.lines as mlines\n",
    "import seaborn\n",
    "seaborn.set_style('darkgrid')\n",
    "import logging\n",
    "import importlib\n",
    "importlib.reload(logging) # see https://stackoverflow.com/a/21475297/1469195\n",
    "log = logging.getLogger()\n",
    "log.setLevel('DEBUG')\n",
    "import sys\n",
    "logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',\n",
    "                     level=logging.DEBUG, stream=sys.stdout)\n",
    "seaborn.set_palette('colorblind')\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "\n",
    "import itertools\n",
    "from reversible.plot import create_bw_image, create_rgb_image\n",
    "from PIL import Image\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Your celeba 64x64 images\n",
    "image_folder = '/data/schirrmr/schirrmr/celeba/CELEB_64/'\n",
    "# For  Frechet Inception distance computation within this training\n",
    "# https://drive.google.com/open?id=1H5dnLqe1UgwDZK3wewdqJvbl5pmSZxh1\n",
    "# These are not the eventual Frechet Inception Distances reported in paper \n",
    "# as final results, they are based on https://github.com/bioinf-jku/TTUR\n",
    "# in tensorflow\n",
    "statistic_file = '/data/schirrmr/schirrmr/reversible-icml/fid_stats_celeb64_resize_normalize.npz'\n",
    "# For loading the previously model trained by clamping, also load means stds \n",
    "load_folder = '/data//schirrmr/schirrmr/reversible-icml/models/CelebA_More_Examples_Only_Clamp/218//'\n",
    "# For saving the continued trained model\n",
    "model_save_folder = '/data//schirrmr/schirrmr/reversible-icml/models/celeba/Continue_Adversarial/'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log.setLevel('WARNING')\n",
    "n_examples = 150000\n",
    "images = []\n",
    "for i_image in range(n_examples):\n",
    "    images.append(np.array(Image.open(os.path.join(image_folder, 'celeb_{:d}.png'.format(\n",
    "        i_image)))))\n",
    "    if i_image % 1000 == 0:\n",
    "        print(\"Loaded {:d}...\".format(i_image))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log.setLevel('INFO')\n",
    "x = np.array(images).astype(np.float32).transpose(0,3,1,2) /255.0\n",
    "plt.imshow(x.transpose(0,2,3,1)[0], interpolation='nearest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from reversible.util import np_to_var, var_to_np\n",
    "inputs = np_to_var(x, dtype=np.float32)\n",
    "# we put all faces into one class \n",
    "targets = np_to_var(np.ones((len(x), 1)), dtype=np.float32)\n",
    "del x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from reversible.models import create_celebA_model\n",
    "from reversible.revnet import SubsampleSplitter, ViewAs\n",
    "from reversible.util import set_random_seeds\n",
    "import torch as th\n",
    "torch.backends.cudnn.benchmark = True\n",
    "set_random_seeds(3049, True)\n",
    "feature_model = create_celebA_model()\n",
    "feature_model = feature_model.cuda()\n",
    "n_wanted_stds = 64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load old params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "means_per_cluster = th.load(load_folder + '/means.pkl')\n",
    "\n",
    "stds_per_cluster = th.load(load_folder + '/stds.pkl')\n",
    "feature_model.load_state_dict(th.load(load_folder + '/model_dict.pkl'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "optimizer = th.optim.Adam(\n",
    "                          [\n",
    "    {'params': list(feature_model.parameters()),\n",
    "    'lr': 0.0001,\n",
    "    'weight_decay': 0},], betas=(0,0.9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create adversary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from reversible.spectral_norm import SpectralNorm\n",
    "from torch import nn\n",
    "class ConcatReLU(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ConcatReLU, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return th.cat((nn.functional.relu(x), -nn.functional.relu(-x)), dim=1)\n",
    "\n",
    "def create_adv_2_layer(dim, intermediate_dim, snorm):\n",
    "    if snorm is None:\n",
    "        model = nn.Sequential(\n",
    "            nn.Linear(dim, intermediate_dim),\n",
    "            ConcatReLU(),\n",
    "            nn.Linear(intermediate_dim*2, intermediate_dim*2),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(intermediate_dim*2, 1))\n",
    "    else:\n",
    "        model = nn.Sequential(\n",
    "            SpectralNorm(nn.Linear(dim, intermediate_dim), power_iterations=1, to_norm=snorm),\n",
    "            ConcatReLU(),\n",
    "            SpectralNorm(nn.Linear(intermediate_dim*2, intermediate_dim*2), power_iterations=1, to_norm=snorm),\n",
    "            nn.ReLU(),\n",
    "            SpectralNorm(nn.Linear(intermediate_dim*2, 1), power_iterations=1, to_norm=snorm))\n",
    "\n",
    "    model = model.cuda()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adversary = create_adv_2_layer(n_wanted_stds, 200, snorm=1)\n",
    "\n",
    "optim_adv = th.optim.Adam([{\n",
    "    'params': [p for p in adversary.parameters() if p.requires_grad],\n",
    "    'lr:': 0.0004, 'weight_decay': 0.00}],\n",
    "                         betas=(0,0.9))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from reversible.gaussian import get_gauss_samples\n",
    "from reversible.revnet import invert \n",
    "\n",
    "def take_only_large_stds(l_out, std, n_wanted_stds):\n",
    "    i_stds = th.sort(std)[1][-n_wanted_stds:]\n",
    "    l_out = l_out.index_select(index=i_stds, dim=1)\n",
    "    return l_out\n",
    "\n",
    "def train_critic_on_batch(b_X, n_wanted_stds):\n",
    "    start_time = time.time()\n",
    "    result = {}\n",
    "    i_class = 0\n",
    "    mean = means_per_cluster[i_class]\n",
    "    std = stds_per_cluster[i_class] * stds_per_cluster[i_class]\n",
    "    outs_real = feature_model(b_X)\n",
    "    samples = get_gauss_samples(750, mean, std)\n",
    "    # first sample_loss\n",
    "    score_real = adversary(take_only_large_stds(outs_real,std, n_wanted_stds=n_wanted_stds))\n",
    "    score_fake = adversary(take_only_large_stds(samples, std, n_wanted_stds=n_wanted_stds))\n",
    "    d_loss = nn.functional.relu(1.0 - score_real).mean() + nn.functional.relu(\n",
    "                1.0 + score_fake).mean()\n",
    "    optim_adv.zero_grad()\n",
    "    d_loss.backward()\n",
    "    all_params = [p for group in optim_adv.param_groups for p in group['params']]\n",
    "    grad_norm = th.nn.utils.clip_grad_norm_(all_params, 1000, 2)\n",
    "    optim_adv.step()\n",
    "    runtime = time.time() - start_time\n",
    "    result['d_loss'] = var_to_np(d_loss)\n",
    "    result['runtime_d'] = runtime\n",
    "    result['grad_norm_D'] = grad_norm\n",
    "    return result\n",
    "\n",
    "def train_gen_on_batch(b_X, only_clamp, n_wanted_stds):\n",
    "    start_time = time.time()\n",
    "    result = {}\n",
    "    i_class = 0\n",
    "    mean = means_per_cluster[i_class]\n",
    "    std = stds_per_cluster[i_class] * stds_per_cluster[i_class]\n",
    "    outs_real = feature_model(b_X)\n",
    "    if only_clamp == False:\n",
    "        samples = get_gauss_samples(len(outs_real) * 3, mean, std)\n",
    "        # sample loss\n",
    "        score_real = adversary(take_only_large_stds(outs_real, std, n_wanted_stds=n_wanted_stds))\n",
    "        g_loss_samples = th.mean(score_real)\n",
    "    else:\n",
    "        g_loss_samples = th.zeros(1).cuda()\n",
    "    \n",
    "    min_vals = (mean - 20 * std).unsqueeze(0)\n",
    "    max_vals = (mean + 20 * std).unsqueeze(0)\n",
    "    outs_clamped = th.min(th.max(outs_real,min_vals.detach()), max_vals.detach())\n",
    "    \n",
    "    out_clamp_loss = th.mean((outs_clamped - outs_real) ** 2) * 10\n",
    "    inverted_clamped = invert(feature_model, outs_clamped)\n",
    "    in_clamp_loss = th.mean(th.abs(b_X - inverted_clamped)) * 10\n",
    "    \n",
    "    if only_clamp == 'in':\n",
    "        out_clamp_loss = th.autograd.Variable(th.zeros(1).cuda())\n",
    "    g_loss = in_clamp_loss + out_clamp_loss + g_loss_samples\n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    g_loss.backward()\n",
    "    all_params = [p for group in optimizer.param_groups for p in group['params']]\n",
    "    grad_norm = th.nn.utils.clip_grad_norm_(all_params, 5, 2)\n",
    "    optimizer.step()\n",
    "    runtime = time.time() - start_time\n",
    "    result['g_loss'] = var_to_np(g_loss)\n",
    "    result['g_loss_samples'] = var_to_np(g_loss_samples)\n",
    "    result['in_loss'] = var_to_np(in_clamp_loss)\n",
    "    result['out_loss'] = var_to_np(out_clamp_loss)\n",
    "    result['runtime_g'] = runtime\n",
    "    result['grad_norm_g'] = grad_norm\n",
    "    return result\n",
    "\n",
    "from reversible.sliced import sliced_from_samples\n",
    "# let's add sliced metric\n",
    "def compute_sliced_dist_on_outputs(n_wanted_stds):\n",
    "    # Sliced distance outputs/gaussian samples\n",
    "    # normalized by sliced distance of two gaussian samples\n",
    "    i_class = 0\n",
    "    mean = means_per_cluster[i_class]\n",
    "    std = stds_per_cluster[i_class] * stds_per_cluster[i_class]\n",
    "    with th.no_grad():\n",
    "        reduced_outs = [var_to_np(take_only_large_stds(\n",
    "            feature_model(inputs[i:i+500].cuda()), stds_per_cluster[0] * stds_per_cluster[0], n_wanted_stds))\n",
    "                for i in range(0, len(inputs) // 10, 500)]\n",
    "    reduced_outs = np.concatenate(reduced_outs, axis=0)\n",
    "    reduced_mean = take_only_large_stds(mean.unsqueeze(0),std,n_wanted_stds=n_wanted_stds).squeeze(0)\n",
    "    reduced_std = take_only_large_stds(std.unsqueeze(0),std,n_wanted_stds=n_wanted_stds).squeeze(0)\n",
    "\n",
    "    gauss_samples = get_gauss_samples(len(reduced_outs), reduced_mean, reduced_std)\n",
    "\n",
    "    sliced_dist = sliced_from_samples(np_to_var(reduced_outs, dtype=np.float32).cuda(),\n",
    "                        gauss_samples, n_dirs=2, adv_dirs=None, orthogonalize=True,\n",
    "                        dist='sqw2')\n",
    "\n",
    "    gauss_samples_2 = get_gauss_samples(len(reduced_outs), reduced_mean, reduced_std)\n",
    "    sliced_ref = sliced_from_samples(gauss_samples_2,\n",
    "                        gauss_samples, n_dirs=2, adv_dirs=None, orthogonalize=True,\n",
    "                        dist='sqw2')\n",
    "    sliced_rel = sliced_dist / sliced_ref\n",
    "    return sliced_rel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### for inception distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from reversible.fid_score import calculate_activation_statistics\n",
    "from reversible.fid_score import calculate_frechet_distance\n",
    "from reversible.inception import InceptionV3\n",
    "\n",
    "ref_vals = np.load(statistic_file)\n",
    "mu_ref = ref_vals['mu']\n",
    "sig_ref = ref_vals['sig']\n",
    "\n",
    "model = InceptionV3(resize_input=True, normalize_input=True)\n",
    "model = model.cuda()\n",
    "\n",
    "def generate_examples(n_examples):\n",
    "    samples = get_gauss_samples(n_examples, means_per_cluster[0], stds_per_cluster[0],)\n",
    "    examples = var_to_np(invert(feature_model, samples)).astype(np.float64)\n",
    "    return examples\n",
    "def calculate_current_fid():\n",
    "    examples = [generate_examples(500) for _ in range(10)]\n",
    "    examples = np.concatenate(examples, axis=0)\n",
    "    mu, sig = calculate_activation_statistics(np.clip(examples, a_min=0, a_max=1), model, cuda=True,\n",
    "                                         verbose=True)\n",
    "    fid = calculate_frechet_distance(mu, sig, mu_ref, sig_ref)\n",
    "    return fid\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calculate_current_fid() # should be  276.5135649535918"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from reversible.iterators import BalancedBatchSizeIterator\n",
    "\n",
    "batch_size = 250\n",
    "iterator = BalancedBatchSizeIterator(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "epochs_dataframe = pd.DataFrame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "min_fid = np.inf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from reversible.gaussian import get_gauss_samples\n",
    "for i_epoch in range(2000):\n",
    "    start_time = time.time()\n",
    "    g_results = []\n",
    "    d_results = []\n",
    "    i_update = 0\n",
    "    for b_X, b_y in iterator.get_batches(inputs, targets,shuffle=True):\n",
    "        b_X = b_X.cuda()\n",
    "        if (i_update % 2 == 1):\n",
    "            g_result = train_gen_on_batch(b_X, only_clamp=False, n_wanted_stds=n_wanted_stds)\n",
    "            g_results.append(g_result)\n",
    "        else:\n",
    "            d_result = train_critic_on_batch(b_X, n_wanted_stds=n_wanted_stds)\n",
    "            d_results.append(d_result)\n",
    "        i_update += 1\n",
    "    result =  {**pd.DataFrame(g_results).mean(), **pd.DataFrame(d_results).mean()}\n",
    "    relative_sliced = compute_sliced_dist_on_outputs(n_wanted_stds)\n",
    "    starttime_fid = time.time()\n",
    "    fid = calculate_current_fid()\n",
    "    runtime_fid = time.time() - starttime_fid\n",
    "    \n",
    "    epoch_time = time.time() - start_time\n",
    "    result['runtime_fid'] = runtime_fid\n",
    "    result['runtime'] = epoch_time\n",
    "    result['sliced_rel'] = var_to_np(relative_sliced)\n",
    "\n",
    "    result['fid'] = fid\n",
    "    epochs_dataframe = epochs_dataframe.append(result, ignore_index=True)\n",
    "    if i_epoch % 1 == 0:\n",
    "        display(epochs_dataframe.iloc[-1:])\n",
    "        fig = plt.figure(figsize=(8,3))\n",
    "        df_copy = epochs_dataframe.copy()\n",
    "        df_copy.fid = df_copy.fid / 10\n",
    "        df_copy = df_copy.drop('runtime', axis=1)\n",
    "        df_copy = df_copy.drop('runtime_fid', axis=1)\n",
    "        df_copy.plot(ax=fig.gca())\n",
    "        display(fig)\n",
    "        plt.close(fig)\n",
    "        i_class = 0\n",
    "        mean = means_per_cluster[i_class]\n",
    "        std = stds_per_cluster[i_class] * stds_per_cluster[i_class]\n",
    "        outs_real = feature_model(inputs[:1000].cuda())\n",
    "        fig = plt.figure(figsize=(8,3))\n",
    "        plt.plot(var_to_np(stds_per_cluster[0]))\n",
    "        plt.plot(var_to_np(th.std(outs_real, dim=0)))\n",
    "        display(fig)\n",
    "        plt.close(fig)\n",
    "        \n",
    "        i_std_1, i_std_2 = np.argsort(var_to_np(stds_per_cluster[0]))[::-1][:2]\n",
    "        feature_a_values = th.linspace(float(mean[i_std_1].data - 2 * std[i_std_1].data),\n",
    "                               float(mean[i_std_1].data + 2 * std[i_std_1].data), 9)\n",
    "        feature_b_values = th.linspace(float(mean[i_std_2].data - 2 * std[i_std_2].data),\n",
    "                               float(mean[i_std_2].data + 2 * std[i_std_2].data), 9)\n",
    "        image_grid = np.zeros((len(feature_a_values), len(feature_b_values), 3, 64,64))\n",
    "\n",
    "        for i_f_a_val, f_a_val in enumerate(feature_a_values):\n",
    "            for i_f_b_val, f_b_val in enumerate(feature_b_values):\n",
    "                this_out = mean.clone()\n",
    "                this_out.data[i_std_1] = f_a_val\n",
    "                this_out.data[i_std_2] = f_b_val\n",
    "                inverted = var_to_np(invert(feature_model, this_out.unsqueeze(0))[0]).squeeze()\n",
    "\n",
    "                image_grid[i_f_a_val, i_f_b_val] = np.copy(inverted)\n",
    "        im = create_rgb_image(image_grid[::-1]).resize((6*100,6*100))\n",
    "        display(im)\n",
    "        \n",
    "        \n",
    "        samples = get_gauss_samples(40, means_per_cluster[0], stds_per_cluster[0],)\n",
    "\n",
    "        inverted = var_to_np(invert(feature_model, samples)).astype(np.float64)\n",
    "\n",
    "        inverted = inverted.reshape(5,8,3,64,64)\n",
    "        im = create_rgb_image(inverted).resize((8*64,5*64))\n",
    "        display(im)\n",
    "    if (i_epoch % 30 == 0) or (fid < min_fid):\n",
    "        folder =  os.path.join(model_save_folder, str(len(epochs_dataframe)))\n",
    "        os.makedirs(folder, exist_ok=False)\n",
    "        pickle.dump(epochs_dataframe, open(os.path.join(folder, 'epochs_df.pkl'), 'wb'))\n",
    "        th.save(optimizer.state_dict(), os.path.join(folder, 'optim_dict.pkl'))\n",
    "        th.save(feature_model.state_dict(), os.path.join(folder, 'model_dict.pkl'))\n",
    "        th.save(means_per_cluster, os.path.join(folder, 'means.pkl'))\n",
    "        th.save(stds_per_cluster, os.path.join(folder, 'stds.pkl'))\n",
    "        log.info(\"Saved to {:s}\".format(folder))\n",
    "        if fid < min_fid:\n",
    "            min_fid = fid"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit 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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}