map_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.distributions import Normal, Gamma, MultivariateNormal\n",
    "\n",
    "from torchvision.datasets.mnist import MNIST\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tt_model import TTModel\n",
    "from tt_model import vectorize_params, unvectorize_params\n",
    "from tt_model import calc_ranks_from_particles, get_lambdas_vec\n",
    "\n",
    "config = {\n",
    "    'resize_shape': (32, 32),\n",
    "    \n",
    "    'in_factors': (4, 4, 4, 4, 4),\n",
    "    'l1_ranks': (8, 8, 8, 8),\n",
    "    'hidd_out_factors': (2, 2, 2, 2, 2),\n",
    "    'ein_string1': \"nabcde,aoiv,bijw,cjkx,dkly,elpz\",\n",
    "    \n",
    "    'hidd_in_factors': (4, 8),\n",
    "    'l2_ranks': (16,),\n",
    "    'out_factors': (5, 2),\n",
    "    'ein_string2': 'nab,aoix,bipy',\n",
    "}\n",
    "\n",
    "class AttrDict(dict):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(AttrDict, self).__init__(*args, **kwargs)\n",
    "        self.__dict__ = self\n",
    "        \n",
    "cfg = AttrDict(config)\n",
    "\n",
    "model = TTModel(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_LABELS = 10\n",
    "MNIST_TRANSFORM = transforms.Compose((\n",
    "    transforms.Pad(2),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1,), (0.2752,))\n",
    "))\n",
    "\n",
    "device = torch.device('cpu')\n",
    "\n",
    "train_size = 40000\n",
    "batch_size = 100\n",
    "dataset = MNIST('mnist', train=True, download=True, transform=MNIST_TRANSFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset, val_dataset = random_split(dataset, (train_size, len(dataset) - train_size))\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=(device.type == \"cuda\"))\n",
    "val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, pin_memory=(device.type == \"cuda\"))\n",
    "\n",
    "model = model.to(device)\n",
    "\n",
    "#optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.95, weight_decay=0.0005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_prior(model, lambdas=None, a_l=1, b_l=5):\n",
    "    log_prior_sum = 0\n",
    "    for name, core_tensor in model.named_parameters():\n",
    "        if 'tt' not in name:\n",
    "            continue\n",
    "            \n",
    "        core_mean = torch.zeros_like(core_tensor)\n",
    "        \n",
    "        if lambdas is None:\n",
    "            core_std = torch.ones_like(core_tensor)\n",
    "        else:\n",
    "            layer_idx = int(name.split('tt')[-1].split('.')[0])\n",
    "            core_idx = int(name.split('cores.')[-1])\n",
    "            \n",
    "            prev_rank = core_tensor.shape[1]\n",
    "            next_rank = core_tensor.shape[2]\n",
    "        \n",
    "            if prev_rank == 1:\n",
    "                l_next = lambdas[layer_idx][core_idx]\n",
    "                l_prev = l_next\n",
    "            elif next_rank  == 1:\n",
    "                l_prev = lambdas[layer_idx][core_idx - 1]\n",
    "                l_next = l_prev\n",
    "            else:\n",
    "                l_prev = lambdas[layer_idx][core_idx - 1]\n",
    "                l_next = lambdas[layer_idx][core_idx]\n",
    "            \n",
    "            core_std = torch.einsum('p,q->pq', l_prev, l_next)\n",
    "            core_std = core_std.repeat(core_tensor.shape[0], core_tensor.shape[3], 1, 1).permute(0, 2, 3, 1)\n",
    "            \n",
    "        log_prior_sum += Normal(core_mean, core_std).log_prob(core_tensor).sum()\n",
    "    \n",
    "    if lambdas is not None:\n",
    "        for layer_lambdas in lambdas:\n",
    "            for l in layer_lambdas:\n",
    "                log_prior_sum += Gamma(a_l, b_l).log_prob(l).sum()\n",
    "        \n",
    "    return log_prior_sum\n",
    "        \n",
    "\n",
    "\n",
    "def log_posterior(model, input, gt, lambdas=None, likelihood_coef=1.):\n",
    "    \"\"\"Calculate log-posterior for core tensors and lambdas (optional)\n",
    "\n",
    "    Parameters:   \n",
    "        model : TT-model with core tensors as parameters, \n",
    "        input : Model input\n",
    "        gt : Ground truth\n",
    "        lambdas : LR-parameters \\lambda, if any\n",
    "        likelihood_coef : Coefficient to multiply log-likelihood by (for batches)\n",
    "    \n",
    "    Returns:\n",
    "        Log-posterior \n",
    "    \"\"\"\n",
    "    \n",
    "    log_g = torch.nn.functional.log_softmax(model(input), dim=1)\n",
    "    log_likelihood = (gt * log_g).sum()\n",
    "    \n",
    "    #print(out, log_g)\n",
    "    \n",
    "    log_prior_sum = log_prior(model, lambdas, 1, 5)\n",
    "    \n",
    "    # not including margnial log-likelihood log(p(D))\n",
    "    return likelihood_coef * log_likelihood + log_prior_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efe6deb984ad4788849d0f13c8fedce8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/site-packages/ipykernel_launcher.py:80: FutureWarning: arrays to stack must be passed as a \"sequence\" type such as list or tuple. Support for non-sequence iterables such as generators is deprecated as of NumPy 1.16 and will raise an error in the future.\n"
     ]
    }
   ],
   "source": [
    "from time import time\n",
    "def svgd(particles, dataloader=train_loader, kernel_sigma_squared=1., eps=1e-4, n_steps=100):\n",
    "    \n",
    "    # to get RBF exactly, not just probability\n",
    "    kernel_const = particles[0].shape[0] / 2 * (np.log(np.pi) + np.log(kernel_sigma_squared))\n",
    "    \n",
    "    \n",
    "    for step in tqdm(range(n_steps)):\n",
    "        t0 = time()\n",
    "        log_posterior_grads = []\n",
    "        kernels = []\n",
    "        phi = []\n",
    "\n",
    "        tk = time()\n",
    "        \n",
    "        for theta, (b, gt) in zip(particles, dataloader):\n",
    "            t1 = time()\n",
    "            model, lambdas = unvectorize_params(theta, cfg)\n",
    "            #print('unvectorize', time() - t1)\n",
    "            onehot_gt = torch.zeros(gt.shape[0], NUM_LABELS).scatter_(1, gt.view(-1, 1), 1)\n",
    "            log_posterior_i = log_posterior(model, b, onehot_gt, lambdas)\n",
    "            log_posterior_i.backward()\n",
    "            log_posterior_grad = theta.grad\n",
    "\n",
    "            log_posterior_grads.append(log_posterior_grad)\n",
    "            theta.grad.data.zero_()\n",
    "\n",
    "            kernels.append(MultivariateNormal(theta, kernel_sigma_squared * torch.eye(theta.shape[0])))\n",
    "        #print('tk', time() - tk)\n",
    "        t2f = time()\n",
    "        for k, theta_k in enumerate(particles):\n",
    "            phi_k = 0\n",
    "            for i, theta_i in enumerate(particles):\n",
    "                #theta_i.grad.data.zero_()\n",
    "\n",
    "                kernel_ik = torch.exp(kernels[k].log_prob(theta_i) + kernel_const)\n",
    "                kernel_ik.backward()\n",
    "                kernel_ik_grad = theta_i.grad\n",
    "                theta_i.grad.data.zero_()\n",
    "\n",
    "                phi_k += kernel_ik * log_posterior_grads[i] + kernel_ik_grad\n",
    "            phi_k /= len(particles)\n",
    "            phi.append(phi_k)\n",
    "        #print('t2f', time() - t2f)\n",
    "\n",
    "        for k, phi_k in enumerate(phi):\n",
    "            particles[k].data.add_(eps * phi_k)\n",
    "        #print('step', time() - t0)\n",
    "\n",
    "    return particles\n",
    "\n",
    "def get_svgd_ranks(n_particles=10, n_steps=100, rank_as_mean=True, a_l=1, b_l=5, kernel_sigma_squared=1.):\n",
    "    # sample particles:\n",
    "    # * initialize empty model for G's\n",
    "    # * sample lambdas from priors\n",
    "    # * vectorize the particle\n",
    "    particles = []\n",
    "    for n in range(n_particles):\n",
    "        model = TTModel(cfg)\n",
    "        lambdas = []\n",
    "        layer_ranks = [cfg.l1_ranks, cfg.l2_ranks]\n",
    "        for layer in range(2):\n",
    "            lambdas.append([])\n",
    "            for j, rank in enumerate(layer_ranks[layer]):\n",
    "                lambdas[layer].append(Gamma(a_l, b_l).rsample([rank]))\n",
    "        \n",
    "        particle = vectorize_params(model, lambdas).detach()\n",
    "        particle.requires_grad_(True)\n",
    "        particles.append(particle)\n",
    "            \n",
    "    # n_steps times perform svgd_step\n",
    "    svgd_particles = svgd(particles, kernel_sigma_squared=kernel_sigma_squared, n_steps=n_steps)\n",
    "    #return svgd_particles\n",
    "\n",
    "    if rank_as_mean:\n",
    "        # calculate mean of posterior lambdas\n",
    "        threshold = 1e-2\n",
    "\n",
    "        lambdas = tuple(unvectorize_params(particle, cfg)[1] for particle in svgd_particles)\n",
    "        lambdas_vecs = np.column_stack((get_lambdas_vec(lambda_).detach().cpu().numpy() for lambda_ in lambdas)).T\n",
    "\n",
    "        sample_ranks = []\n",
    "        for i in range(len(svgd_particles)):\n",
    "            # sample the same number of times as number of particles\n",
    "\n",
    "            sample_mean = torch.tensor(lambdas_vecs[np.random.choice(range(len(lambdas_vecs)))])\n",
    "            lambdas_sample = MultivariateNormal(sample_mean, kernel_sigma_squared * torch.eye(sample_mean.shape[0])).sample()\n",
    "\n",
    "\n",
    "            cur = 0\n",
    "            ranks = []\n",
    "            for layer_max_ranks in (cfg.l1_ranks, cfg.l2_ranks):\n",
    "                this_layer_ranks = []\n",
    "                for max_rank in layer_max_ranks:\n",
    "                    this_layer_ranks.append((lambdas_sample[cur:cur+max_rank] > threshold).sum())\n",
    "                    cur += max_rank\n",
    "                ranks.append(this_layer_ranks)        \n",
    "            sample_ranks.append(ranks)\n",
    "\n",
    "        ranks = []\n",
    "        for l in range(2):\n",
    "            layer_ranks = []\n",
    "            for j in range(len(sample_ranks[0][l])):\n",
    "                core_ranks = [sample_ranks[k][l][j] for k in range(len(sample_ranks))]\n",
    "                layer_ranks.append(sum(core_ranks) / len(core_ranks))\n",
    "            ranks.append(layer_ranks)\n",
    "                    \n",
    "    else:\n",
    "        # calculate mode of posterior lambdas\n",
    "        ranks = calc_ranks_from_particles(cfg, svgd_particles, \n",
    "                                          kernel_sigma_squared=kernel_sigma_squared, \n",
    "                                          how='mode')\n",
    "\n",
    "    \n",
    "    return ranks, svgd_particles\n",
    "    \n",
    "ranks, svgd_particles = get_svgd_ranks(10, kernel_sigma_squared=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[tensor(7), tensor(7), tensor(6), tensor(6)], [tensor(13)]]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ranks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MAP training for new ranks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def acc(model, loader):\n",
    "    accs = []\n",
    "    with torch.no_grad():\n",
    "        for b, gt in tqdm(loader):\n",
    "            out = model(b.to(device)).argmax(1).cpu().numpy()\n",
    "            gt = gt.numpy()\n",
    "            accs.append(sum(out == gt) / len(out))\n",
    "    return sum(accs) / len(accs) \n",
    "\n",
    "learning_rate = 1e-3\n",
    "n_epochs = 100\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "for ep in range(n_epochs):\n",
    "    for b, gt in tqdm(train_loader):\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "        onehot_gt = torch.zeros(gt.shape[0], NUM_LABELS).scatter_(1, gt.view(-1, 1), 1)\n",
    "\n",
    "        likelihood_coef = len(train_dataset) / batch_size\n",
    "        loss = -log_posterior(model, b.to(device), onehot_gt.to(device), likelihood_coef=likelihood_coef)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    print(loss.item())\n",
    "    #print(acc(model, val_loader))\n",
    "    "
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}