train_mnist.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import gzip\n",
    "import time\n",
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import ndimage\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from IPython.display import clear_output\n",
    "from datetime import datetime\n",
    "\n",
    "from lib.utils import SamplePool, make_seed, make_circle_masks\n",
    "from lib.utils import get_living_mask, get_sobel, softmax\n",
    "from lib.utils import mnist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images, train_labels, test_images, test_labels = mnist(\"./data/mnist\")\n",
    "map_size = (28,28)\n",
    "ALPHA_CHANNEL = 0\n",
    "\n",
    "print(train_images.shape)\n",
    "print(train_labels.shape)\n",
    "print(test_images.shape)\n",
    "print(test_labels.shape)\n",
    "\n",
    "print(np.min(train_images))\n",
    "print(np.max(train_images))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE_encoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, dim_out):\n",
    "        super(VAE_encoder, self).__init__()\n",
    "        self.conv_mu    = nn.Sequential(nn.Conv2d(1,8,3,padding=1), nn.ReLU(),\n",
    "                                        nn.Conv2d(8,8,3,padding=1), nn.ReLU(),\n",
    "                                        nn.MaxPool2d(2,stride=2),\n",
    "                                        nn.Conv2d(8,16,3,padding=1), nn.ReLU(),\n",
    "                                        nn.Conv2d(16,16,3,padding=1), nn.ReLU(),\n",
    "                                        nn.MaxPool2d(2,stride=2),\n",
    "                                        nn.Conv2d(16,32,3,padding=1), nn.ReLU(),\n",
    "                                        nn.Conv2d(32,64,5,stride=2,padding=2), nn.ReLU(),\n",
    "                                        nn.Conv2d(64,64,4))\n",
    "        self.lin_mu     = nn.Linear(64,dim_out)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x_ = torch.reshape(x, [-1,ALPHA_CHANNEL+1,map_size[0],map_size[1]])\n",
    "        c_mu = torch.reshape(self.conv_mu(x_), [-1,64])\n",
    "        return self.lin_mu(c_mu)\n",
    "    \n",
    "class VAE_decoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, dim_in, dim_out):\n",
    "        super(VAE_decoder, self).__init__()\n",
    "        fc0 = nn.Linear(dim_in,512)\n",
    "        fc1 = nn.Linear(512,512)\n",
    "        fc2 = nn.Linear(512,dim_out)\n",
    "        self.lin = nn.Sequential(fc0, nn.ReLU(), fc1, nn.ReLU(), fc2)\n",
    "        \n",
    "    def forward(self, z):\n",
    "        r = self.lin(z)\n",
    "        return r\n",
    "    \n",
    "class GNCAModel(nn.Module):\n",
    "\n",
    "    def __init__(self, channel_n, alpha_channel,\n",
    "                 fire_rate=0.5, calibration=1.0, device=torch.device(\"cpu\")):\n",
    "        super(GNCAModel, self).__init__()\n",
    "\n",
    "        self.device = device\n",
    "        self.channel_n = channel_n\n",
    "        self.alpha_channel = alpha_channel\n",
    "        \n",
    "        self.pool = torch.nn.MaxPool2d(kernel_size=3, padding=1, stride=1)\n",
    "\n",
    "        self.fire_rate = fire_rate\n",
    "        self.calibration = calibration\n",
    "        self.to(self.device)\n",
    "\n",
    "    def perceive(self, x, angle):\n",
    "\n",
    "        def _perceive_with(x, weight):\n",
    "            size = weight.shape[0]\n",
    "            padding = (size-1)/2\n",
    "            conv_weights = torch.from_numpy(weight.astype(np.float32)).to(self.device)\n",
    "            conv_weights = conv_weights.view(1,1,size,size).repeat(self.channel_n, 1, 1, 1)\n",
    "            return F.conv2d(x, conv_weights, padding=int(padding), groups=self.channel_n)\n",
    "\n",
    "        wa_1, wa_2 = get_sobel(3)\n",
    "        wa_3, wa_4 = get_sobel(7)\n",
    "        \n",
    "        wa_1/=np.sum(np.abs(wa_1))\n",
    "        wa_2/=np.sum(np.abs(wa_2))\n",
    "        wa_3/=np.sum(np.abs(wa_3))\n",
    "        wa_4/=np.sum(np.abs(wa_4))\n",
    "\n",
    "        y1 = _perceive_with(x, wa_1)\n",
    "        y2 = _perceive_with(x, wa_2)\n",
    "        y3 = _perceive_with(x, wa_3)\n",
    "        y4 = _perceive_with(x, wa_4)\n",
    "        y5 = self.pool(x)\n",
    "        y = torch.cat((x,y1,y2,y3,y4,y5),1)\n",
    "        return y\n",
    "    \n",
    "    def linear(self, x, w, b=None):\n",
    "        original_shape = x.size()\n",
    "        batch = x.size(0)\n",
    "        y = torch.reshape(x, [batch,-1,original_shape[-1]]).to(self.device)\n",
    "        if b is None:\n",
    "            y = torch.bmm(y, w)\n",
    "        else:\n",
    "            y = torch.bmm(y, w)+b\n",
    "        y = torch.reshape(y, list(original_shape[:-1])+[y.size(-1)])\n",
    "        return y\n",
    "\n",
    "    def update(self, x, params, fire_rate, angle):\n",
    "        # (channel_n*8, 256), (256,), (256, channel_n). The params should be trainable while training.\n",
    "        w0, b0, w1 = params\n",
    "        \n",
    "        x = x.transpose(1,3)\n",
    "        pre_life_mask = get_living_mask(x, self.alpha_channel, 3)\n",
    "\n",
    "        dx = self.perceive(x, angle)\n",
    "        dx = dx.transpose(1,3)\n",
    "        dx = self.linear(dx, w0, b0)\n",
    "        dx = F.relu(dx)\n",
    "        dx = self.linear(dx, w1)\n",
    "\n",
    "        if fire_rate is None:\n",
    "            fire_rate=self.fire_rate\n",
    "        stochastic = torch.rand([dx.size(0),dx.size(1),dx.size(2),1])>fire_rate\n",
    "        stochastic = stochastic.float().to(self.device)\n",
    "        dx = dx * stochastic\n",
    "        dx = dx.transpose(1,3)\n",
    "\n",
    "        x = x+dx\n",
    "\n",
    "        post_life_mask = get_living_mask(x, self.alpha_channel, 3)\n",
    "        life_mask = (pre_life_mask & post_life_mask).float()\n",
    "        \n",
    "        x = x * life_mask\n",
    "        x = x.transpose(1,3)\n",
    "        return x\n",
    "\n",
    "    def forward(self, x, params, steps, calibration_map=None, fire_rate=None, angle=0.0):\n",
    "        history = [x.detach().cpu().clamp(0.0, 1.0).numpy(),]\n",
    "        for step in range(steps):\n",
    "            x = self.update(x, params, fire_rate, angle)\n",
    "            if calibration_map is not None:\n",
    "                h = x[..., :(self.alpha_channel+1)]\n",
    "                t = calibration_map[..., :(self.alpha_channel+1)]\n",
    "                _delta = t*(h-1)\n",
    "                delta = _delta * self.calibration * (calibration_map!=0).float()\n",
    "                _x = x[..., :(self.alpha_channel+1)]-delta\n",
    "                x = torch.cat((_x,x[..., (self.alpha_channel+1):]), -1)\n",
    "            history.append(x.detach().cpu().clamp(0.0, 1.0).numpy())\n",
    "        return x, history\n",
    "    \n",
    "class model_VAE(nn.Module):\n",
    "    \n",
    "    def __init__(self, hidden_dim, channel_n, alpha_channel, hidden_channel_n,\n",
    "                 device=torch.device(\"cpu\")):\n",
    "        super(model_VAE, self).__init__()\n",
    "        self.channel_n = channel_n\n",
    "        self.hidden_channel_n = hidden_channel_n\n",
    "        self.alpha_channel = alpha_channel\n",
    "        self.eps = 1e-3\n",
    "        \n",
    "        self.encoder = VAE_encoder(hidden_dim)\n",
    "        self.decoder_w0 = VAE_decoder(hidden_dim, channel_n*6*self.hidden_channel_n)\n",
    "        self.decoder_b0 = VAE_decoder(hidden_dim, self.hidden_channel_n)\n",
    "        self.decoder_w1 = VAE_decoder(hidden_dim, self.hidden_channel_n*channel_n)\n",
    "        self.GNCA = GNCAModel(self.channel_n, self.alpha_channel, device=device)\n",
    "        \n",
    "        self.device = device\n",
    "        self.to(self.device)\n",
    "    \n",
    "    def encode(self, x):\n",
    "        z = self.encoder(x)\n",
    "        return z\n",
    "    \n",
    "    def decode(self, z):\n",
    "        w0 = self.decoder_w0(z)\n",
    "        b0 = self.decoder_b0(z)\n",
    "        w1 = self.decoder_w1(z)\n",
    "        params = (torch.reshape(w0,[-1,self.channel_n*6,self.hidden_channel_n]),\n",
    "                  torch.reshape(b0,[-1,1,self.hidden_channel_n]),\n",
    "                  torch.reshape(w1,[-1,self.hidden_channel_n,self.channel_n]))\n",
    "        return params\n",
    "    \n",
    "    def infer(self, x0, x, steps, calibration_map=None):\n",
    "        with torch.no_grad():\n",
    "            z = self.encode(x)\n",
    "            params = self.decode(z)\n",
    "            y, history = self.GNCA(x0, params, steps, calibration_map=calibration_map)\n",
    "            y = y[..., :(self.alpha_channel+1)].clamp(self.eps, 1.0-self.eps)\n",
    "        return y, history\n",
    "    \n",
    "    def train(self, x0, x, steps, beta, calibration_map=None):\n",
    "        z = self.encode(x)\n",
    "        params = self.decode(z)\n",
    "        y, _ = self.GNCA(x0, params, steps, calibration_map=calibration_map)\n",
    "        y = y[..., :(self.alpha_channel+1)].clamp(self.eps, 1.0-self.eps)\n",
    "        \n",
    "        target = torch.reshape(x, y.size())\n",
    "        marginal = -(target*torch.log(y)+(1-target)*torch.log(1-y))\n",
    "        marginal_likelihood = torch.mean(torch.sum(marginal,-1))\n",
    "        \n",
    "        l2 = torch.sum(torch.pow(z, 2))\n",
    "        loss = marginal_likelihood+beta*l2\n",
    "        \n",
    "        return y, loss, (marginal_likelihood.item(), l2.item())\n",
    "    \n",
    "def plot_loss(loss_log):\n",
    "    plt.figure(figsize=(10, 4))\n",
    "    plt.title('Loss history (log10)')\n",
    "    plt.plot(np.log10(loss_log), '.', alpha=0.1)\n",
    "    plt.show()\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = torch.device(\"cpu\")\n",
    "model_path = \"models/my_model.pth\"\n",
    "map_size = (28,28)\n",
    "init_coord = (map_size[0]//2, map_size[1]//2)\n",
    "\n",
    "HIDDEN = 32\n",
    "CHANNEL_N = 8\n",
    "HIDDEN_CHANNEL_N = 64\n",
    "\n",
    "LR = 1e-4\n",
    "BETAS = (0.9, 0.999)\n",
    "AE_BETA = 1e-5\n",
    "GAMMA = 1.0\n",
    "N_EPOCH = 400\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "N_STEPS = 64\n",
    "\n",
    "my_model = model_VAE(HIDDEN, CHANNEL_N, ALPHA_CHANNEL, HIDDEN_CHANNEL_N, device=DEVICE)\n",
    "# my_model.load_state_dict(torch.load(model_path))\n",
    "optimizer = optim.Adam(my_model.parameters(), lr=LR, betas=BETAS)\n",
    "scheduler = optim.lr_scheduler.ExponentialLR(optimizer, GAMMA)\n",
    "\n",
    "loss_log = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "for i_epoch in range(N_EPOCH):\n",
    "    # n_batch = int(np.ceil(train_images.shape[0]/BATCH_SIZE))\n",
    "    n_batch = 50\n",
    "    \n",
    "    seq = np.random.choice(train_images.shape[0], train_images.shape[0], replace=False)\n",
    "    for index in range(n_batch):\n",
    "        indices = seq[index*BATCH_SIZE:(index+1)*BATCH_SIZE]\n",
    "        x_np = train_images[indices].astype(np.float32)\n",
    "        x = torch.from_numpy(x_np).to(DEVICE)\n",
    "        \n",
    "        seed = make_seed(map_size, CHANNEL_N, np.arange(CHANNEL_N-ALPHA_CHANNEL)+ALPHA_CHANNEL, init_coord)\n",
    "        x0 = np.repeat(seed[None, ...], len(indices), 0)\n",
    "        x0 = torch.from_numpy(x0.astype(np.float32)).to(DEVICE)\n",
    "        \n",
    "        y, loss, (l_margin, l_kl) = my_model.train(x0, x, N_STEPS, AE_BETA)\n",
    "        print(\"Loss =\", loss.item())\n",
    "        loss_log.append(loss.item())\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    scheduler.step()\n",
    "        \n",
    "    y = y.detach().cpu().numpy()\n",
    "    n_show = min(8,len(indices))\n",
    "\n",
    "    clear_output()\n",
    "\n",
    "    plot_loss(loss_log)\n",
    "    torch.save(my_model.state_dict(), model_path)\n",
    "    time_cost = (time.time()-start_time)/60\n",
    "    print(\"Epoch:\", i_epoch+1, \", Time Cost:\", time_cost, \"min\")\n",
    "    print(\"model saved in\", model_path)\n",
    "    plt.figure(figsize=(18,2))\n",
    "    for i in range(n_show):\n",
    "        plt.subplot(2,n_show,i+1)\n",
    "        plt.imshow(x_np[i].reshape(map_size), cmap=plt.cm.gray, vmin=0, vmax=1)\n",
    "        plt.axis('off')\n",
    "    for i in range(n_show):\n",
    "        plt.subplot(2,n_show,i+n_show+1)\n",
    "        plt.imshow(y[i,...,0], cmap=plt.cm.gray, vmin=0, vmax=1)\n",
    "        plt.axis('off')\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:py37_torch]",
   "language": "python",
   "name": "conda-env-py37_torch-py"
  },
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}