examples/Equilibrium models (RNN-style model on MNIST).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch as th\n",
    "import hypergrad as hg\n",
    "\n",
    "TARGET_DEVICE = th.device('cuda') if th.cuda.is_available() else th.device('cpu')"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Code for training a simple equilibrium network with \"RNN-style\" dynamics on a subset of the MNIST data.\n",
    "\n",
    "For more details refer to Section 3.2 of the paper\n",
    "\"On the iteration complexity of hypergradient computation\" (https://arxiv.org/abs/2006.16218)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# --------------------------------------------\n",
    "# UTILS\n",
    "# --------------------------------------------\n",
    "\n",
    "def to_numpy(tensor):\n",
    "    if isinstance(tensor, list):\n",
    "        return [to_numpy(v) for v in tensor]\n",
    "    else:\n",
    "        return tensor.detach().to(th.device('cpu')).numpy()\n",
    "\n",
    "\n",
    "def set_requires_grad(lst): [l.requires_grad_(True) for l in lst]\n",
    "\n",
    "\n",
    "def acc(preds, targets):\n",
    "    \"\"\"Computes the accuracy\"\"\"\n",
    "    return preds.argmax(dim=1).eq(targets).float().mean()\n",
    "\n",
    "\n",
    "class NamedLists(list):\n",
    "    def __init__(self, lst, names) -> None:\n",
    "        super().__init__(lst)\n",
    "        assert len(lst) == len(names)\n",
    "        self.names = names\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        if isinstance(i, str):\n",
    "            return self.__getattribute__(i)\n",
    "        else:\n",
    "            return super().__getitem__(i)\n",
    "\n",
    "\n",
    "class TVT(NamedLists):  # train val & test\n",
    "    def __init__(self, lst) -> None:\n",
    "        super().__init__(lst, ['train', 'val', 'test'])\n",
    "        self.train, self.val, self.test = lst\n",
    "\n",
    "\n",
    "class DT(NamedLists):  # data & targets\n",
    "    def __init__(self, lst) -> None:\n",
    "        super().__init__(lst, ['data', 'targets'])\n",
    "        self.data, self.targets = lst\n",
    "\n",
    "\n",
    "class LA(NamedLists):  # loss and accuracy\n",
    "    def __init__(self, lst):\n",
    "        super().__init__(lst, ['loss', 'acc'])\n",
    "        self.loss, self.acc = lst\n",
    "\n",
    "\n",
    "def load_mnist(seed=0, num_train=50000, num_valid=10000):\n",
    "    \"\"\"Load MNIST dataset with given number of training and validation examples\"\"\"\n",
    "    from torchvision import datasets\n",
    "    rnd = np.random.RandomState(seed)\n",
    "    mnist_train = datasets.MNIST('../data', download=True, train=True)\n",
    "    train_indices = rnd.permutation(list(range(60000)))\n",
    "    dta, targets = mnist_train.data, mnist_train.targets\n",
    "\n",
    "    # print(train_indices)\n",
    "    tr_inds = train_indices[:num_train]\n",
    "    mnist_tr1 = DT([dta[tr_inds], targets[tr_inds]])\n",
    "\n",
    "    val_inds = train_indices[num_train:num_train + num_valid]\n",
    "    mnist_valid = DT([dta[val_inds], targets[val_inds]])\n",
    "\n",
    "    mnist_test = datasets.MNIST('../data', download=True, train=False)\n",
    "\n",
    "    def _process_dataset(dts):\n",
    "        dt, tgt = np.array(dts.data.numpy(), dtype=np.float32), dts.targets.numpy()\n",
    "        return DT([th.from_numpy(\n",
    "            np.reshape(dt / 255., (-1, 28 * 28))).to(TARGET_DEVICE),\n",
    "                   th.from_numpy(tgt).to(TARGET_DEVICE)])\n",
    "\n",
    "    return TVT([_process_dataset(dtt) for dtt in [mnist_tr1, mnist_valid, mnist_test]])\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "i_sig = 0.01  # initialization\n",
    "dw = 200  # dimensionality of the hidden state\n",
    "\n",
    "lr = 0.5\n",
    "\n",
    "th.manual_seed(0)\n",
    "data = load_mnist(0, num_train=5000, num_valid=5000)\n",
    "num_exp, dim_x = data.train.data.shape"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "do_projection = True\n",
    "\n",
    "T = K = 20  # number of iterations; T for forward iterations, K for backward;"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# choose between\n",
    "# rm (reverse-mode iterative differentiation),\n",
    "# fp (fixed point implicit differentiation) and\n",
    "# cg (conjugate gradient implicit differentiation)\n",
    "\n",
    "hg_mode = 'rm'\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "def matrix_projection_on_spectral_ball(a, radius=0.99, project=True):\n",
    "    A = a.detach()\n",
    "    if A.is_cuda: A = A.cpu()\n",
    "    A = A.numpy()\n",
    "    U, S, V = np.linalg.svd(A)\n",
    "    if project:\n",
    "        S1 = np.minimum(S, radius)\n",
    "        a = U @ np.diag(S1) @ V\n",
    "    else:\n",
    "        a = A\n",
    "    return th.from_numpy(a).type(th.FloatTensor).to(TARGET_DEVICE).requires_grad_(True), S\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "initial_states = TVT([th.zeros(d.data.shape[0], dw, device=TARGET_DEVICE) for d in data])\n",
    "\n",
    "if hg == 'rm': set_requires_grad(initial_states)  # necessary only for reverse-mode with unrolling\n",
    "\n",
    "# define model's parameters\n",
    "parameters = [\n",
    "    i_sig * th.randn(dw, dw, device=TARGET_DEVICE),\n",
    "    i_sig * th.randn(dim_x, dw, device=TARGET_DEVICE),\n",
    "    i_sig * th.randn(dw, device=TARGET_DEVICE),\n",
    "    i_sig * th.randn(dw, 10, device=TARGET_DEVICE),\n",
    "    th.zeros(10, device=TARGET_DEVICE)\n",
    "]\n",
    "set_requires_grad(parameters)\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def get_fully_connected_dynamics(x):\n",
    "    def fully_connected_dynamics(state_list, params):\n",
    "        # RNNs like dynamics (the fp_map of the bi-level problem)\n",
    "        A, B, c = params[:3]\n",
    "        state = state_list[0]\n",
    "        return [th.tanh(state @ A + x @ B + c)]\n",
    "\n",
    "    return fully_connected_dynamics\n",
    "\n",
    "\n",
    "get_dynamics = get_fully_connected_dynamics  # change this line for changing type of dynamics\n",
    "\n",
    "# obtain one dynamics per set (training, validation and test) which is a callable\n",
    "tvt_dynamics = TVT([get_dynamics(dt.data) for dt in data])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "\n",
    "def linear(state_list, params):\n",
    "    return state_list[0] @ params[-2] + params[-1]\n",
    "\n",
    "\n",
    "def get_loss(targets):\n",
    "    def loss(state_list, params):\n",
    "        # cross entropy loss (the outer loss of the bi-level problem)\n",
    "        outputs = linear(state_list, params)\n",
    "        criterion = th.nn.CrossEntropyLoss()\n",
    "        return th.mean(criterion(outputs, targets))\n",
    "\n",
    "    return loss\n",
    "\n",
    "\n",
    "# obtain one loss per dataset (note: the losses remain callable as well as the dynamics!).\n",
    "tvt_losses = TVT([get_loss(dt.targets) for dt in data])\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# forward pass\n",
    "def get_forward(initial_state, dynamics):\n",
    "    def forward():\n",
    "        states = [[initial_state]]\n",
    "        for _ in range(T):\n",
    "            states.append(dynamics(states[-1], parameters))\n",
    "        return states\n",
    "\n",
    "    return forward\n",
    "\n",
    "\n",
    "# one per dataset\n",
    "tvt_forward = TVT([get_forward(s, dyna) for s, dyna in zip(initial_states, tvt_dynamics)])\n",
    "\n",
    "\n",
    "def metric_after_fw(forward, metric):\n",
    "    def _f():\n",
    "        states = forward()\n",
    "        return metric(states[-1], parameters)\n",
    "\n",
    "    return _f\n",
    "\n",
    "\n",
    "def accuracy(targets):\n",
    "    def _f(states, params):\n",
    "        return acc(linear(states, params), targets)\n",
    "\n",
    "    return _f\n",
    "\n",
    "\n",
    "# obtain callables for loss and accuracy for each set (after executing the model's dynamics)\n",
    "tvt_metrics = TVT([\n",
    "    LA([metric_after_fw(fww, lss), metric_after_fw(fww, accuracy(dt.targets))])\n",
    "    for fww, lss, dt in zip(tvt_forward, tvt_losses, data)\n",
    "])"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# optimizer\n",
    "opt = th.optim.SGD(parameters, lr, momentum=0.9)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# training!\n",
    "\n",
    "for t in range(1000):\n",
    "    opt.zero_grad()\n",
    "    states = tvt_forward.train()\n",
    "\n",
    "    # compute the hypergradient (with different methods)\n",
    "    if hg_mode == 'fp':\n",
    "        hg.fixed_point(states[-1], parameters, K, tvt_dynamics.train, tvt_losses.train)\n",
    "    elif hg_mode == 'cg':\n",
    "        hg.CG_normaleq(states[-1], parameters, K, tvt_dynamics.train, tvt_losses.train)\n",
    "    elif hg_mode == 'rm':\n",
    "        hg.reverse_unroll(states[-1], parameters, tvt_losses.train)\n",
    "    else:\n",
    "        raise NotImplementedError('{} not available!'.format(hg_mode))\n",
    "\n",
    "    opt.step()\n",
    "\n",
    "    try:  # perform projection\n",
    "        A_proj, svl = matrix_projection_on_spectral_ball(parameters[0], project=do_projection)\n",
    "        parameters[0].data = A_proj.data\n",
    "    except (ValueError, np.linalg.LinAlgError) as e:\n",
    "        print('there were nans most probably: aborting all')\n",
    "        break\n",
    "\n",
    "    if t % 20 == 0:\n",
    "        valid_acc = to_numpy(tvt_metrics.val.acc())\n",
    "        hgs = to_numpy([l.grad for l in parameters])\n",
    "\n",
    "        print('Validation accuracy at iteration {}:'.format(t), valid_acc)  # update early stopping\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}