tt_vs_nn.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "53ff8cc6",
   "metadata": {},
   "source": [
    "'''\n",
    "    \n",
    "    Copyright (c) 2022 Idiap Research Institute, http://www.idiap.ch/\n",
    "    Written by Suhan Shetty <suhan.shetty@idiap.ch>,\n",
    "   \n",
    "    This file is part of TTGO.\n",
    "\n",
    "    TTGO is free software: you can redistribute it and/or modify\n",
    "    it under the terms of the GNU General Public License version 3 as\n",
    "    published by the Free Software Foundation.\n",
    "\n",
    "    TTGO is distributed in the hope that it will be useful,\n",
    "    but WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the\n",
    "    GNU General Public License for more details.\n",
    "\n",
    "    You should have received a copy of the GNU General Public License\n",
    "    along with TTGO. If not, see <http://www.gnu.org/licenses/>.\n",
    "'''\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4552dce4",
   "metadata": {},
   "source": [
    "### Comparision of performance of TT vs NN\n",
    "In this notebook,  we compare the approximation accuracy and speed of training between TT and NN. NN is a great tool for data-driven function approximation. However, it is not that great when the function to be approximated is given. On the other hand, TT is equipped with powerful technique called TT-Cross that can approximate a given function in TT format more efficiently. It directly takes the function to be approximated as input and outputs the function in TT format. Moreover, TT representation, unlike NN, offers other benefits like fast ways to sample, optimize, do algebra etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1ddacdf5-a59b-48b1-8431-57c9032fb439",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/idiap/temp/sshetty/miniconda/envs/pyml/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from tt_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ff5e3456-9ee6-4705-a749-de3abd085edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "65330fb7-0821-4cd6-94dd-b49fef8c643f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def gmm(n=2,nmix=3,L=1,mx_coef=None,mu=None,s=0.1, device='cpu'):\n",
    "    \"\"\"\n",
    "        Mixture of spherical Gaussians (un-normalized)\n",
    "        nmix: number of mixture coefficients\n",
    "        n: dimension of the domain\n",
    "        s: variance\n",
    "        mu: the centers assumed to be in : [-L,L]^n\n",
    "    \"\"\"\n",
    "    n_sqrt = torch.sqrt(torch.tensor([n]).to(device))\n",
    "    if mx_coef is None: # if centers and mixture coef are not given, generate them randomly\n",
    "        mx_coef = torch.rand(nmix).to(device)\n",
    "        mx_coef = mx_coef/torch.sum(mx_coef)\n",
    "        mu = (torch.rand(nmix,n).to(device)-0.5)*2*L\n",
    "\n",
    "    def pdf(x):\n",
    "        result = torch.tensor([0]).to(device)\n",
    "        for k in range(nmix):\n",
    "            l = torch.linalg.norm(mu[k]-x, dim=1)/n_sqrt\n",
    "            result = result + mx_coef[k]*torch.exp(-(l/s)**2)\n",
    "        return 1.+100*result\n",
    "\n",
    "    return pdf\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "29d66285-42f9-4bb4-bc0f-57d82e20bd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 10\n",
    "L = 1\n",
    "nmix = 1\n",
    "s = 0.2\n",
    "\n",
    "# generate an arbitrary function (gmm with centers and covariances chosen randomly)\n",
    "pdf = gmm(n=dim,nmix=nmix,L=L,mx_coef=None,mu=None,s=s, device=device) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4b86ba96-9280-4706-9f55-820ea55ce398",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cross device is cuda\n",
      "Cross-approximation over a 10D domain containing 1.024e+23 grid points:\n",
      "iter: 0  | tt-error: 1.074e+00, test-error:9.370e-01 | time:   1.3723 | largest rank:   1\n",
      "iter: 1  | tt-error: 2.653e+00, test-error:6.331e-15 | time:   1.5148 | largest rank:   4\n",
      "iter: 2  | tt-error: 0.000e+00, test-error:1.367e-14 | time:   1.6755 | largest rank:   7 <- converged: eps < 0.001\n",
      "Did 218400 function evaluations, which took 0.03448s (6.335e+06 evals/s)\n",
      "\n",
      "time taken:  1.8815488815307617\n"
     ]
    }
   ],
   "source": [
    "# Represent the function in TT format (unsupervised learning and kind of non-parametric)\n",
    "n_discretization = torch.tensor([200]*dim).to(device)\n",
    "domain = [torch.linspace(-L,L,n_discretization[i]).to(device) for i in range(dim)] \n",
    "\n",
    "import time \n",
    "t1 = time.time()\n",
    "tt_gmm = cross_approximate(fcn=pdf,  max_batch=10**6, domain=domain, \n",
    "                        rmax=200, nswp=20, eps=1e-3, verbose=True, \n",
    "                        kickrank=3, device=device)\n",
    "t2 = time.time()\n",
    "print(\"time taken: \", t2-t1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7239100c-5790-43f7-bbc3-bf6a3aeedb35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For testing and training NN\n",
    "ndata_train = 200000\n",
    "ndata_test = 1000\n",
    "\n",
    "x_train = 2*L*(-0.5 + torch.rand((ndata_train,dim)).to(device))\n",
    "y_train = pdf(x_train)\n",
    "\n",
    "x_test = 2*L*(-0.5 + torch.rand((ndata_test,dim)).to(device))\n",
    "y_test = pdf(x_test)\n",
    "\n",
    "data_train = torch.cat((x_train.view(-1,dim),y_train.view(-1,1)),dim=-1)\n",
    "data_test = torch.cat((x_test.view(-1,dim),y_test.view(-1,1)),dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bd26218e-27c2-4c03-9e52-25f2c0640edd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_tt:  tensor(5.6480e-11, device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "# Test the accuracy of TT over the test set \n",
    "y_tt =  get_value(tt_model=tt_gmm, x=x_test.to(device),  domain=domain, \n",
    "                    n_discretization=n_discretization , max_batch=10**5, device=device)\n",
    "\n",
    "mse_tt = (((y_tt.view(-1)-y_test.view(-1))/(1e-9+y_test.view(-1).abs()))**2).mean()\n",
    "print(\"mse_tt: \", mse_tt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1e70ec4e",
   "metadata": {},
   "source": [
    "#### Represent the function as a NN "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "30911676-7815-4463-bd24-37f12d4ca2a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Matplotlib created a temporary config/cache directory at /tmp/matplotlib-wmlzg2qv because the default path (/idiap/home/sshetty/.cache/matplotlib) is not a writable directory; it is highly recommended to set the MPLCONFIGDIR environment variable to a writable directory, in particular to speed up the import of Matplotlib and to better support multiprocessing.\n"
     ]
    }
   ],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6817e869-498f-4627-8ad5-4876c271bf01",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self, dim=2, width=32):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(dim, width),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(width, width),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(width, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork(dim=dim, width=dim*nmix*10).to(device)\n",
    "def train_loop(data, model, loss_fn, optimizer, batch_size):\n",
    "    size = data.shape[0]\n",
    "    counter = 0\n",
    "    for i in range(int(size/batch_size)-1):\n",
    "        # Compute prediction and loss\n",
    "        next_counter = (counter+batch_size)\n",
    "        x_data = data[counter:next_counter,:-1]\n",
    "        y_data = data[counter:next_counter,-1].view(-1,1)\n",
    "        y_pred = model(x_data)\n",
    "        loss = loss_fn(y_pred, y_data)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        counter = 1*next_counter\n",
    "\n",
    "        if (i % int(0.25*size/batch_size)) == 0 :\n",
    "            loss = loss.item()\n",
    "            print(f\"loss: {loss:>7f}\")\n",
    "\n",
    "\n",
    "def test_loop(data, model, loss_fn):\n",
    "    x_data = data[:,:-1]\n",
    "    y_data = data[:,-1]\n",
    "    with torch.no_grad():\n",
    "        pred = model(x_data)\n",
    "        test_loss = loss_fn(pred, y_data).item()\n",
    "    print(f\"Test Error: \", test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "66181567-4394-4706-9c8f-9fb526bd7aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 100\n",
    "epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "88f91213-495d-4c42-bd26-77a548caf7f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn_0:  tensor(1.1368, device='cuda:0', grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "y_nn_0 = model(x_test)\n",
    "mse_nn_0 = (((y_nn_0.view(-1)-y_test.view(-1))/(1e-9+y_test.view(-1).abs()))**2).mean()\n",
    "print(\"mse_nn_0: \", mse_nn_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9f980e0b-41d4-46d2-a540-648a1daf118a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 1.178049\n",
      "loss: 0.002727\n",
      "loss: 0.130106\n",
      "loss: 0.001618\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/idiap/temp/sshetty/miniconda/envs/pyml/lib/python3.9/site-packages/torch/nn/modules/loss.py:530: UserWarning: Using a target size (torch.Size([1000])) that is different to the input size (torch.Size([1000, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
      "  return F.mse_loss(input, target, reduction=self.reduction)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Error:  0.19129082062105746\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.014067\n",
      "loss: 0.001272\n",
      "loss: 0.296887\n",
      "loss: 0.001081\n",
      "Test Error:  0.24745713466026872\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.001960\n",
      "loss: 0.000607\n",
      "loss: 0.107745\n",
      "loss: 0.000396\n",
      "Test Error:  0.2681934520381166\n",
      "Epoch 4\n",
      "-------------------------------\n",
      "loss: 0.003532\n",
      "loss: 0.000300\n",
      "loss: 0.015622\n",
      "loss: 0.000641\n",
      "Test Error:  0.2947253271853954\n",
      "Epoch 5\n",
      "-------------------------------\n",
      "loss: 0.003216\n",
      "loss: 0.000287\n",
      "loss: 0.003205\n",
      "loss: 0.000881\n",
      "Test Error:  0.28090638369930526\n",
      "Epoch 6\n",
      "-------------------------------\n",
      "loss: 0.004636\n",
      "loss: 0.000319\n",
      "loss: 0.000631\n",
      "loss: 0.000360\n",
      "Test Error:  0.3405355712848097\n",
      "Epoch 7\n",
      "-------------------------------\n",
      "loss: 0.002147\n",
      "loss: 0.000153\n",
      "loss: 0.000551\n",
      "loss: 0.000163\n",
      "Test Error:  0.29791631463540097\n",
      "Epoch 8\n",
      "-------------------------------\n",
      "loss: 0.002407\n",
      "loss: 0.000159\n",
      "loss: 0.001045\n",
      "loss: 0.000148\n",
      "Test Error:  0.3015412034751193\n",
      "Epoch 9\n",
      "-------------------------------\n",
      "loss: 0.001409\n",
      "loss: 0.000373\n",
      "loss: 0.003092\n",
      "loss: 0.000480\n",
      "Test Error:  0.30012809493840126\n",
      "Epoch 10\n",
      "-------------------------------\n",
      "loss: 0.000574\n",
      "loss: 0.000890\n",
      "loss: 0.001775\n",
      "loss: 0.000553\n",
      "Test Error:  0.29518859088854293\n",
      "time taken:  50.70554804801941\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# NN training (Note: compare the time it takes with TT-Cross)\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "batch_size = 100\n",
    "epochs = 10\n",
    "t1 = time.time()\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(data_train, model, loss_fn, optimizer, batch_size)\n",
    "    test_loop(data_test, model, loss_fn)\n",
    "t2 = time.time()\n",
    "print(\"time taken: \", t2-t1)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7539baea-ecda-4204-ad3b-58b246262033",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn:  tensor(0.0003, device='cuda:0', grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "# Test the accuracy over the test set\n",
    "y_nn = model(x_test)\n",
    "mse_nn = (((y_nn.view(-1)-y_test.view(-1))/(1e-9+y_test.view(-1).abs()))**2).mean()\n",
    "print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0124e4d3-5e3a-45ec-a321-de88a81a02c6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn:  tensor(0.0003, device='cuda:0', grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "y_nn = model(x_test)\n",
    "mse_nn = (((y_nn.view(-1)-y_test.view(-1))/(1e-9+y_test.view(-1).abs()))**2).mean()\n",
    "print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c542ebc4-ce13-4403-a75b-3a43300add3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_tt:  tensor(5.6480e-11, device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "mse_tt = (((y_tt.view(-1)-y_test.view(-1))/(1e-9+y_test.view(-1).abs()))**2).mean()\n",
    "print(\"mse_tt: \", mse_tt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "04b08e8e-ecfb-4ccd-9d24-bfd871eb1748",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn:  tensor(0.0006, device='cuda:0', grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "y_nn = model(x_train)\n",
    "mse_nn = (((y_nn.view(-1)-y_train.view(-1))/(1e-6+y_train.view(-1).abs()))**2).mean()\n",
    "print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "37936493-c0a5-406e-a595-5917051d7b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# y_nn = model(x_train)\n",
    "# mse_nn = ((y_nn-y_train)**2).mean()\n",
    "# print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d815d40d-ca43-4732-a534-98b7a6a9fe96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn:  tensor(0.2952, device='cuda:0', grad_fn=<MeanBackward0>)\n"
     ]
    }
   ],
   "source": [
    "y_nn = model(x_test)\n",
    "mse_nn = ((y_nn-y_test)**2).mean()\n",
    "print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cf370189-7146-431a-a7e7-1b38ed504902",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_tt:  tensor(6.4013e-09, device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "mse_tt = (((y_tt.view(-1)-y_test.view(-1)))**2).mean()\n",
    "print(\"mse_tt: \", mse_tt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75724813-e417-49aa-9510-2a9042977193",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "941b4c73",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c30acd-efeb-4493-a41d-b9c3f1e1a61e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}