MAP_lambdas_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MAP-training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## In this notebook we perform MAP-training with additional constraint on ranks "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.distributions import Normal, Gamma\n",
    "\n",
    "from random import shuffle\n",
    "\n",
    "from torchvision.datasets.mnist import MNIST\n",
    "from torch.utils.data import DataLoader, random_split, Subset\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "#from tqdm import tqdm_notebook as tqdm\n",
    "from tqdm import tqdm\n",
    "from tt_model import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 7666 12\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",
    "parameters_config = {\n",
    "    'batch_size': 200,\n",
    "    'train_size': 40000,\n",
    "    'device': torch.device('cuda'),\n",
    "    'learning_rate': 1e-4,\n",
    "    'n_epochs': 100,\n",
    "    \n",
    "    # gamma distribution parameters from paper\n",
    "#     'a_l': 1,\n",
    "#     'b_l': 5,\n",
    "    'a_l': 1,\n",
    "    'b_l': 15,\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",
    "model_cfg = AttrDict(config)\n",
    "params_cfg = AttrDict(parameters_config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TTModel(model_cfg)\n",
    "model = model.to(params_cfg.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "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",
    "dataset = MNIST('mnist', train=True, download=True, transform=MNIST_TRANSFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset, val_dataset = random_split(dataset, (params_cfg.train_size, len(dataset) - params_cfg.train_size))\n",
    "train_subset_indices = list(range(len(train_dataset)))\n",
    "shuffle(train_subset_indices)\n",
    "train_subset_indices = train_subset_indices[:10000]\n",
    "train_subset = Subset(train_dataset, train_subset_indices)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=params_cfg.batch_size, shuffle=True, pin_memory=(params_cfg.device.type == \"cuda\"))\n",
    "val_loader = DataLoader(val_dataset, batch_size=params_cfg.batch_size, shuffle=True, pin_memory=(params_cfg.device.type == \"cuda\"))\n",
    "train_subset_loader = DataLoader(train_subset, batch_size=params_cfg.batch_size, shuffle=True, pin_memory=(params_cfg.device.type == \"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(params_cfg.device)\n",
    "\n",
    "#optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.95, weight_decay=0.0005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "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",
    "        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",
    "#             print(l_prev.shape, l_next.shape)\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",
    "    log_g = log_prior_sum\n",
    "    log_lambda = 0\n",
    "    if lambdas is not None:\n",
    "        for layer_lambdas in lambdas:\n",
    "            for l in layer_lambdas:\n",
    "                log_lambda += Gamma(a_l, b_l).log_prob(l).sum()\n",
    "        \n",
    "    return log_g, log_lambda\n",
    "        \n",
    "\n",
    "\n",
    "def log_posterior(model, input, gt, lambdas=None, likelihood_coef=1., a_l=1, b_l=5):\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",
    "    model_out = model(input)\n",
    "    \n",
    "    log_g = torch.nn.functional.log_softmax(model_out, dim=1)\n",
    "    log_likelihood = (gt * log_g).sum()\n",
    "    \n",
    "    log_g_prior, log_lambda = log_prior(model, lambdas, a_l, b_l)\n",
    "    \n",
    "    # not including margnial log-likelihood log(p(D))\n",
    "    global global_step\n",
    "#     if global_step % 2000 == 0:\n",
    "#         print('Likelihood loss', -(likelihood_coef * log_likelihood).item())\n",
    "#         print('Prior G', -log_g_prior.item())\n",
    "#         print('Prior lambda', -log_lambda.item())\n",
    "    global_step += 1\n",
    "    return likelihood_coef * log_likelihood + (log_g_prior + log_lambda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lambdas proposed shapes: [8 6 3 8] [16]\n",
      "val_acc = 0.8630999999999998\n",
      "train acc = 0.861\n",
      "lambdas proposed shapes: [8 6 2 8] [16]\n",
      "val_acc = 0.8575499999999995\n",
      "train acc = 0.8530999999999999\n",
      "lambdas proposed shapes: [8 5 2 8] [16]\n",
      "val_acc = 0.8553000000000003\n",
      "train acc = 0.8516000000000002\n",
      "lambdas proposed shapes: [8 4 2 8] [16]\n",
      "val_acc = 0.8528500000000001\n",
      "train acc = 0.8484999999999998\n",
      "lambdas proposed shapes: [8 4 2 8] [16]\n",
      "val_acc = 0.84455\n",
      "train acc = 0.8405000000000002\n",
      "lambdas proposed shapes: [8 4 2 8] [16]\n",
      "val_acc = 0.8282499999999997\n",
      "train acc = 0.8225999999999997\n",
      "lambdas proposed shapes: [8 4 2 8] [16]\n",
      "val_acc = 0.8211999999999996\n",
      "train acc = 0.8170999999999999\n",
      "lambdas proposed shapes: [8 4 2 8] [16]\n",
      "val_acc = 0.8182999999999997\n",
      "train acc = 0.8136\n",
      "lambdas proposed shapes: [8 3 2 8] [16]\n",
      "val_acc = 0.81395\n",
      "train acc = 0.8107000000000002\n",
      "lambdas proposed shapes: [8 3 2 8] [16]\n",
      "val_acc = 0.7966500000000003\n",
      "train acc = 0.7930999999999999\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7972499999999999\n",
      "train acc = 0.7943000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7919500000000004\n",
      "train acc = 0.7885000000000002\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7891499999999999\n",
      "train acc = 0.7844\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7818500000000003\n",
      "train acc = 0.7783999999999999\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7809\n",
      "train acc = 0.7787999999999999\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7837999999999998\n",
      "train acc = 0.7815000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7802500000000001\n",
      "train acc = 0.7796000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7771000000000005\n",
      "train acc = 0.7771999999999999\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7781999999999999\n",
      "train acc = 0.7775000000000003\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7750499999999997\n",
      "train acc = 0.7764000000000002\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7751000000000003\n",
      "train acc = 0.7746000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7720000000000004\n",
      "train acc = 0.7710000000000002\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7680500000000002\n",
      "train acc = 0.7655000000000005\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7676999999999999\n",
      "train acc = 0.7659000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7652000000000001\n",
      "train acc = 0.7641000000000001\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7634500000000002\n",
      "train acc = 0.7614\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7618999999999997\n",
      "train acc = 0.7603000000000002\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7579000000000002\n",
      "train acc = 0.7569\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7601000000000002\n",
      "train acc = 0.7572999999999999\n",
      "lambdas proposed shapes: [8 2 2 8] [16]\n",
      "val_acc = 0.7601500000000001\n",
      "train acc = 0.7571000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7565\n",
      "train acc = 0.7541999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7540999999999998\n",
      "train acc = 0.7528999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7523500000000003\n",
      "train acc = 0.7512000000000003\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7487500000000006\n",
      "train acc = 0.7472\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7496499999999999\n",
      "train acc = 0.7494000000000003\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7479499999999999\n",
      "train acc = 0.7468000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7475999999999999\n",
      "train acc = 0.7464999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7441999999999998\n",
      "train acc = 0.7397999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7460499999999999\n",
      "train acc = 0.743\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7440999999999997\n",
      "train acc = 0.7413000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7455499999999999\n",
      "train acc = 0.7428000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7449000000000001\n",
      "train acc = 0.7436999999999998\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7458\n",
      "train acc = 0.7436999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7413499999999997\n",
      "train acc = 0.738\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.74605\n",
      "train acc = 0.7420999999999998\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.74775\n",
      "train acc = 0.7427999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7465000000000004\n",
      "train acc = 0.7424\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.74755\n",
      "train acc = 0.7440000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.74845\n",
      "train acc = 0.7454000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.74855\n",
      "train acc = 0.7453000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7514999999999996\n",
      "train acc = 0.7515000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7525499999999998\n",
      "train acc = 0.7499000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7503000000000002\n",
      "train acc = 0.7482\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7556999999999998\n",
      "train acc = 0.7515999999999998\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7564500000000003\n",
      "train acc = 0.7515\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7508\n",
      "train acc = 0.7479999999999998\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7549499999999999\n",
      "train acc = 0.7526\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7561500000000001\n",
      "train acc = 0.7541000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.75865\n",
      "train acc = 0.7553999999999997\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7571499999999994\n",
      "train acc = 0.7541999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7602500000000005\n",
      "train acc = 0.7577000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7616499999999995\n",
      "train acc = 0.7590000000000003\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7636\n",
      "train acc = 0.7621000000000001\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7669500000000002\n",
      "train acc = 0.7631000000000003\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7681500000000002\n",
      "train acc = 0.7650999999999999\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7678999999999996\n",
      "train acc = 0.7663\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7714999999999997\n",
      "train acc = 0.7694\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7743499999999994\n",
      "train acc = 0.7726000000000002\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7775999999999996\n",
      "train acc = 0.7752\n",
      "lambdas proposed shapes: [8 2 3 8] [16]\n",
      "val_acc = 0.7777000000000003\n",
      "train acc = 0.7757\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-84-2ae8e34e381b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[0mglobal_step\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     22\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mep\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparams_cfg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mn_epochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 23\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    558\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_workers\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# same-process loading\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    559\u001b[0m             \u001b[0mindices\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample_iter\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 560\u001b[0;31m             \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollate_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mindices\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    561\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    562\u001b[0m                 \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m<listcomp>\u001b[0;34m(.0)\u001b[0m\n\u001b[1;32m    558\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum_workers\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# same-process loading\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    559\u001b[0m             \u001b[0mindices\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnext\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msample_iter\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 560\u001b[0;31m             \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollate_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mindices\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    561\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    562\u001b[0m                 \u001b[0mbatch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torch/utils/data/dataset.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, idx)\u001b[0m\n\u001b[1;32m    105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    106\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__getitem__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0midx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 107\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mindices\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0midx\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    108\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    109\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__len__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torchvision/datasets/mnist.py\u001b[0m in \u001b[0;36m__getitem__\u001b[0;34m(self, index)\u001b[0m\n\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 95\u001b[0;31m             \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     96\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     97\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtarget_transform\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torchvision/transforms/transforms.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, img)\u001b[0m\n\u001b[1;32m     59\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     60\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtransforms\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 61\u001b[0;31m             \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     62\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mimg\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     63\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torchvision/transforms/transforms.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, pic)\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0mTensor\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mConverted\u001b[0m \u001b[0mimage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m         \"\"\"\n\u001b[0;32m---> 92\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__repr__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/torchvision/transforms/functional.py\u001b[0m in \u001b[0;36mto_tensor\u001b[0;34m(pic)\u001b[0m\n\u001b[1;32m     77\u001b[0m         \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m255\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_numpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muint8\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     78\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 79\u001b[0;31m         \u001b[0mimg\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mByteTensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mByteStorage\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfrom_buffer\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpic\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtobytes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     80\u001b[0m     \u001b[0;31m# PIL image mode: L, LA, P, I, F, RGB, YCbCr, RGBA, CMYK\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     81\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mpic\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmode\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'YCbCr'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/PIL/Image.py\u001b[0m in \u001b[0;36mtobytes\u001b[0;34m(self, encoder_name, *args)\u001b[0m\n\u001b[1;32m    747\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    748\u001b[0m         \u001b[0;31m# unpack data\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 749\u001b[0;31m         \u001b[0me\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_getencoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoder_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    750\u001b[0m         \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msetimage\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mim\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    751\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/uniconf-Cjmz1t59/lib/python3.7/site-packages/PIL/Image.py\u001b[0m in \u001b[0;36m_getencoder\u001b[0;34m(mode, encoder_name, args, extra)\u001b[0m\n\u001b[1;32m    474\u001b[0m         \u001b[0;31m# get encoder\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    475\u001b[0m         \u001b[0mencoder\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcore\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mencoder_name\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\"_encoder\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 476\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mencoder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mextra\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    477\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mAttributeError\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    478\u001b[0m         \u001b[0;32mraise\u001b[0m \u001b[0mIOError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"encoder %s not available\"\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0mencoder_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "def acc(model, loader):\n",
    "    accs = []\n",
    "    with torch.no_grad():\n",
    "        for b, gt in loader:\n",
    "            out = model(b.to(params_cfg.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",
    "\n",
    "lambdas1 = nn.ParameterList([nn.Parameter(torch.distributions.Gamma(params_cfg.a_l, params_cfg. b_l).sample([r])) for r in model_cfg.l1_ranks]).to(params_cfg.device)\n",
    "lambdas2 = nn.ParameterList([nn.Parameter(torch.distributions.Gamma(params_cfg.a_l, params_cfg.b_l).sample([r])) for r in model_cfg.l2_ranks]).to(params_cfg.device)\n",
    "                     \n",
    "\n",
    "optimizer = torch.optim.Adam(list(model.parameters())\n",
    "                             + list(lambdas1)\n",
    "                             + list(lambdas2)\n",
    "                             , lr=params_cfg.learning_rate)\n",
    "\n",
    "\n",
    "global_step = 0\n",
    "for ep in range(params_cfg.n_epochs):\n",
    "    for b, gt in 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) / params_cfg.batch_size\n",
    "        loss = -log_posterior(model, b.to(params_cfg.device), onehot_gt.to(params_cfg.device),\n",
    "                              lambdas=[lambdas1,lambdas2], \n",
    "                              a_l=params_cfg.a_l, b_l=params_cfg.b_l, likelihood_coef=likelihood_coef)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "#     print(loss.item())\n",
    "    print('lambdas proposed shapes:', \n",
    "          np.array(model_cfg['l1_ranks']) - [(q.cpu().numpy() < 1e-2).sum() for q in lambdas1.state_dict().values()],\n",
    "          np.array(model_cfg['l2_ranks']) - [(q.cpu().numpy() < 1e-2).sum() for q in lambdas2.state_dict().values()])\n",
    "    print(f\"val_acc = {acc(model, val_loader)}\")\n",
    "    print(f\"train acc = {acc(model, train_subset_loader)}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see fair rank constraints obtained by optimization by lambdas, but accuracy drop is too high.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "uniconf",
   "language": "python",
   "name": "uniconf"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}