wgan_ssl_wn.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import time\n",
    "import csv\n",
    "\n",
    "os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID' \n",
    "os.environ['CUDA_VISIBLE_DEVICES']='1'\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import h5py\n",
    "import imageio\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.utils import weight_norm\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch.optim as optim\n",
    "# from torchvision import datasets, transforms\n",
    "# from torch.autograd import Variable\n",
    "import torch.autograd as autograd\n",
    "import torchvision.utils as vutils\n",
    "\n",
    "from torch.utils.data import DataLoader, TensorDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data():\n",
    "    x_train = (h5py.File('camelyonpatch_level_2_split_train_x.h5', 'r')['x'][:, 16:80,16:80] - 127.5) / 127.5\n",
    "    y_train = h5py.File('camelyonpatch_level_2_split_train_y.h5', 'r')['y'][:].reshape(-1,1)\n",
    "    x_test = (h5py.File('camelyonpatch_level_2_split_test_x.h5', 'r')['x'][:, 16:80,16:80] - 127.5) / 127.5\n",
    "    y_test = h5py.File('camelyonpatch_level_2_split_test_y.h5', 'r')['y'][:].reshape(-1,1)\n",
    "    x_valid = (h5py.File('camelyonpatch_level_2_split_valid_x.h5', 'r')['x'][:, 16:80,16:80] - 127.5) / 127.5\n",
    "    y_valid = h5py.File('camelyonpatch_level_2_split_valid_y.h5', 'r')['y'][:].reshape(-1,1)\n",
    "              \n",
    "    return x_train, y_train, x_test, y_test, x_valid, y_valid\n",
    "\n",
    "def plot_samples(samples, folder=None, epoch=None, i=None):\n",
    "    rt = int(np.sqrt(samples.shape[0]))\n",
    "    r, c = rt, rt\n",
    "    # r, c = 6, 12\n",
    "\n",
    "    generatedImage = 0.5 * samples + 0.5\n",
    "\n",
    "    fig = plt.figure(figsize=(10,10))\n",
    "\n",
    "    axs = [fig.add_subplot(r,c,i+1) for i in range(r*c)]\n",
    "    cnt = 0\n",
    "    for ax in axs:\n",
    "        ax.imshow(generatedImage[cnt],interpolation='nearest')\n",
    "        ax.axis('off')\n",
    "        ax.set_aspect('equal')\n",
    "        cnt+=1\n",
    "    fig.subplots_adjust(wspace=.004, hspace=.02)\n",
    "\n",
    "    if folder:\n",
    "        path = 'results/'+folder+'/samples'\n",
    "        if not os.path.exists('results'):\n",
    "            os.mkdir('results')\n",
    "        if not os.path.exists('results/'+folder):\n",
    "            os.mkdir('results/'+folder)\n",
    "        if not os.path.exists(path):\n",
    "            os.mkdir(path)\n",
    "        step = \"\"\n",
    "        if i:\n",
    "            step = '_'+str(i)\n",
    "        fig.savefig(path+'/epoch_%d%s.png' % (epoch, step))\n",
    "        plt.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/aray/pathology_gan\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(array([13054., 13160.]),\n",
       " array([0. , 0.5, 1. ], dtype=float32),\n",
       " <a list of 2 Patch objects>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYcAAAD8CAYAAACcjGjIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAEgNJREFUeJzt3H2MnedZ5/Hvj5gUCqVOm9moaxvsVQ27bnZXzY7SoEos1ChxAoojUSpHsHGLhSUILAsISJY/vGqJ1IiFQERf8BJvnarbJGRZYtGUrJWminaF00xICXkhZDZpG5u0GWon7G7UFpdr/zh3uqe+ZzKTOcdzPJ7vRxrN81zP/Zznuj22f/O8nJOqQpKkYd8y6QYkSWcew0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEmddZNuYLnOP//82rx586TbkKRV5aGHHvrbqppabNyqDYfNmzczMzMz6TYkaVVJ8vmljPOykiSpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySps2rfIS2dqTZf94lJt6Cz2Ofe/yMrcpw1GQ7+45WkV+ZlJUlSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUWDYckB5I8n+TRodpvJvmrJI8k+W9J1g9tuz7JbJInk1w2VN/RarNJrhuqb0nyQKvfnuTccU5QkvTqLeXM4SPAjlNqh4ELq+pfAH8NXA+QZBuwC3hL2+eDSc5Jcg7wAeByYBtwdRsLcCNwU1W9GTgB7BlpRpKkkS0aDlV1P3D8lNp/r6qTbfUIsLEt7wRuq6qvVtUzwCxwcfuaraqnq+prwG3AziQB3gHc2fY/CFw14pwkSSMaxz2HnwI+2ZY3AM8ObTvaagvV3wi8MBQ0L9clSRM0Ujgk+XXgJPCx8bSz6PH2JplJMjM3N7cSh5SkNWnZ4ZDk3cCPAj9RVdXKx4BNQ8M2ttpC9S8D65OsO6U+r6raX1XTVTU9NTW13NYlSYtYVjgk2QH8KnBlVb00tOkQsCvJa5JsAbYCnwEeBLa2J5POZXDT+lALlfuAd7b9dwN3LW8qkqRxWcqjrB8H/gz4viRHk+wBfg94HXA4yWeTfBigqh4D7gAeB/4UuLaqvt7uKfwccA/wBHBHGwvwa8AvJZllcA/ilrHOUJL0qq1bbEBVXT1PecH/wKvqBuCGeep3A3fPU3+awdNMkqQzhO+QliR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1Fg2HJAeSPJ/k0aHaG5IcTvJU+35eqyfJzUlmkzyS5KKhfXa38U8l2T1U/1dJ/rLtc3OSjHuSkqRXZylnDh8BdpxSuw64t6q2Ave2dYDLga3tay/wIRiECbAPeBtwMbDv5UBpY356aL9TjyVJWmGLhkNV3Q8cP6W8EzjYlg8CVw3Vb62BI8D6JG8CLgMOV9XxqjoBHAZ2tG3fVVVHqqqAW4deS5I0Icu953BBVT3Xlr8IXNCWNwDPDo072mqvVD86T31eSfYmmUkyMzc3t8zWJUmLGfmGdPuNv8bQy1KOtb+qpqtqempqaiUOKUlr0nLD4UvtkhDt+/OtfgzYNDRuY6u9Un3jPHVJ0gQtNxwOAS8/cbQbuGuofk17aukS4MV2+eke4NIk57Ub0ZcC97Rtf5fkkvaU0jVDryVJmpB1iw1I8nHgB4Hzkxxl8NTR+4E7kuwBPg+8qw2/G7gCmAVeAt4DUFXHk7wPeLCNe29VvXyT+2cZPBH17cAn25ckaYIWDYequnqBTdvnGVvAtQu8zgHgwDz1GeDCxfqQJK0c3yEtSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkjuEgSeoYDpKkzkjhkOQXkzyW5NEkH0/ybUm2JHkgyWyS25Oc28a+pq3Ptu2bh17n+lZ/Msllo01JkjSqZYdDkg3AvwWmq+pC4BxgF3AjcFNVvRk4Aexpu+wBTrT6TW0cSba1/d4C7AA+mOSc5fYlSRrdqJeV1gHfnmQd8FrgOeAdwJ1t+0Hgqra8s63Ttm9Pkla/raq+WlXPALPAxSP2JUkawbLDoaqOAf8R+AKDUHgReAh4oapOtmFHgQ1teQPwbNv3ZBv/xuH6PPtIkiZglMtK5zH4rX8L8I+B72BwWei0SbI3yUySmbm5udN5KEla00a5rPTDwDNVNVdVfw/8EfB2YH27zASwETjWlo8BmwDa9tcDXx6uz7PPN6mq/VU1XVXTU1NTI7QuSXolo4TDF4BLkry23TvYDjwO3Ae8s43ZDdzVlg+1ddr2T1VVtfqu9jTTFmAr8JkR+pIkjWjd4kPmV1UPJLkT+HPgJPAwsB/4BHBbkt9otVvaLrcAH00yCxxn8IQSVfVYkjsYBMtJ4Nqq+vpy+5IkjW7Z4QBQVfuAfaeUn2aep42q6ivAjy/wOjcAN4zSiyRpfHyHtCSpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpYzhIkjqGgySpM1I4JFmf5M4kf5XkiSTfn+QNSQ4neap9P6+NTZKbk8wmeSTJRUOvs7uNfyrJ7lEnJUkazahnDr8L/GlV/VPgXwJPANcB91bVVuDetg5wObC1fe0FPgSQ5A3APuBtwMXAvpcDRZI0GcsOhySvB34AuAWgqr5WVS8AO4GDbdhB4Kq2vBO4tQaOAOuTvAm4DDhcVcer6gRwGNix3L4kSaMb5cxhCzAH/OckDyf5gyTfAVxQVc+1MV8ELmjLG4Bnh/Y/2moL1TtJ9iaZSTIzNzc3QuuSpFcySjisAy4CPlRVbwX+L///EhIAVVVAjXCMb1JV+6tquqqmp6amxvWykqRTjBIOR4GjVfVAW7+TQVh8qV0uon1/vm0/Bmwa2n9jqy1UlyRNyLLDoaq+CDyb5PtaaTvwOHAIePmJo93AXW35EHBNe2rpEuDFdvnpHuDSJOe1G9GXtpokaULWjbj/zwMfS3Iu8DTwHgaBc0eSPcDngXe1sXcDVwCzwEttLFV1PMn7gAfbuPdW1fER+5IkjWCkcKiqzwLT82zaPs/YAq5d4HUOAAdG6UWSND6+Q1qS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEkdw0GS1DEcJEmdkcMhyTlJHk7yJ219S5IHkswmuT3Jua3+mrY+27ZvHnqN61v9ySSXjdqTJGk04zhz+AXgiaH1G4GbqurNwAlgT6vvAU60+k1tHEm2AbuAtwA7gA8mOWcMfUmSlmmkcEiyEfgR4A/aeoB3AHe2IQeBq9ryzrZO2769jd8J3FZVX62qZ4BZ4OJR+pIkjWbUM4ffAX4V+Ie2/kbghao62daPAhva8gbgWYC2/cU2/hv1efaRJE3AssMhyY8Cz1fVQ2PsZ7Fj7k0yk2Rmbm5upQ4rSWvOKGcObweuTPI54DYGl5N+F1ifZF0bsxE41paPAZsA2vbXA18ers+zzzepqv1VNV1V01NTUyO0Lkl6JcsOh6q6vqo2VtVmBjeUP1VVPwHcB7yzDdsN3NWWD7V12vZPVVW1+q72NNMWYCvwmeX2JUka3brFh7xqvwbcluQ3gIeBW1r9FuCjSWaB4wwChap6LMkdwOPASeDaqvr6aehLkrREYwmHqvo08Om2/DTzPG1UVV8BfnyB/W8AbhhHL5Kk0fkOaUlSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSx3CQJHUMB0lSZ9nhkGRTkvuSPJ7ksSS/0OpvSHI4yVPt+3mtniQ3J5lN8kiSi4Zea3cb/1SS3aNPS5I0ilHOHE4Cv1xV24BLgGuTbAOuA+6tqq3AvW0d4HJga/vaC3wIBmEC7APeBlwM7Hs5UCRJk7HscKiq56rqz9vy/waeADYAO4GDbdhB4Kq2vBO4tQaOAOuTvAm4DDhcVcer6gRwGNix3L4kSaMbyz2HJJuBtwIPABdU1XNt0xeBC9ryBuDZod2OttpC9fmOszfJTJKZubm5cbQuSZrHyOGQ5DuB/wr8u6r6u+FtVVVAjXqModfbX1XTVTU9NTU1rpeVJJ1ipHBI8q0MguFjVfVHrfyldrmI9v35Vj8GbBrafWOrLVSXJE3IKE8rBbgFeKKqfnto0yHg5SeOdgN3DdWvaU8tXQK82C4/3QNcmuS8diP60laTJE3IuhH2fTvwb4C/TPLZVvv3wPuBO5LsAT4PvKttuxu4ApgFXgLeA1BVx5O8D3iwjXtvVR0foS9J0oiWHQ5V9T+ALLB5+zzjC7h2gdc6ABxYbi+SpPHyHdKSpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqGA6SpI7hIEnqnDHhkGRHkieTzCa5btL9SNJadkaEQ5JzgA8AlwPbgKuTbJtsV5K0dp0R4QBcDMxW1dNV9TXgNmDnhHuSpDXrTAmHDcCzQ+tHW02SNAHrJt3Aq5FkL7C3rf6fJE8u86XOB/52PF2tGs55bVhrc15r8yU3jjzn71nKoDMlHI4Bm4bWN7baN6mq/cD+UQ+WZKaqpkd9ndXEOa8Na23Oa22+sHJzPlMuKz0IbE2yJcm5wC7g0IR7kqQ164w4c6iqk0l+DrgHOAc4UFWPTbgtSVqzzohwAKiqu4G7V+hwI1+aWoWc89qw1ua81uYLKzTnVNVKHEeStIqcKfccJElnkLM6HBb7SI4kr0lye9v+QJLNK9/l+Cxhvr+U5PEkjyS5N8mSHmk7ky31Y1eS/FiSSrLqn2xZypyTvKv9rB9L8l9WusdxW8Lf7e9Ocl+Sh9vf7ysm0ee4JDmQ5Pkkjy6wPUlubn8ejyS5aOxNVNVZ+cXgxvb/Av4JcC7wF8C2U8b8LPDhtrwLuH3SfZ/m+f4Q8Nq2/DOreb5LnXMb9zrgfuAIMD3pvlfg57wVeBg4r63/o0n3vQJz3g/8TFveBnxu0n2POOcfAC4CHl1g+xXAJ4EAlwAPjLuHs/nMYSkfybETONiW7wS2J8kK9jhOi863qu6rqpfa6hEG7ydZzZb6sSvvA24EvrKSzZ0mS5nzTwMfqKoTAFX1/Ar3OG5LmXMB39WWXw/8zQr2N3ZVdT9w/BWG7ARurYEjwPokbxpnD2dzOCzlIzm+MaaqTgIvAm9cke7G79V+BMkeBr95rGaLzrmdbm+qqk+sZGOn0VJ+zt8LfG+S/5nkSJIdK9bd6bGUOf8H4CeTHGXw1OPPr0xrE3PaP3LojHmUVSsnyU8C08C/nnQvp1OSbwF+G3j3hFtZaesYXFr6QQZnh/cn+edV9cJEuzq9rgY+UlW/leT7gY8mubCq/mHSja1WZ/OZw1I+kuMbY5KsY3A6+uUV6W78lvQRJEl+GPh14Mqq+uoK9Xa6LDbn1wEXAp9O8jkG12YPrfKb0kv5OR8FDlXV31fVM8BfMwiL1Wopc94D3AFQVX8GfBuDz106Wy3p3/sozuZwWMpHchwCdrfldwKfqna3ZxVadL5J3gr8PoNgWO3XoWGROVfVi1V1flVtrqrNDO6zXFlVM5NpdyyW8vf6jxmcNZDkfAaXmZ5eySbHbClz/gKwHSDJP2MQDnMr2uXKOgRc055augR4saqeG+cBztrLSrXAR3IkeS8wU1WHgFsYnH7OMrj5s2tyHY9mifP9TeA7gT9s992/UFVXTqzpES1xzmeVJc75HuDSJI8DXwd+papW6xnxUuf8y8B/SvKLDG5Ov3sV/6JHko8zCPjz232UfcC3AlTVhxncV7kCmAVeAt4z9h5W8Z+fJOk0OZsvK0mSlslwkCR1DAdJUsdwkCR1DAdJUsdwkCR1DAdJUsdwkCR1/h9NJUWDkw52+AAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%cd \"~/pathology_gan\"\n",
    "\n",
    "x_train, y_train, x_test, y_test, x_valid, y_valid = load_data()\n",
    "\n",
    "x_train = torch.from_numpy(np.moveaxis(x_train.astype(np.float32), -1, 1))\n",
    "x_test  = torch.from_numpy(np.moveaxis(x_test.astype(np.float32), -1, 1))\n",
    "x_valid  = torch.from_numpy(np.moveaxis(x_valid.astype(np.float32), -1, 1))\n",
    "\n",
    "y_train = torch.from_numpy(y_train).float() \n",
    "y_test = torch.from_numpy(y_test).float() \n",
    "y_valid = torch.from_numpy(y_valid).float() \n",
    "\n",
    "# X = torch.from_numpy(np.moveaxis(np.concatenate([x_train, x_test, x_valid]).astype(np.float32), -1, 1))\n",
    "# y = torch.from_numpy(np.concatenate([y_train, y_test, y_valid]))\n",
    "\n",
    "# trainloader = DataLoader(TensorDataset(x_train, y_train), batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\n",
    "validloader = DataLoader(TensorDataset(x_valid, y_valid), batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\n",
    "\n",
    "percent = int(x_train.shape[0]*.1)\n",
    "np.random.seed(17)\n",
    "idx_small = np.random.choice(range(x_train.shape[0]), percent, replace=False)\n",
    "\n",
    "x_train_small = x_train[idx_small]\n",
    "y_train_small = y_train[idx_small]\n",
    "trainloader = DataLoader(TensorDataset(x_train_small, y_train_small), batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\n",
    "\n",
    "plt.hist(y_train_small.numpy(), bins=2)\n",
    "# print(percent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_train_small = y_train_small.type(torch.LongTensor)\n",
    "y_valid=y_valid.type(torch.LongTensor)\n",
    "trainloader = DataLoader(TensorDataset(x_train_small, y_train_small), batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\n",
    "validloader = DataLoader(TensorDataset(x_valid, y_valid), batch_size=128, shuffle=True, num_workers=0, pin_memory=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# noise = np.random.normal(0,1,size=(16,100))\n",
    "# imgs = saved_actor.predict(noise)\n",
    "# plot_samples(imgs)\n",
    "\n",
    "# imgs_real = .5+.5*np.moveaxis(x_train[np.random.choice(range(x_train.shape[0]), 12*6, replace=False)], -1, 1)\n",
    "# plot_samples(x_train[np.random.choice(range(x_train.shape[0]), 12*6, replace=False)])\n",
    "# print(imgs_real.shape)\n",
    "# plt.figure(figsize=(24,12))\n",
    "# plt.imshow(np.transpose(make_grid(torch.from_numpy(imgs_real), nrow=12, ).numpy(), (1, 2, 0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, w, h, c, latent_dim):\n",
    "        super(Generator, self).__init__()\n",
    "        \n",
    "        self.w = w\n",
    "        self.h = h\n",
    "        self.c = c\n",
    "        self.latent_dim = latent_dim\n",
    "        \n",
    "        self.input = nn.Sequential(\n",
    "            nn.Linear(latent_dim, 32 * 16 * 16),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(32 * 16 * 16),\n",
    "        )\n",
    "        \n",
    "        self.deconv = nn.Sequential(\n",
    "            nn.Upsample(size=[32, 32], mode='nearest'),\n",
    "            nn.Conv2d(32, 64, 3, stride=1, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(64),\n",
    "            \n",
    "            nn.Upsample(size=[64, 64], mode='nearest'),\n",
    "            nn.Conv2d(64, 48, 4, stride=2, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(48),\n",
    "            \n",
    "            nn.Upsample(size=[64, 64], mode='nearest'),\n",
    "            nn.Conv2d(48, 32, 3, stride=1, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(32),\n",
    "\n",
    "            nn.Upsample(size=[128, 128], mode='nearest'),\n",
    "            nn.Conv2d(32, 16, 4, stride=2, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(16),\n",
    "\n",
    "            nn.Conv2d(16, 8, 3, stride=1, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(8),\n",
    "\n",
    "            nn.Upsample(size=[128, 128], mode='nearest'),\n",
    "            nn.Conv2d(8, 16, 4, stride=2, padding=1),\n",
    "            nn.ReLU(True),\n",
    "            nn.BatchNorm2d(16),\n",
    "            \n",
    "            nn.Conv2d(16, c, 3, stride=1, padding=1),\n",
    "            nn.Tanh()            \n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        output = self.input(x)\n",
    "        # print(output.shape)\n",
    "        output = output.view(-1, 32, 16, 16)\n",
    "        # print(output.shape)\n",
    "        return self.deconv(output) #.view(-1, self.w, self.h, self.c)\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, h, w, c):\n",
    "        super(Discriminator, self).__init__()\n",
    "        \n",
    "        self.w = w\n",
    "        self.h = h\n",
    "        self.c = c\n",
    "        n_filters = 32\n",
    "\n",
    "        self.conv = nn.Sequential(\n",
    "            weight_norm(nn.Conv2d(c, n_filters, 4, stride=2, padding=1)),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Dropout2d(p=0.5),\n",
    "            \n",
    "            weight_norm(nn.Conv2d(n_filters, 2*n_filters, 4, stride=2, padding=1)),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Dropout2d(p=0.5),\n",
    "                        \n",
    "            weight_norm(nn.Conv2d(2*n_filters, 4*n_filters, 4, stride=2, padding=1)),\n",
    "            nn.LeakyReLU(0.2),\n",
    "            nn.Flatten(),\n",
    "            nn.Dropout(p=0.5),\n",
    "       )\n",
    "        self.fc = weight_norm(nn.Linear(4*n_filters*int(w/2**3)*int(h/2**3), 2, bias=True))\n",
    "\n",
    "    def forward(self, x, clf=True, dropout=True):\n",
    "        flatten = self.conv(x)\n",
    "#         if dropout:\n",
    "#             flatten = F.dropout(flatten, p=0.5)\n",
    "        \n",
    "        out = self.fc(flatten)\n",
    "        \n",
    "        if clf:\n",
    "            return F.softmax(out, dim=1)\n",
    "\n",
    "#         expsum = torch.exp(out).sum(dim=1)\n",
    "#         out = expsum/(expsum+1)\n",
    "        return out, flatten\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SSL_WGAN:\n",
    "    def __init__(self, w, h, c, model_name, latent_dim=100):\n",
    "        self.model_name = model_name\n",
    "        self.latent_dim = latent_dim \n",
    "        self.lambda_gp = 10\n",
    "        self.lambda_ct = 2\n",
    "        self.d_iterations = 5\n",
    "        self.print_every = 10\n",
    "        # CT multiplier\n",
    "        self.M = .1\n",
    "        self.use_cuda = True\n",
    "        self.D = Discriminator(w, h, c)\n",
    "        self.G = Generator(w, h, c, latent_dim)\n",
    "\n",
    "        lr = 1e-4\n",
    "        betas = (.9, .99)\n",
    "\n",
    "        self.Clf_opt = optim.Adam(self.D.parameters(), lr=lr)\n",
    "        self.Clf_criterion = torch.nn.CrossEntropyLoss()\n",
    "        \n",
    "        self.D_opt = optim.Adam(self.D.parameters(), lr=lr)\n",
    "        self.G_opt = optim.Adam(self.G.parameters(), lr=lr)\n",
    "\n",
    "\n",
    "        if self.use_cuda:\n",
    "            self.D = self.D.cuda()\n",
    "            self.G = self.G.cuda()\n",
    "\n",
    "        if not os.path.exists('results/'+model_name):\n",
    "            os.mkdir('results/'+model_name)\n",
    "\n",
    "    def train(self, train_loader, validation_loader, X_unlabled, epochs, save_training_gif=True):\n",
    "        writer = SummaryWriter('/home/aray/runs/'+self.model_name)  \n",
    "        if save_training_gif:\n",
    "            # Fix latents to see how image generation improves during training\n",
    "            self.fixed_latents = torch.randn((12*6, self.latent_dim))\n",
    "            if self.use_cuda:\n",
    "                self.fixed_latents = self.fixed_latents.cuda()\n",
    "            self.training_progress_images = []\n",
    "\n",
    "        self.stats = {\n",
    "            'clf_loss': [],\n",
    "            'clf_acc': [],\n",
    "            'clf_loss_val': [],\n",
    "            'clf_acc_val': [],\n",
    "            'g_loss': [],\n",
    "            'd_loss': [],\n",
    "            'd_loss_real': [],\n",
    "            'd_loss_fake': [],\n",
    "            'gp': [],\n",
    "            'ct': [],\n",
    "            \n",
    "        }\n",
    "\n",
    "        for epoch in range(epochs):\n",
    "            clf_loss = []\n",
    "            clf_acc = []\n",
    "            g_loss = []\n",
    "            d_loss = []\n",
    "            d_loss_fake = []\n",
    "            d_loss_real = []\n",
    "            gradient_penalty = []\n",
    "            consistency_term = []\n",
    "\n",
    "#             for i, data in tqdm(enumerate(train_loader), desc=\"epoch \"+str(epoch)):\n",
    "            for i, data in enumerate(train_loader):\n",
    "#                 if i%((n_samples//128)//100)==0:\n",
    "#                     print(\".\", end=\"\", flush=True)\n",
    "                X = data[0]\n",
    "                y = data[1].view(-1)\n",
    "                batch_size = X.shape[0]\n",
    "                idx_u = np.random.choice(X_unlabled.shape[0], batch_size, replace=False)\n",
    "                X_u = X_unlabled[idx_u]\n",
    "                if self.use_cuda:\n",
    "                    X_u = X_u.cuda()\n",
    "                    X = X.cuda()\n",
    "                    y = y.cuda()\n",
    "                    \n",
    "                loss, acc = self._train_Clf(X, y)\n",
    "                clf_loss.append(loss)\n",
    "                clf_acc.append(acc)                                \n",
    "                    \n",
    "                ct, loss = self._train_D(X_u)\n",
    "#                 d_loss_real.append(r)\n",
    "#                 d_loss_fake.append(g)\n",
    "#                 gradient_penalty.append(gp)\n",
    "                consistency_term.append(ct)\n",
    "                d_loss.append(loss)\n",
    "                \n",
    "                # Only update generator every |d_iterations| iterations\n",
    "#                 if i % self.d_iterations == 0:\n",
    "                g_loss.append(self._train_G1(X_u))\n",
    "            \n",
    "            if epoch % 10 == 0:\n",
    "                acc_val = self._eval_Clf(validation_loader)\n",
    "                clf_loss_m = sum(clf_loss)/len(clf_loss)\n",
    "                clf_acc_m = sum(clf_acc)/len(clf_acc)     \n",
    "                g_loss_m = sum(g_loss)/len(g_loss)\n",
    "                d_loss_m = sum(d_loss)/len(d_loss)\n",
    "#                 real_m = sum(d_loss_real)/len(d_loss_real)\n",
    "#                 fake_m = sum(d_loss_fake)/len(d_loss_fake)\n",
    "#                 gp_m = sum(gradient_penalty)/len(gradient_penalty)\n",
    "                ct_m = sum(consistency_term)/len(consistency_term)\n",
    "                self.stats['clf_loss'].append(clf_loss_m)\n",
    "                self.stats['clf_acc'].append(clf_acc_m)\n",
    "    #             self.stats['clf_loss_val'].append(loss_val)\n",
    "                self.stats['clf_acc_val'].append(acc_val)\n",
    "                self.stats['g_loss'].append(g_loss_m)\n",
    "                self.stats['d_loss'].append(d_loss_m)\n",
    "#                 self.stats['d_loss_real'].append(real_m)\n",
    "#                 self.stats['d_loss_fake'].append(fake_m)\n",
    "#                 self.stats['gp'].append(gp_m)\n",
    "                self.stats['ct'].append(ct_m)\n",
    "                print(\"Epoch: %d, G loss: %f\"%(epoch, g_loss_m))\n",
    "                print(\"Clf loss: %f, acc: %.3f, acc_val: %.3f\"%(clf_loss_m, clf_acc_m, acc_val))\n",
    "                print(\"D loss: %f; ct: %f\"%(d_loss_m, ct_m))\n",
    " \n",
    "                writer.add_scalar('wgan_weight_nomr_ct', ct_m, epoch)\n",
    "                writer.add_scalar('wgan_weight_nomr_d_loss', d_loss_m, epoch)\n",
    "                writer.add_scalar('wgan_weight_nomr_g_loss', g_loss_m, epoch)\n",
    "                writer.add_scalar('wgan_weight_nomr_clf_loss', clf_loss_m, epoch)\n",
    "                writer.add_scalar('wgan_weight_nomr_clf_acc', clf_acc_m, epoch)\n",
    "                writer.add_scalar('wgan_weight_nomr_acc_val', acc_val, epoch)\n",
    "\n",
    "#             if epoch % 50 == 0:\n",
    "#                 with torch.no_grad():\n",
    "#                     self.G.eval()\n",
    "# #                     self.D.eval()\n",
    "# #                     print(self.D(X_unlabled[:10].cuda()).view(-1))\n",
    "# #                     self.D.train()\n",
    "#                     imgs = self.G(torch.randn((16, self.latent_dim)).cuda()).cpu().numpy()\n",
    "#                     plot_samples(np.moveaxis(imgs, 1,-1), self.model_name, epoch)\n",
    "#                     self.G.train()\n",
    "\n",
    "            if  epoch % 10 == 0 and save_training_gif:\n",
    "                with torch.no_grad():\n",
    "                    self.G.eval()\n",
    "                    img_grid = vutils.make_grid(self.G(self.fixed_latents).cpu(), nrow=12).numpy()\n",
    "                    # (width, height, channels)\n",
    "                    img_grid = .5+.5*np.transpose(img_grid, (1, 2, 0))\n",
    "                    self.training_progress_images.append(img_grid)\n",
    "                    writer.add_image('wgan_weight_nomr_images', img_grid, epoch, dataformats='HWC')\n",
    "                    self.G.train()\n",
    "\n",
    "            # if i % self.print_every == 0:\n",
    "            #     print(\"Iteration {}\".format(i + 1))\n",
    "            #     print(\"D: {}\".format(self.losses['D'][-1]))\n",
    "            #     print(\"GP: {}\".format(self.losses['GP'][-1]))\n",
    "            #     print(\"Gradient norm: {}\".format(self.losses['gradient_norm'][-1]))\n",
    "            #     if self.num_steps > self.critic_iterations:\n",
    "            #         print(\"G: {}\".format(self.losses['G'][-1]))\n",
    "\n",
    "        if save_training_gif:\n",
    "            imageio.mimsave('results/'+self.model_name+'/training_{}_epochs.gif'.format(epochs), self.training_progress_images)\n",
    "\n",
    "    def _eval_Clf(self, validation_loader):\n",
    "        self.D.eval()        \n",
    "        with torch.no_grad():\n",
    "            acc = .0\n",
    "            for i, data in enumerate(validation_loader):\n",
    "                X = data[0]\n",
    "                y = data[1].view(-1)\n",
    "                if self.use_cuda:\n",
    "                    X = X.cuda()\n",
    "                    y = y.cuda()\n",
    "                predicted = torch.argmax(self.D(X), dim=1)\n",
    "                acc+=(predicted == y).sum()/float(predicted.shape[0])       \n",
    "        self.D.train()\n",
    "        return (acc/(i+1)).detach().item()\n",
    "\n",
    "    def _train_Clf(self, data, labels):\n",
    "        self.Clf_opt.zero_grad()\n",
    "        \n",
    "        predicted = self.D(data, clf=True, dropout=True)\n",
    "    \n",
    "        loss = self.Clf_criterion(predicted, labels)\n",
    "        \n",
    "        loss.backward()\n",
    "        self.Clf_opt.step()    \n",
    "      \n",
    "        acc = (torch.argmax(predicted.detach().cpu(), dim=1) == labels.detach().cpu()).sum()/float(predicted.shape[0])\n",
    "\n",
    "        return loss.detach().item(), acc\n",
    "        \n",
    "#     def _D_loss(self, out):\n",
    "#         expsum = torch.exp(out).sum(dim=1)\n",
    "#         return expsum/(expsum+1)    \n",
    "    \n",
    "    def _train_D(self, data):\n",
    "        self.D_opt.zero_grad()\n",
    "\n",
    "        batch_size = data.shape[0]\n",
    "        generated_data = self.sample_generator(batch_size)\n",
    "        \n",
    "#         l_lab = output_before_softmax_lab[T.arange(args.batch_size),labels] \n",
    "#         l_unl = nn.log_sum_exp(output_before_softmax_unl) \n",
    "#         l_gen = nn.log_sum_exp(output_before_softmax_gen)\n",
    "#         loss_lab = -T.mean(l_lab) + T.mean(T.mean(nn.log_sum_exp(output_before_softmax_lab)))\n",
    "\n",
    "#         loss_comp = T.mean(lasagne.objectives.squared_error(T.nnet.softmax(output_before_softmax_unl),T.nnet.softmax(output_before_softmax_unl2)))\n",
    "#         loss_comp_ = T.mean(lasagne.objectives.squared_error(output_before_softmax_unl_,output_before_softmax_unl2_))\n",
    "\n",
    "#         loss_unl = 0.05*loss_comp_ + 0.5*loss_comp -0.5*T.mean(l_unl) + 0.5*T.mean(T.nnet.softplus(l_unl)) -0.5*np.log(1) + 0.5*T.mean(T.nnet.softplus(l_gen))  \n",
    "\n",
    "#         disc_param_updates = nn.adam_updates(disc_params, loss_lab + args.unlabeled_weight*loss_unl, lr=lr, mom1=0.5)\n",
    "        \n",
    "        real_fc1, real_flatten1 = self.D(data, clf=False)\n",
    "        real_fc2, real_flatten2 = self.D(data, clf=False)\n",
    "        gen_fc, _ = self.D(generated_data, clf=False)\n",
    "\n",
    "#         ct1 = ((torch.softmax(real_fc1, dim=1) - torch.softmax(real_fc2, dim=1))**2).mean()\n",
    "#         ct2 = ((torch.softmax(real_flatten1, dim=1) - torch.softmax(real_flatten2, dim=1))**2).mean()\n",
    "    \n",
    "        ct1 = (torch.softmax(real_fc1, dim=1) - torch.softmax(real_fc2, dim=1))**2\n",
    "#         ct2 = (real_flatten1 - real_flatten2)**2\n",
    "        \n",
    "#         ct = 0.1(*0.5*ct1 + 0.05*ct2.mean())\n",
    "        ct = 0.1*ct1.mean()\n",
    "\n",
    "        real_lse = torch.logsumexp(real_fc1, dim=1)\n",
    "        \n",
    "        real_mean = F.softplus(real_lse).mean()\n",
    "        gen_mean = F.softplus(torch.logsumexp(gen_fc, dim=1)).mean()\n",
    "\n",
    "        d_loss = 0.5*(ct+real_mean + gen_mean - real_lse.mean())\n",
    "        \n",
    "        d_loss.backward()\n",
    "        self.D_opt.step()\n",
    "\n",
    "\n",
    "        return ct.detach().item(), d_loss.detach().item()\n",
    "\n",
    "    \n",
    "    def _train_G1(self, real_data):\n",
    "        for p in self.D.parameters():\n",
    "            p.requires_grad = False\n",
    "            \n",
    "        self.G_opt.zero_grad()\n",
    "\n",
    "        batch_size = real_data.shape[0]\n",
    "        generated_data = self.sample_generator(batch_size)\n",
    "\n",
    "        _, d_real_output = self.D(real_data, clf=False, dropout=False)\n",
    "        _, d_gen_output = self.D(generated_data, clf=False, dropout=False)\n",
    "\n",
    "        g_loss = ((d_gen_output.mean(dim=0)-d_real_output.mean(dim=0))**2).mean()\n",
    "        g_loss.backward()\n",
    "        self.G_opt.step()\n",
    "        \n",
    "        for p in self.D.parameters():\n",
    "            p.requires_grad = True\n",
    "\n",
    "        return g_loss.detach().item()    \n",
    "\n",
    "    def sample_generator(self, num_samples):\n",
    "        latent_samples = torch.randn((num_samples, self.latent_dim), requires_grad=True)\n",
    "        if self.use_cuda:\n",
    "            latent_samples = latent_samples.cuda()\n",
    "        generated_data = self.G(latent_samples)\n",
    "        return generated_data\n",
    "\n",
    "    def sample(self, num_samples):\n",
    "        generated_data = self.sample_generator(num_samples)\n",
    "        # Remove color channel\n",
    "        return generated_data.data.cpu().numpy()[:, 0, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_stats(model):\n",
    "    with open('results/'+model.model_name+'/clf_acc.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['clf_acc'])\n",
    "\n",
    "    with open('results/'+model.model_name+'/clf_loss.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['clf_loss'])\n",
    "\n",
    "    with open('results/'+model.model_name+'/clf_acc_val.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['clf_acc_val'])\n",
    "        \n",
    "    with open('results/'+model.model_name+'/d_loss.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['d_loss'])\n",
    "\n",
    "    with open('results/'+model.model_name+'/g_loss.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['g_loss'])\n",
    "\n",
    "    with open('results/'+model.model_name+'/ct.csv', 'w', newline='') as myfile:\n",
    "        wr = csv.writer(myfile)\n",
    "        wr.writerow(model.stats['ct'])\n",
    "\n",
    "def save_plots(model):\n",
    "    path = 'results/'+model.model_name\n",
    "    \n",
    "    fig = plt.figure(figsize=(18,5))\n",
    "    plt.plot(wgan_ct.stats['ct'], label='CT')\n",
    "    plt.xlabel('Epochs x10')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "    fig.savefig(path+'/ct.png')\n",
    "\n",
    "    fig = plt.figure(figsize=(18,5))\n",
    "    plt.plot(wgan_ct.stats['d_loss'], label='D')\n",
    "    plt.plot(wgan_ct.stats['g_loss'], label='G')\n",
    "    plt.xlabel('Epochs x10')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "    fig.savefig(path+'/gan_loss.png')\n",
    "\n",
    "    fig = plt.figure(figsize=(18,5))\n",
    "    plt.plot(wgan_ct.stats['clf_loss'], label='Clf train')\n",
    "    plt.xlabel('Epochs x10')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.legend()\n",
    "    fig.savefig(path+'/clf_train_loss.png')\n",
    "\n",
    "    fig = plt.figure(figsize=(18,5))\n",
    "    acc_max = np.max(wgan_ct.stats['clf_acc_val'])\n",
    "    acc_max_e = np.argmax(wgan_ct.stats['clf_acc_val'])\n",
    "    plt.axhline(y=acc_max, color='k', linestyle='-', alpha=0.2, linewidth=1)\n",
    "    plt.axvline(x=acc_max_e, color='k', linestyle='-', alpha=0.2, linewidth=1)\n",
    "    plt.plot(wgan_ct.stats['clf_acc'], label='Train')\n",
    "    plt.plot(wgan_ct.stats['clf_acc_val'], label='Validation')\n",
    "    plt.xlabel('Epochs x10')\n",
    "    plt.ylabel('Acc')\n",
    "    plt.legend()\n",
    "    fig.savefig(path+'/clf_acc.png')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, G loss: 0.008431\n",
      "Clf loss: 0.632586, acc: 0.639, acc_val: 0.736\n",
      "D loss: 0.416151; ct: 0.006047\n",
      "Epoch: 10, G loss: 0.004059\n",
      "Clf loss: 0.548845, acc: 0.752, acc_val: 0.765\n",
      "D loss: 0.428434; ct: 0.003698\n",
      "Epoch: 20, G loss: 0.006013\n",
      "Clf loss: 0.544301, acc: 0.756, acc_val: 0.766\n",
      "D loss: 0.399273; ct: 0.003922\n",
      "Epoch: 30, G loss: 0.004186\n",
      "Clf loss: 0.544524, acc: 0.758, acc_val: 0.770\n",
      "D loss: 0.467871; ct: 0.003925\n",
      "Epoch: 40, G loss: 0.003086\n",
      "Clf loss: 0.540375, acc: 0.765, acc_val: 0.784\n",
      "D loss: 0.495108; ct: 0.004237\n",
      "Epoch: 50, G loss: 0.002383\n",
      "Clf loss: 0.536181, acc: 0.769, acc_val: 0.791\n",
      "D loss: 0.475838; ct: 0.004495\n"
     ]
    }
   ],
   "source": [
    "wgan_ct = SSL_WGAN(64, 64, 3, 'ssl_wn_10')\n",
    "wgan_ct.train(trainloader, validloader, x_train, 500, x_train.shape[0])\n",
    "save_stats(wgan_ct)\n",
    "save_plots(wgan_ct)\n",
    "torch.save(wgan_ct.D.state_dict(), 'results/'+wgan_ct.model_name+'/D.pt')\n",
    "torch.save(wgan_ct.G.state_dict(), 'results/'+wgan_ct.model_name+'/G.pt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}