CV_14_NN_compression_practice_LAB3.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "AIMS_NN_compression_practice_LAB3.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Gedeon-m-gedus/Image_Processing/blob/master/CV_14_NN_compression_practice_LAB3.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gtDXorK6UI6p",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import argparse\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "from torch.optim.lr_scheduler import StepLR\n",
        "!pip install tensorly\n",
        "import tensorly as tl\n",
        "#Need to set Tensorly's backend as pytorch\n",
        "tl.set_backend('pytorch')\n",
        "from tensorly import random, tenalg\n",
        "import numpy as np\n",
        "np.random.seed(0)\n",
        "import copy \n",
        "import time\n",
        "def tic():\n",
        "  return time.time()\n",
        "def toc(t):\n",
        "  return time.time()-t\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PlTzu2GM8sxK",
        "colab_type": "text"
      },
      "source": [
        "#Background\n",
        "In recent researches, we have seen numerous successes with neural networks, epseicially deep learning methods. However, for these deep architectures, the number of parameters is often extremely big. This overparameterized structure, although often has state-of-the-art performance, also suffers from  long processing time and the need of expensive hardwares. Therefore, many efforts have been made into deriving deep learning models with less number of parameters but still remains similar performance. One interesting idea is to leverage the structure of matrix product operator (MPO) and represent the weights of the NN using this MPO format and directly do gradient descent w.r.t this structure. In this practice, we will explore how to train a fully connected neural network in its matrix product operator format. \\\\\n",
        "This practice is inspired by the paper: https://arxiv.org/abs/1509.06569"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I10qf-zR9pmY",
        "colab_type": "text"
      },
      "source": [
        "# Directly learning the MPO represented NN (MPO-NN)\n",
        "## Detour: What is MPO and how to contract a MPO with an input tensor?\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jQM9NFm3xtAs",
        "colab_type": "text"
      },
      "source": [
        "Matrix product operator (MPO) is essentially an MPS with one extra leg on each core. For example, in the following figure, we show a typical MPO structure of the tensor $\\mathcal{G} = [\\mathcal{G}_1, \\mathcal{G}_2, \\mathcal{G}_3, \\mathcal{G}_4]$, which is of size $c_1\\times c_2\\times c_3\\times c_4\\times d_1\\times d_2\\times d_3\\times d_4$. As a similar notion to MPS, the leg connecting each core representing the MPO rank. In this example, the rank of the MPO is a list of $[r_1, r_2, r_3]. $![MPO_strcurtre](https://docs.google.com/uc?export=download&id=1a9GRYOhrC8Zn_rcumigt8gc4EZkKJJuz)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DT2zecGPnjRu",
        "colab_type": "text"
      },
      "source": [
        "### 1. \n",
        "In this first problem, you will need to implement the MPO contraction operation between a MPO and an input tensor $\\mathcal{X}$. Assuming you are given a MPO $\\mathcal{G}$ which is a list of MPO cores, i.e. $[\\mathcal{G}_1, \\cdots \\mathcal{G}_m]$ and the MPO is of the rank $[r_1 \\cdots r_{m-1}]$. We also assume that the i-th core is of the shape $r_{i-1}\\times d_i\\times c_i\\times r_i$, where for the first core we let $r_{0}= 1$ and for the last core $r_m = 1$ (Note a leg with dimension 1 is equivalent to an empty leg), as shown in the following figure. \n",
        "\n",
        "![MPO_core](https://docs.google.com/uc?export=download&id=1OAWiLk3kCBarhzQ6-BhmFyabmZi65eSh) \n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Y0TN5k08-qS",
        "colab_type": "text"
      },
      "source": [
        "Now, given an input tensor $\\mathcal{X}\\in \\mathbb{R}^{n \\times d_1 \\times \\cdots \\times d_m}$, where n is the batch size, implement the contraction between the corresponding legs such that the contracted tensor will be of size $n \\times c_1\\times \\cdots \\times c_m$. As an example, for an MPO with 4 cores, the contraction between it and an input tensor $\\mathcal{X}$ is shown as the following:\n",
        "\n",
        "![MPO_strcurtre](https://docs.google.com/uc?export=download&id=1pOQn5YeqfwWWD7mIemC3sYjecp59IYue) "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SMTW4PnF_-Ti",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def MPO_contraction(mpo, x):\n",
        "  '''\n",
        "  mpo: a list of MPO cores (tensorly format), the i-th core is of shape (r_{i-1}, d_i, c_i, r_i), \n",
        "       for the first core r_0 =1, for the last core r_m = 1.\n",
        "  x: input tensor to be contracted (tensorly format), should be of size (n, d_1, d_2, ..., d_m).\n",
        "  return: a tensor of shape (n, c_1, c_2, ..., c_m)\n",
        "  '''\n",
        "  return output_tmp.squeeze()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_xwmReDKA97n",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "results = [[410.3555, 367.0687, 394.2908, 358.6689, 381.4024, 341.9384, 360.2169,\n",
        "        329.8646, 407.1425, 363.4720, 393.6679, 355.0994, 390.3836, 349.0190,\n",
        "        371.2960, 335.3856, 381.8200, 349.6505, 372.9323, 343.3012, 366.3085,\n",
        "        336.4729, 351.3506, 325.3335], \n",
        "       [356.8043, 418.2853, 361.2211, 410.5959, 316.2948, 370.2491, 312.2206,\n",
        "        352.7069, 382.1384, 444.2850, 380.6153, 436.5235, 355.7083, 413.0649,\n",
        "        344.6199, 393.3047, 396.4348, 455.7986, 393.3049, 447.2321, 361.2296,\n",
        "        414.9194, 348.0826, 394.7057],\n",
        "        [354.9081, 293.2105, 359.8825, 296.4861, 377.5464, 315.0949, 377.2471,\n",
        "        312.9066, 370.3658, 318.1749, 367.7635, 312.0098, 404.4339, 350.1533,\n",
        "        393.8055, 334.7884, 376.3939, 327.9788, 379.5774, 330.4409, 408.0018,\n",
        "        352.9874, 404.4712, 350.2234],\n",
        "        [318.2197, 296.4062, 346.0204, 326.2782, 296.7090, 278.4173, 326.6268,\n",
        "        307.8708, 329.1500, 309.8562, 356.8482, 340.7281, 308.3035, 291.0971,\n",
        "        335.1161, 318.2343, 336.9507, 307.0458, 366.0646, 340.7538, 315.1257,\n",
        "        290.0126, 346.2084, 322.9559],\n",
        "        [318.1615, 349.5595, 331.5495, 369.5155, 328.6678, 364.3328, 342.3642,\n",
        "        382.9929, 330.7670, 360.3952, 347.7114, 379.8224, 354.8564, 389.9683,\n",
        "        370.4451, 407.6407, 355.8427, 385.7378, 372.6829, 404.4416, 379.2988,\n",
        "        412.8916, 393.8806, 429.8946],\n",
        "        [270.0208, 350.8890, 262.4508, 333.1976, 293.0804, 382.3908, 288.4987,\n",
        "        368.0071, 282.7515, 381.6471, 271.3544, 362.4180, 303.4720, 409.4315,\n",
        "        296.4285, 395.9372, 333.2086, 421.1110, 319.1843, 399.1218, 358.9009,\n",
        "        454.4263, 350.2224, 437.3429],\n",
        "        [352.3250, 368.0678, 358.6050, 376.5540, 373.9196, 390.7939, 386.3511,\n",
        "        400.3182, 359.8570, 371.7061, 368.3690, 378.2644, 375.4347, 388.3662,\n",
        "        389.0728, 397.3719, 391.6757, 409.0408, 401.8670, 415.1198, 410.2174,\n",
        "        428.3144, 426.0149, 437.9070],\n",
        "        [389.6984, 366.2637, 408.0897, 381.6472, 333.4066, 313.0393, 353.7299,\n",
        "        331.9791, 415.8656, 384.7394, 432.8951, 405.1320, 356.6425, 330.8488,\n",
        "        379.1675, 357.9731, 456.8054, 414.1855, 482.8067, 437.2590, 393.1959,\n",
        "        357.1537, 421.8253, 384.6655],\n",
        "        [343.7931, 335.9971, 359.5031, 351.5452, 298.8585, 291.3354, 321.7639,\n",
        "        311.5385, 319.9868, 315.0025, 331.4951, 329.4894, 281.8465, 278.9174,\n",
        "        299.7816, 297.7877, 338.9509, 337.0882, 350.7429, 348.8928, 305.9270,\n",
        "        303.5083, 326.6728, 323.9245],\n",
        "        [542.8569, 527.5298, 541.1932, 531.6619, 430.4786, 419.5898, 421.7222,\n",
        "        409.8663, 523.4540, 494.5858, 518.2686, 499.6876, 395.8323, 376.9006,\n",
        "        383.6669, 367.5051, 512.1805, 493.6794, 503.1376, 493.2400, 402.5128,\n",
        "        393.7093, 390.9211, 382.5303] ]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RYEElL8OEd9N",
        "colab_type": "text"
      },
      "source": [
        "You can run the following program to make sure the output of your function is correct. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g1N5vsYOBAfG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#Test for MPO_contraction\n",
        "np.random.seed(0)\n",
        "for k in range(10):\n",
        "  x = tl.tensor(np.random.rand(3, 4, 5, 6))\n",
        "  mpo = []\n",
        "  rank = 7\n",
        "  out_dim = 2\n",
        "  for i in range(1, x.ndim):\n",
        "    if i == 1:\n",
        "      dim_0 = 1 \n",
        "    else:\n",
        "      dim_0 = rank\n",
        "    if i == x.ndim - 1:\n",
        "      dim_3 = 1\n",
        "    else:\n",
        "      dim_3 = rank\n",
        "    dim_1 = x.shape[i]\n",
        "    dim_2 = out_dim\n",
        "    tmp = tl.tensor(np.random.rand(dim_0, dim_1, dim_2, dim_3))\n",
        "    mpo.append(tmp)\n",
        "  out = tl.to_numpy(tl.tensor_to_vec(MPO_contraction(mpo, x))).round(decimals = 4)\n",
        "  results[k] = np.asarray(results[k])\n",
        "  if ((out - results[k]) < 1e-3).all():\n",
        "    print('Test '+str(k)+' passed!')\n",
        "  else:\n",
        "    print('Test '+str(k)+' failed!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BcKhP8gYBFtZ",
        "colab_type": "text"
      },
      "source": [
        "## Define the MPO-NN stucture the forward pass\n",
        "### 2. \n",
        "From the structure of MPO we can see that an MPO essentially has two set of legs: using the previous notations, one set is $\\{c_1,\\cdots c_m\\}$, the other one is $\\{d_1,\\cdots d_n\\}$. In the above problem, we implemented the contraction between the input tensor and one set of legs of the MPO. The resulted tensor is of shape $(n, c_1, \\cdots, c_m)$. Note this can be seen as another input tensor and can be contracted with another MPO. One can see that this layer-by-layer structure is very similar to multi-layer perceptron (MLP). The idea of MPO neural networks, or the so-called tensorized neural networks, is to replace the fully-connected (or convolutional layer) weights (i.e. weight matrices) by an MPO structure, then use gradient descent to learn these MPO layers. This replacement will drastically reduce the number of parameters that the neural networks need. In this problem, we will be implementing MPO-NN to replace a fully connected neural networks. The dataset we use here is the classic MNist hand-written dataset, which contains pictures of size $28\\times 28$. \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hLk5fwtny94v",
        "colab_type": "text"
      },
      "source": [
        "The network structure is as the following:\n",
        "\n",
        "![MPO_NN](https://docs.google.com/uc?export=download&id=1lfy6lBzPj_t2ji2o8d56UKkiebonsTRK) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B-f2tNWL-l09",
        "colab_type": "text"
      },
      "source": [
        "# Prepare the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2NRCWJZW-hKQ",
        "colab_type": "text"
      },
      "source": [
        "In this practice, we will be using the classic Mnist dataset. The first step is of course to load the dataset from Pytorch. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IY4DEnur_ETz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "use_cuda = True # we will be using gpu to train the model, therefore setting use_cuda to True\n",
        "batch_size = 512\n",
        "test_batch_size = 2048\n",
        "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}\n",
        "# Load MNIST dataset from pytorch using dataloader\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('../data', train=True, download=True,\n",
        "                    transform=transforms.Compose([\n",
        "                        transforms.ToTensor(),\n",
        "                        transforms.Normalize((0.1307,), (0.3081,))\n",
        "                    ])),\n",
        "    batch_size=batch_size, shuffle=True, **kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
        "                        transforms.ToTensor(),\n",
        "                        transforms.Normalize((0.1307,), (0.3081,))\n",
        "                    ])),\n",
        "    batch_size=test_batch_size, shuffle=True, **kwargs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1t2JGY_R_Lvz",
        "colab_type": "text"
      },
      "source": [
        "Then let us create training_loader, validation_loader and test_loader for our data to facilitate our training process later on. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4ifXboVa_dFN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Use the first 50000 data for training and the rest for validation\n",
        "train_ds  = torch.utils.data.Subset(train_loader.dataset, range(0, 50000-1))\n",
        "train_ds = torch.utils.data.DataLoader(train_ds, batch_size=batch_size, shuffle=True, **kwargs)\n",
        "valid_ds = torch.utils.data.Subset(train_loader.dataset, range(50000, 60000-1))\n",
        "valid_ds = torch.utils.data.DataLoader(valid_ds, batch_size=batch_size, shuffle=True, **kwargs)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYMwzIg9Garw",
        "colab_type": "text"
      },
      "source": [
        "#### (a)\n",
        "In this problem, you need to implement the function `MPO_layer`, which takes a list of mpo ranks, a list of input tensor dimensions and a list of output tensor dimensions. The function outputs a list of MPO cores with corresponding dimensions determined by the input argument. For example, if `mpo_rank = [2,3]`, `input_dims = [4,5,6]`, `output_dims = [7,8,9]`, then the output of this function should be a list of 3 tensors, which are of shape `(1,4,7,2)`, `(2,5,8,3)` and `(3,6,9,1)`. Each core should be initialized with weights drawn from a uniform distribution $(-1/sqrt(a), 1/sqrt(a))$, where $a$ is the number of parameters for the corresponding core."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z6Td5Vo4XitE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "b343960f-05eb-45d7-85f7-28e7ec834429"
      },
      "source": [
        "def MPO_layer(mpo_rank, input_dims, output_dims, seed = 0, device = 'cuda'):\n",
        "  '''\n",
        "  mpo_rank: a list of ranks for the MPO\n",
        "  input_dims: a list of input dimensions of the MPO layer, corresponding to [d_1,d_2, ...] in the above figures\n",
        "  output_dims: a list of output dimensions of the MPO layer, corresponding to [c_1,c_2, ...] in the above figures\n",
        "  '''\n",
        "  np.random.seed(0)\n",
        "  mpo_layer = []\n",
        "  for i in range(len(input_dims)):\n",
        "    '''\n",
        "    Insert here to finish initializing `Core`\n",
        "    '''\n",
        "    mpo_layer.append(tl.tensor(core, device = device, requires_grad=True))\n",
        "  return mpo_layer"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Qw8OYiNvmNp",
        "colab_type": "text"
      },
      "source": [
        "#### (b)\n",
        "Finish writing the class `Net_MPO`, you will need to implement the `__init__` function and the `forward` function. In `__init__` you will need to define the network structure using the function `MPO_layer` you have just implemented in the previous question. \n",
        "\n",
        "For `__init__`, you will need to write the strcutre of the MPO-NN. For this particular problem, the network you will be implementing has two MPO layers and then a fully  connected layer to connect to the output. After each MPO layers, there are two `ReLu` layer, which takes the intermediate output and pass it through a `ReLu` function. You will need to construct a list of MPO layers, which is called `self.mpo_layers`, this should contain 2 lists, each of the list correspond to the MPO cores for that particular MPO layer. The function takes 5 arguements, `number_layers` indicates how many layers you want for the MPO-NN; `mpo_rank` is a list of MPO ranks for each MPO layer. For example, if there are 2 MPO layers, then `mpo_rank` is a list of two lists, where each list contains the MPO ranks for the respective MPO layer. `hidden_dims`  is also in this format, it is a list of all hidden layer dimensions. These hidden layer dimensions correspond to the output tensor shape of each MPO layer that you have implemented in the `MPO_layer` function. `input_shape` is the input tensor's shape and finally `output_dim` is the output dimension. For this particular dataset, the `output_dim` is 10 (10 classes in total). **Use the default parameters for your network structure and use the function you implemented in (a).**\n",
        " \n",
        "\n",
        "For `forward` function, it only takes the input tensor of size $(n, 28, 28)$ as argument, where $n$ is the batch size and $28\\times 28$ is the input image's dimension. The output shape of this function should be $n \\times 10$. You will first need to reshpe the input tensor to the input_shape defined in the `__init__` function. Then you will need to write a forward pass of the desginated MPO networks, which contains two contractions with the MPO layers and use `ReLu` function to activate after each MPO contraction and finally multiply by the last fully connected layer. **Use the function you implemented in (1)**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7nbxXQjhd3Y0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Net_MPO(nn.Module):\n",
        "  def __init__(self, number_layers = 2, mpo_ranks = [[5, 5, 5], [5,5,5]], \n",
        "               input_shape = [4, 7, 4, 7], output_dim = 10, \n",
        "               hidden_dims = [[5, 5, 5, 5], [5, 5, 5, 5]], seed = 0, device = 'cuda'):\n",
        "    \n",
        "    '''\n",
        "      mpo_rank: a list of lists of ranks for each MPO layer\n",
        "      input_shape: a list of input dimensions of the MPO layer, corresponding to [d_1,d_2, ...] in the above figures\n",
        "      hidden_dims: a list of lists of intermidate output dimensions of each MPO layer, corresponding to [c_1,c_2, ...] in the above figures\n",
        "      output_dim: the dimension of the final output, set to 10 for this problem.\n",
        "    '''\n",
        "    super(Net_MPO, self).__init__()\n",
        "\n",
        "    # Construct MPO layers\n",
        "    self.mpo_layers = []\n",
        "    for i in range(number_layers):\n",
        "      '''\n",
        "      Insert here to construct your MPO layers, 'layer'\n",
        "      '''\n",
        "      self.mpo_layers.append(layer)\n",
        "    \n",
        "    \n",
        "\n",
        "    #Construct final fully-connected layer\n",
        "    fc_in = get_tensor_size_given_shape(hidden_dims[-1])\n",
        "    self.fc = nn.Linear(int(fc_in), output_dim).to(device)\n",
        "\n",
        "    #Some class properties that might come handy\n",
        "    self.output_dim = output_dim\n",
        "    self.input_shape = input_shape\n",
        "    self.device = device\n",
        "\n",
        "  def forward(self, x):\n",
        "    '''\n",
        "    Reshape the input x here.\n",
        "    '''\n",
        "    inter_out = tl.tensor(x_reshaped).to(self.device)\n",
        "\n",
        "    '''\n",
        "    Write your forward pass of the two MPO layers\n",
        "    '''\n",
        "    \n",
        "    #FC layer\n",
        "    inter_out = tl.unfold(inter_out, 0).to(self.device)\n",
        "    out = self.fc(inter_out)\n",
        "    out = F.log_softmax(out, dim=1)\n",
        "    return out\n",
        "\n",
        "  def get_number_parameters(self):\n",
        "    size = 0\n",
        "    for mpo_layer in self.mpo_layers:\n",
        "      for core in mpo_layer:\n",
        "        size += get_tensor_size_given_shape(core.shape)\n",
        "    fc_layer_shape = hidden_dims[-1]\n",
        "    fc_layer_shape.append(self.output_dim)\n",
        "    size += get_tensor_size_given_shape(fc_layer_shape)\n",
        "    return size\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IyDaqngRxDj8",
        "colab_type": "text"
      },
      "source": [
        "## **Run the following code to train your MPO-NN**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alUW8KV7Dmq0",
        "colab_type": "text"
      },
      "source": [
        "## Training MPO-NN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gl1PvS6qDEK2",
        "colab_type": "text"
      },
      "source": [
        "The training and testing methods for the MPO-NN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "soOnKuPZrjSk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#The train and test function of the Tensorized Network\n",
        "def train(model, device, train_loader, optimizer, epoch, log_interval = 10):\n",
        "    model.train()\n",
        "    error = []\n",
        "    for batch_idx, (data, target) in enumerate(train_loader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        output = model(data).to(device)\n",
        "        loss = F.nll_loss(output, target)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if batch_idx % log_interval == 0:\n",
        "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
        "                100. * batch_idx / len(train_loader), loss.item()))\n",
        "          \n",
        "        error.append(loss.item())\n",
        "    return sum(error)/len(error)\n",
        "\n",
        "\n",
        "def test(model, device, test_loader, return_acc = False):\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data).to(device)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss\n",
        "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
        "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
        "\n",
        "    test_loss /= len(test_loader.dataset)\n",
        "\n",
        "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
        "        test_loss, correct, len(test_loader.dataset),\n",
        "        100. * correct / len(test_loader.dataset)))\n",
        "    if return_acc:\n",
        "      return 100. * correct / len(test_loader.dataset)\n",
        "    else:\n",
        "      return test_loss\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TwxDLYgfD0ZJ",
        "colab_type": "text"
      },
      "source": [
        "Hyper-parameters setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1vZPycJ9Dzm1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "lr = 0.01\n",
        "save_model = True\n",
        "epochs = 10\n",
        "torch.manual_seed(1)\n",
        "gamma = 0.5\n",
        "log_interval = 100\n",
        "step_size = 3\n",
        "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
        "number_layers = 2\n",
        "mpo_ranks = [[5, 5, 5], [5,5,5]] \n",
        "input_shape = [4, 7, 4, 7]\n",
        "output_dim = 10\n",
        "hidden_dims = [[5, 5, 5, 5], [5, 5, 5, 5]]\n",
        "seed = 0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rXURvdCeD3qA",
        "colab_type": "text"
      },
      "source": [
        "Training process"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TWQKqognD4qC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "model_MPO = Net_MPO(number_layers = number_layers, mpo_ranks =mpo_ranks, \n",
        "               input_shape = input_shape, output_dim = output_dim, \n",
        "               hidden_dims = hidden_dims, seed = seed)\n",
        "params = model_MPO.mpo_layers[0] + model_MPO.mpo_layers[0]\n",
        "#print(model_MPO.get_number_parameters())\n",
        "params.append(model_MPO.fc.weight)\n",
        "\n",
        "optimizer = optim.Adam(params , lr=lr)\n",
        "# scheduler for automatic decaying the learning rate\n",
        "scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)\n",
        "                   \n",
        "train_loss_tt = []\n",
        "test_loss_tt = []\n",
        "\n",
        "# Training\n",
        "for epoch in range(1, epochs + 1):\n",
        "    train_loss_tt.append(train(model_MPO, device, train_ds, optimizer, epoch, log_interval))\n",
        "    test_loss_tt.append(test(model_MPO, device, valid_ds))\n",
        "    scheduler.step()\n",
        "\n",
        "# Testing\n",
        "t = tic()\n",
        "MPO_acc = test(model_MPO, device, test_loader, return_acc = True)\n",
        "time_MPO = toc(t)\n",
        "if save_model:\n",
        "    torch.save(model_MPO.state_dict(), \"mnist_mpo.pt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YrOdLPw1-CHx",
        "colab_type": "text"
      },
      "source": [
        "## Learning fully connected NN\n",
        "### FCNN structure\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CLDSQCVEUYxJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "class Net(nn.Module):\n",
        "    def __init__(self, hidden_units = 300, fc_weights = None):\n",
        "        super(Net, self).__init__()\n",
        "        self.hidden_units = hidden_units\n",
        "        if fc_weights is None:\n",
        "          self.fc1 = nn.Linear(28*28, hidden_units)\n",
        "          self.fc2 = nn.Linear(hidden_units, hidden_units)\n",
        "          self.fc3 = nn.Linear(hidden_units, 10)\n",
        "        else:\n",
        "          self.fc1 = nn.Linear(28*28, hidden_units)\n",
        "          self.fc2 = nn.Linear(hidden_units, hidden_units)\n",
        "          self.fc3 = nn.Linear(hidden_units, 10)\n",
        "          self.fc1.weight.data = torch.tensor(fc_weights[0])\n",
        "          self.fc2.weight.data = torch.tensor(fc_weights[1])\n",
        "          self.fc3.weight.data = torch.tensor(fc_weights[2])\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.flatten(x, 1)\n",
        "        x = self.fc1(x)\n",
        "        x = F.leaky_relu(x)\n",
        "        x = self.fc2(x)\n",
        "        x = F.leaky_relu(x)\n",
        "        x = self.fc3(x)\n",
        "        output = F.log_softmax(x, dim=1)\n",
        "        #print(output[0])\n",
        "        return output\n",
        "        \n",
        "    def get_number_parameters(self):\n",
        "      return 28*28 * self.hidden_units + self.hidden_units**2 + self.hidden_units * 10"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TJze8E1fFJbl",
        "colab_type": "text"
      },
      "source": [
        "### Hyper-parameters setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b2cFJXdkyUzQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "use_cuda = True\n",
        "batch_size = 512\n",
        "test_batch_size = 2048\n",
        "lr = 0.001\n",
        "save_model = True\n",
        "epochs = 10\n",
        "torch.manual_seed(1)\n",
        "gamma = 0.1\n",
        "log_interval = 100\n",
        "hidden_units = 625\n",
        "step_size = 5\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FR2GmmuaFR6d",
        "colab_type": "text"
      },
      "source": [
        "### Training process"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mWhoVp0FFRVg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model_fc = Net(hidden_units).to(device)\n",
        "optimizer = optim.Adam(model_fc.parameters(), lr=lr)\n",
        "scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)\n",
        "\n",
        "train_loss_fc = []\n",
        "test_loss_fc = []\n",
        "\n",
        "for epoch in range(1, epochs + 1):\n",
        "    train_loss_fc.append(train(model_fc, device, train_ds, optimizer, epoch, log_interval))\n",
        "    test_loss_fc.append(test(model_fc, device, valid_ds))\n",
        "    scheduler.step()\n",
        "\n",
        "t = tic()\n",
        "fc_acc = test(model_fc, device, test_loader, return_acc = True)\n",
        "time_fc = toc(t)\n",
        "if save_model:\n",
        "    torch.save(model_fc.state_dict(), \"mnist_fc.pt\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GK4D3T5K-rew",
        "colab_type": "text"
      },
      "source": [
        "## Training and testing loss for TNN and FCNN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7LDPqX5AqVdE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline  \n",
        "plt.plot(train_loss_tt, label = 'Training_loss_MPO')\n",
        "plt.plot(test_loss_tt, label = 'Test_loss_MPO')\n",
        "plt.plot(train_loss_fc, label = 'Training_loss_FC')\n",
        "plt.plot(test_loss_fc, label = 'Test_loss_FC')\n",
        "plt.title('Training Loss and Test Loss for MPO-NN')\n",
        "plt.legend()\n",
        "plt.xlabel('Epoch')\n",
        "plt.ylabel('Cross Entropy')\n",
        "plt.show()\n",
        "MPO_size = int(model_MPO.get_number_parameters())\n",
        "print('Number of Parameters (MPO): ', MPO_size)\n",
        "print('Accuracy (MPO): ', MPO_acc)\n",
        "print('Forward pass time (MPO): ', time_MPO)\n",
        "\n",
        "print('Number of Parameters (FC): ', model_fc.get_number_parameters())\n",
        "print('Accuracy (FC): ', fc_acc)\n",
        "print('Forward pass time (FC): ', time_fc)\n",
        "\n",
        "print('Compression Rate:', MPO_size/model_fc.get_number_parameters())\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vqrK4WrU-61R",
        "colab_type": "text"
      },
      "source": [
        "# Decompose a fullly connected NN\n",
        "Now that we have showed that by using MPO format to train the neural networks, it can compress a significant amount of parameters. However, another more straight-forward idea would be to use various tensor decomposition methods to decompose and reconstruct the fully-connected layer. In the following practice we will be exploring if this method is duable and if it is better than MPO-NN. \n",
        "## Functions for decomposing FCNN via CP, Tucker and TT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rXSFojWfQ9p1",
        "colab_type": "text"
      },
      "source": [
        "#### Data prepration"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LULAkvOnQ9IU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def get_newshape(weights, tensor_dim):\n",
        "  total_elements = weights.shape[0]*weights.shape[1]\n",
        "  remaining_elements = total_elements\n",
        "  newshape = []\n",
        "  while (remaining_elements/tensor_dim).is_integer():\n",
        "    newshape.append(tensor_dim)\n",
        "    if remaining_elements/tensor_dim - int(remaining_elements/tensor_dim) == 0:\n",
        "      remaining_elements /= tensor_dim\n",
        "    else:\n",
        "      print('Error getting newshape, cannot reshape the matrix into a tensor')\n",
        "      break\n",
        "  if remaining_elements != 1:\n",
        "    newshape.append(int(remaining_elements))\n",
        "  return newshape\n",
        "\n",
        "layers = [model_fc.fc1.weight.data.cpu().numpy(), model_fc.fc2.weight.data.cpu().numpy(), model_fc.fc3.weight.data.cpu().numpy()]\n",
        "\n",
        "layers_ori_shape = [layer.shape for layer in layers]\n",
        "\n",
        "tensor_dim = 5\n",
        "layers_tensor_shape = [get_newshape(layer, tensor_dim) for layer in layers]\n",
        "\n",
        "layers_tensor = [tl.tensor(layer.reshape(layers_tensor_shape[i])) for i, layer in enumerate(layers)]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o2-IkWhNPIvt",
        "colab_type": "text"
      },
      "source": [
        "### (a) \n",
        "Complete the following functions, where the input is the fulled connected layer's weights for FC-NN in its tensor format and the output should be the reconstructed tensor of corresponding tensor decomposition methods and the number of parameters of that decomposition. Note all the `weights` argument is already reshaped into particular tensor shapes and you can directly use tensorly's decomposition functions (e.g. `tensorly.decomposition.parafac` for CP decomposition and  `tensorly.kruskal_to_tensor` to recover the tensor). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SdJ6QOsN_C2V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tensorly.decomposition import parafac\n",
        "from tensorly.decomposition import tucker\n",
        "\n",
        "def get_cp_reconstruction(weights, cp_rank):\n",
        "  '''\n",
        "  weights: Layer weights in its tensor format (already reshaped)\n",
        "  cp_rank: integer, cp rank\n",
        "  '''\n",
        "\n",
        "  return cp_reconstruction, num_params\n",
        "  \n",
        "def get_tt_reconstruction(weights, rank = 10):\n",
        "  '''\n",
        "  weights: Layer weights in its tensor format (already reshaped)\n",
        "  rank: integer, maximum TT rank\n",
        "  '''\n",
        "\n",
        "  return mps_to_tensor(factors), num_params\n",
        "\n",
        "def get_tucker_reconstruction(weights, tucker_ranks):\n",
        "  '''\n",
        "  weights: Layer weights in its tensor format (already reshaped)\n",
        "  tucker_rank: a list of integer, tucker rank\n",
        "  '''\n",
        "\n",
        "  return tl.tucker_to_tensor((core, factors)), num_params\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zsOtbykXxjks",
        "colab_type": "text"
      },
      "source": [
        "## **Run the following code to check your implementations and results**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R4iWuasL_SW5",
        "colab_type": "text"
      },
      "source": [
        "## TT"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SVwIE6ui7CwE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#TT\n",
        "tt_ranks = [10, 20, 30, 40, 50]\n",
        "tt_acc = []\n",
        "num_parameters_TT = []\n",
        "for rank_lim in tt_ranks:\n",
        "  recon_layers = []\n",
        "  num_param_total = 0\n",
        "  for layer in layers_tensor:\n",
        "    recon_layer, num_param_layer = get_tt_reconstruction(layer, rank_lim)\n",
        "    recon_layers.append(recon_layer)\n",
        "    num_param_total += num_param_layer\n",
        "  num_parameters_TT.append(num_param_total)\n",
        "  recon_model = Net(hidden_units, [layer.reshape(layers_ori_shape[i]) for i, layer in enumerate(recon_layers)])\n",
        "  recon_model.cuda()\n",
        "  print('Current TT rank is :' +str(rank_lim))\n",
        "  tt_acc.append(test(recon_model, device, test_loader, return_acc = True))\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_A2opE7A_WAv",
        "colab_type": "text"
      },
      "source": [
        "## Tucker"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KeVUrTZer-Gf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gc\n",
        "gc.collect()\n",
        "#Tucker\n",
        "Tucker_rank_list = [2, 3, 4]\n",
        "tucker_acc = []\n",
        "num_parameters_tucker = []\n",
        "for tucker_rank in Tucker_rank_list:\n",
        "  print('Current Tucker rank is :' +str(tucker_rank))\n",
        "  recon_layers = []\n",
        "  num_param_total = 0\n",
        "  for layer in layers_tensor:\n",
        "    tucker_ranks = np.ones(tl.ndim(layer)).astype(int)*tucker_rank\n",
        "    recon_layer, num_param_layer = get_tucker_reconstruction(layer, tucker_ranks)\n",
        "    recon_layers.append(recon_layer)\n",
        "    num_param_total += num_param_layer\n",
        "  num_parameters_tucker.append(num_param_total)\n",
        "  recon_model = Net(hidden_units, [layer.reshape(layers_ori_shape[i]) for i, layer in enumerate(recon_layers)])\n",
        "  recon_model.cuda()\n",
        "  tucker_acc.append(test(recon_model, device, test_loader, return_acc = True))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C5wDzFO4_UJE",
        "colab_type": "text"
      },
      "source": [
        "## CP"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HsQ9O2_0BQpY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import gc\n",
        "gc.collect()\n",
        "#set up CP pararmeters\n",
        "cp_ranks = [10, 50, 100, 200, 300, 400, 500]\n",
        "cp_acc = []\n",
        "num_parameters_cp = []\n",
        "\n",
        "for cp_rank in cp_ranks:\n",
        "  recon_layers = []\n",
        "  num_param_total = 0\n",
        "  print('Current CP rank is :' +str(cp_rank))\n",
        "  for layer in layers_tensor:\n",
        "    recon_layer, num_param_layer = get_cp_reconstruction(layer, cp_rank)\n",
        "    recon_layers.append(recon_layer)\n",
        "    num_param_total += num_param_layer\n",
        "  num_parameters_cp.append(num_param_total)\n",
        "  recon_model = Net(hidden_units, [layer.reshape(layers_ori_shape[i]) for i, layer in enumerate(recon_layers)])\n",
        "  recon_model.cuda()\n",
        "  cp_acc.append(test(recon_model, device, test_loader, return_acc = True))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VZS3deeN_axd",
        "colab_type": "text"
      },
      "source": [
        "## Results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TER1Li_I4sMJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "plt.plot( num_parameters_cp, cp_acc, label = 'CP')\n",
        "plt.plot( num_parameters_TT, tt_acc, label = 'TT')\n",
        "plt.plot( num_parameters_tucker, tucker_acc, label = 'Tucker')\n",
        "plt.plot( MPO_size, MPO_acc, '-bx', label = 'MPO_NN')\n",
        "plt.legend()\n",
        "plt.xlabel('Number_parameters')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.title('Accuracy on different types of decomposition of the orginal NN')\n",
        "plt.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}