1-MNIST_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the 2L-SPC on MNIST database\n",
    "https://arxiv.org/abs/2002.00892"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from  SPC_2L.DataTools import DataBase\n",
    "from SPC_2L.Network import LayerPC, Network\n",
    "from SPC_2L.Coding import ML_Lasso,ML_FISTA\n",
    "from SPC_2L.DataTools import DataBase, gaussian_kernel\n",
    "from SPC_2L.Monitor import Monitor\n",
    "from SPC_2L.Optimizers import mySGD, myAdam\n",
    "import torch.nn.functional as f\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "import time\n",
    "import tensorboardX\n",
    "from SPC_2L.DataTools import LCN, whitening, z_score, mask, to_cuda, norm\n",
    "from torchvision.utils import make_grid\n",
    "import numpy as np\n",
    "from tensorboardX import SummaryWriter\n",
    "import pickle\n",
    "from torchvision.transforms import ToTensor, Compose\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.optim as optim\n",
    "from SDPC.Monitor import Monitor\n",
    "import math\n",
    "from torchvision.datasets import MNIST\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tools, parameters and database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(img):\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## Database \n",
    "data_path = '../../DataSet/MNIST/'\n",
    "\n",
    "transform = Compose([ToTensor(),\n",
    "                    to_cuda(),\n",
    "                    whitening((28,28),f_0=0.6),\n",
    "                    LCN(kernel_size=3,sigma=0.5,rgb=False),\n",
    "                    z_score()])\n",
    "\n",
    "dataset = MNIST(data_path, transform=transform, train=True, download=False)\n",
    "\n",
    "DataBase =  DataLoader(dataset,batch_size=32,shuffle=True,drop_last=True)\n",
    "\n",
    "## Gaussian masks for the dictionaries\n",
    "mask_g = [gaussian_kernel((64,3,8,8), sigma=30), gaussian_kernel((128,64,8,8), sigma=30)]\n",
    "\n",
    "## Setting parameters\n",
    "l_r = [5e-2,1e-1]#### dictionaries learning rate [Layer1, Layer2]\n",
    "l_rv = [1e-3,1e-3]#### normalizer learning rate [Layer1, Layer2]\n",
    "l = [0.20,0.30]#### Sparsity parameters [Layer1, Layer2]\n",
    "b=0 #### Feedback strength parameter. b=0 --> Hila, b=1 --> 2L-SPC\n",
    "v_i=[2,2] #### Initial normalizer value [Layer1, Layer2]\n",
    "nb_epoch = 100 #### number of training epochs\n",
    "\n",
    "Use_tb = True #### Use to activate tensorboard monitoring\n",
    "save = False #### Use to run the entire simulation : TAKE HOURS. Use False to load previous simulation\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and saving the network (b=0 for Hi-La, b=1 for 2L-SPC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NETWORK STRUCTURE : \n",
      " Input : (32, 1, 28, 28)\n",
      " Layer 1 : [32, 32, 12, 12]\n",
      " Layer 2 : [32, 64, 8, 8]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "## Definition of the layers, network and sparse coding algorithm\n",
    "layer = [LayerPC((32, 1, 5, 5), stride=2, b=b, v=v_i[0], v_size=32 ,out_pad=1),\n",
    "        LayerPC((64, 32, 5, 5), stride=1, b=b, v=v_i[1], v_size=64 ,out_pad=0)]\n",
    "\n",
    "Net = Network(layer, input_size=(32, 1, 28,28))\n",
    "Loss = ML_Lasso(Net, [l[0],l[1]])\n",
    "Pursuit = ML_FISTA(Net, Loss, max_iter=1000, th=5e-4, mode='eigen')\n",
    "\n",
    "## Optimizer initialization\n",
    "opt_dico = [None] * (Net.nb_layers + 1)\n",
    "for i in range(0, Net.nb_layers):\n",
    "    opt_dico[i] = mySGD([{'params': Net.layers[i].dico}], lr=l_r[i], momentum=0.9, normalize=True)\n",
    "\n",
    "opt_v = [myAdam([{'params': Net.layers[i].v}], lr=l_rv[i], normalize=False) \\\n",
    "    for i in range(Net.nb_layers)]\n",
    "\n",
    "\n",
    "L = [None] * (Net.nb_layers)\n",
    "L_v = [None] * (Net.nb_layers)\n",
    "reco = [None] * (Net.nb_layers)\n",
    "\n",
    "model_name = 'MNIST_[{0},{1}]_b={2}'.format(l[0],l[1],b)\n",
    "path = 'Savings/MNIST/' + model_name +'.pkl'\n",
    "if Use_tb : \n",
    "    nrows = [8,8,8,8,8,8,8]\n",
    "    writer = SummaryWriter('Savings/Log/' + model_name)\n",
    "    M = Monitor(Net, writer, n_row=nrows)\n",
    "\n",
    "k=0\n",
    "\n",
    "l2_loss = torch.zeros(2,nb_epoch*len(DataBase))\n",
    "l1_loss = torch.zeros(2,nb_epoch*len(DataBase))\n",
    "if save == True : \n",
    "    for e in range(nb_epoch):\n",
    "        for idx_batch, data in enumerate(DataBase):\n",
    "\n",
    "            batch = data[0].cuda()\n",
    "            gamma, it, Loss_G, delta = Pursuit.coding(batch)\n",
    "\n",
    "\n",
    "            for i in range(Net.nb_layers):\n",
    "                Net.layers[i].dico.requires_grad = True\n",
    "                L[i] = Loss.F(batch, gamma, i, do_feedback=False).div(batch.size()[0])  ## Unsupervised\n",
    "                L[i].backward()\n",
    "                Net.layers[i].dico.requires_grad = False\n",
    "                opt_dico[i].step()\n",
    "                opt_dico[i].zero_grad()\n",
    "                \n",
    "                l2_loss[i,k]= L[i].detach() \n",
    "                l1_loss[i,k] =  gamma[i].detach().sum().div(gamma[i].size(0))\n",
    "                \n",
    "\n",
    "            for i in range(Net.nb_layers):\n",
    "                Net.layers[i].v.requires_grad = True  # turn_on(i)\n",
    "                L_v[i] = Loss.F_v(batch, gamma, i).div(batch.size()[0])\n",
    "                L_v[i].backward()\n",
    "                Net.layers[i].v.requires_grad = False  # turn_off(i)\n",
    "                opt_v[i].step()  \n",
    "                opt_v[i].zero_grad()\n",
    "                \n",
    "            if Use_tb:\n",
    "                if (k%10) == 0:\n",
    "                    writer.add_scalar('FISTA_iterations', it, k)\n",
    "                    M.MonitorGamma(gamma, k, option=['NNZ', '%', 'Sum', 'V'])\n",
    "                    M.MonitorList(L, 'Loss_Dico', k)\n",
    "                    M.MonitorList(L_v, 'Loss_v', k)\n",
    "                    M.MonitorDicoBP(k)\n",
    "                    M.ComputeHisto(gamma)\n",
    "\n",
    "                if (k%100) == 0:\n",
    "                    reco = [None] * (Net.nb_layers)\n",
    "                    for i in range(Net.nb_layers-1,-1,-1):\n",
    "                        reco[i] = gamma[i]\n",
    "                        for j in range(i, -1, -1):\n",
    "                            reco[i] = Net.layers[j].backward(reco[i])\n",
    "                        reco_image = make_grid(reco[i],normalize=True,pad_value=1)\n",
    "                        writer.add_image('Reco/L{0}'.format(i),reco_image,k)\n",
    "\n",
    "            k += 1\n",
    "\n",
    "    output_exp = {'Net': Net,\n",
    "            'Loss': Loss,\n",
    "            'Pursuit': Pursuit,\n",
    "            'l2_loss': l2_loss,\n",
    "            'l1_loss': l1_loss    \n",
    "                 }\n",
    "    path = 'Savings/MNIST/' + model_name +'.pkl'\n",
    "    with open(path, 'wb') as file:\n",
    "        pickle.dump(output_exp, file, pickle.HIGHEST_PROTOCOL)\n",
    "    \n",
    "else :        \n",
    "    with open(path, 'rb') as file:\n",
    "        output_exp = pickle.load(file)\n",
    "\n"
   ]
  },
  {
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}