tt_vs_nn_vs_bgmm.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 vs BGMM\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.\n",
    "\n",
    "We also compare it against Bayesian GMM. Note that unlike, NN and TT, BGMM requires exact samples from the reference pdf to be fit. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1ddacdf5-a59b-48b1-8431-57c9032fb439",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tt_utils import *\n",
    "from fcn_approx_utils import GMM, NeuralNetwork, BGMM\n",
    "import time \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ff5e3456-9ee6-4705-a749-de3abd085edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\"#torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "29d66285-42f9-4bb4-bc0f-57d82e20bd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dim = 5\n",
    "L = 1\n",
    "nmix = 20\n",
    "s = 0.2\n",
    "\n",
    "# generate an arbitrary function (gmm with centers and covariances chosen randomly)\n",
    "gmm = GMM(n=dim,nmix=nmix,L=L,mx_coef=None,mu=None,s=s, device=device) \n",
    "pdf = gmm.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7239100c-5790-43f7-bbc3-bf6a3aeedb35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For testing and training NN\n",
    "ndata_train = int(1e5)\n",
    "ndata_test = 10000\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)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e2f2b2db",
   "metadata": {},
   "source": [
    "### Fit TT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4b86ba96-9280-4706-9f55-820ea55ce398",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cross device is cpu\n",
      "Cross-approximation over a 5D domain containing 3.2e+11 grid points:\n",
      "iter: 0  | tt-error: 1.000e+00, test-error:9.807e-01 | time:   0.0699 | largest rank:   1\n",
      "iter: 1  | tt-error: 2.078e+00, test-error:8.744e-01 | time:   0.2167 | largest rank:   4\n",
      "iter: 2  | tt-error: 1.188e+00, test-error:7.371e-01 | time:   0.3070 | largest rank:   7\n",
      "iter: 3  | tt-error: 7.334e-01, test-error:5.693e-01 | time:   0.4282 | largest rank:  10\n",
      "iter: 4  | tt-error: 5.292e-01, test-error:3.752e-01 | time:   0.5595 | largest rank:  13\n",
      "iter: 5  | tt-error: 1.260e-01, test-error:3.545e-01 | time:   0.7153 | largest rank:  16\n",
      "iter: 6  | tt-error: 3.123e-01, test-error:7.426e-03 | time:   0.9527 | largest rank:  19\n",
      "iter: 7  | tt-error: 5.737e-03, test-error:1.889e-15 | time:   1.2702 | largest rank:  22\n",
      "iter: 8  | tt-error: 2.039e-08, test-error:1.663e-15 | time:   1.6888 | largest rank:  25 <- converged: eps < 0.001\n",
      "Did 2543400 function evaluations, which took 1.37s (1.856e+06 evals/s)\n",
      "\n",
      "time taken:  1.7976365089416504\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",
    "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": "bd26218e-27c2-4c03-9e52-25f2c0640edd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_tt:  tensor(5.2491e-09)\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))**2).mean()\n",
    "print(\"mse_tt: \", mse_tt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "79525b5f",
   "metadata": {},
   "source": [
    "### Fit NN Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "66181567-4394-4706-9c8f-9fb526bd7aab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken:  3.4809112548828125e-05\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "# Fit NN\n",
    "lr= 1e-3\n",
    "batch_size = 128\n",
    "epochs = 1\n",
    "nn = NeuralNetwork(dim, width=64, lr=1e-3, device=device)\n",
    "nn.load_data(data_train, data_test)\n",
    "t1 = time.time()\n",
    "# nn.train(num_epochs=epochs, batch_size=batch_size, verbose=True)\n",
    "t2 = time.time()\n",
    "print(\"time taken: \", t2-t1)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7539baea-ecda-4204-ad3b-58b246262033",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_nn:  tensor(0.0167)\n"
     ]
    }
   ],
   "source": [
    "# Test the accuracy of NN over the test set\n",
    "y_nn = nn.model(x_test)\n",
    "mse_nn = ((y_nn.view(-1)-y_test.view(-1))**2).mean().detach()\n",
    "print(\"mse_nn: \", mse_nn)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "43eff033",
   "metadata": {},
   "source": [
    "### Fit BGMM Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "05c0a3bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mse_bgmm:  tensor(0.0002)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/idiap/temp/sshetty/miniconda/envs/pyml/lib/python3.9/site-packages/sklearn/mixture/_base.py:268: ConvergenceWarning: Initialization 1 did not converge. Try different init parameters, or increase max_iter, tol or check for degenerate data.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# Sample data and Train BGMM\n",
    "X_sample = gmm.generate_sample(x_train.shape[0]) # sample from reference distribution\n",
    "bgmm = BGMM(nmix=nmix)\n",
    "X_numpy = X_sample.detach().cpu().numpy()\n",
    "bgmm.load_data(X_numpy)\n",
    "bgmm.fit()\n",
    "\n",
    "# Test bgmm\n",
    "y_bgmm = bgmm.pdf(x_test.detach().cpu().numpy())\n",
    "y_test_numpy = y_test.detach().cpu().numpy()\n",
    "mse_bgmm = ((y_bgmm.reshape(-1)-y_test_numpy.reshape(-1))**2).mean()\n",
    "print(\"mse_bgmm: \", mse_bgmm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "65914903",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " mse_tt:5.249069685752172e-09,\n",
      " mse_bgmm:0.00016431352560326084,\n",
      " mse_nn:0.016652875400446403\n"
     ]
    }
   ],
   "source": [
    "print(f\" mse_tt:{mse_tt},\\n mse_bgmm:{mse_bgmm},\\n mse_nn:{mse_nn}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec3fb22",
   "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
}