{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Digit recognition using TT neural networks\n",
    "\n",
    "The TT layer is applied to the MNIST dataset.\n",
    "\n",
    "Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch as tn\n",
    "import torch.nn as nn\n",
    "try: \n",
    "    import torchtt as tntt\n",
    "except:\n",
    "    print('Installing torchTT...')\n",
    "    %pip install git+https://github.com/ion-g-ion/torchTT\n",
    "    import torchtt as tntt\n",
    "from torch import optim\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "device = tn.device('cuda' if tn.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the dataset and store it to a subfolder 'data'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = datasets.MNIST(root = 'downloads', train = True, transform = ToTensor(), download = True)\n",
    "test_data = datasets.MNIST(root = 'downloads', train = False, transform = ToTensor())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create 2 dataloaders for the training set and the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader_train = tn.utils.data.DataLoader(train_data, batch_size=1000, shuffle=True, num_workers=10)\n",
    "dataloader_test = tn.utils.data.DataLoader(test_data, batch_size=100, shuffle=True, num_workers=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the neural network arhitecture. I contains 2 hidden TT layers (with RELU activation function) with a linear output layer. A sotmax is applied at the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class BasicTT(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.ttl1 = tntt.nn.LinearLayerTT([1,7,4,7,4], [8,10,10,10,10], [1,4,2,2,2,1])\n",
    "        self.ttl2 = tntt.nn.LinearLayerTT([8,10,10,10,10], [8,3,3,3,3], [1,2,2,2,2,1])\n",
    "        self.linear = nn.Linear(81*8, 10, dtype = tn.float32)\n",
    "        self.logsoftmax = nn.LogSoftmax(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.ttl1(x)\n",
    "        x = tn.relu(x)\n",
    "        x = self.ttl2(x)\n",
    "        x = tn.relu(x)\n",
    "        x = x.view(-1,81*8)\n",
    "        x = self.linear(x)\n",
    "        return self.logsoftmax(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instantiate the model and choose the optimizer and the loss function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BasicTT().to(device)\n",
    "loss_function = nn.CrossEntropyLoss()   \n",
    "optimizer = optim.Adam(model.parameters(), lr = 0.001)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start the training for 30 epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 30\n",
    " \n",
    "for epoch in range(n_epochs):\n",
    "    \n",
    "    for i,(input,label) in enumerate(dataloader_train):\n",
    "        \n",
    "        input = tn.reshape(input.to(device),[-1,1,7,4,7,4])\n",
    "        label = label.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        output = model(input)\n",
    "        loss = loss_function(output, label)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        print('Epoch %d/%d iteration %d/%d loss %e'%(epoch+1,n_epochs,i+1,len(dataloader_train),loss))\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the accuracy over the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_correct = 0\n",
    "n_total = 0\n",
    "for (input,label) in dataloader_test:\n",
    "    input = tn.reshape(input.to(device),[-1,1,7,4,7,4])\n",
    "        \n",
    "    output = model(input).cpu()\n",
    "    \n",
    "    n_correct += tn.sum(tn.max(output,1)[1] == label)   \n",
    "    \n",
    "    n_total += input.shape[0]\n",
    "    \n",
    "print('Test accuracy ',n_correct/n_total)\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "0bb18d59442045223691660a1f8e0079e69ab2417ab62beafa5ba9155c0a563f"
  },
  "kernelspec": {
   "display_name": "Python 3.6.13 ('pytorch2')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
