{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from torch.optim import Adam\n",
    "import torchvision.datasets as datasets # pip install torch vision if you dont have this\n",
    "from sklearn.metrics import classification_report\n",
    "from tamnun.core import TorchEstimator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load data using torch vision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=None)\n",
    "mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the data tensors for train and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X = mnist_trainset.data.reshape(-1, 28*28).numpy()\n",
    "train_y = mnist_trainset.targets.numpy()\n",
    "\n",
    "test_X = mnist_testset.data.reshape(-1, 28*28).numpy()\n",
    "test_y = mnist_testset.targets.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create simple linear classifier with 28x28 (the image size) as input and 10 classes as output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "module = nn.Linear(28*28, 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the tamnun estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "clf = TorchEstimator(module, optimizer=Adam(module.parameters(), lr=1e-4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### fit().predict()!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20:\n",
      "1874/1875 batch loss: 3.3726372718811035    avg loss: 6.262858227777481\n",
      "Epoch 2/20:\n",
      "1874/1875 batch loss: 0.9372385740280151   5 avg loss: 2.11575478618145\n",
      "Epoch 3/20:\n",
      "1874/1875 batch loss: 2.048184633255005 9    avg loss: 1.6549429415345192\n",
      "Epoch 4/20:\n",
      "1874/1875 batch loss: 0.0006771087646484375  avg loss: 1.4070107420285543\n",
      "Epoch 5/20:\n",
      "1874/1875 batch loss: 0.8950674533843994 6  avg loss: 1.234210046708584\n",
      "Epoch 6/20:\n",
      "1874/1875 batch loss: 1.1697653532028198     avg loss: 1.1113991823037466\n",
      "Epoch 7/20:\n",
      "1874/1875 batch loss: 0.6111218333244324     avg loss: 1.0094971430778503\n",
      "Epoch 8/20:\n",
      "1874/1875 batch loss: 0.008098363876342773  avg loss: 0.9214095135052999\n",
      "Epoch 9/20:\n",
      "1874/1875 batch loss: 1.3127366304397583     avg loss: 0.8733639629522959\n",
      "Epoch 10/20:\n",
      "1874/1875 batch loss: 0.2436842918395996    avg loss: 0.8153708161791166\n",
      "Epoch 11/20:\n",
      "1874/1875 batch loss: 0.3469241261482239  6 avg loss: 0.786324020353953\n",
      "Epoch 12/20:\n",
      "1874/1875 batch loss: 0.07478591799736023  6 avg loss: 0.7397853104670843\n",
      "Epoch 13/20:\n",
      "1874/1875 batch loss: 1.0269213914871216    avg loss: 0.7178852183063825\n",
      "Epoch 14/20:\n",
      "1874/1875 batch loss: 0.7970320582389832    avg loss: 0.6849127511382103\n",
      "Epoch 15/20:\n",
      "1874/1875 batch loss: 0.5232244729995728    avg loss: 0.6671388566295305\n",
      "Epoch 16/20:\n",
      "1874/1875 batch loss: 0.805496335029602     avg loss: 0.6524858958641688\n",
      "Epoch 17/20:\n",
      "1874/1875 batch loss: 0.460053026676178    7 avg loss: 0.6337650477568308\n",
      "Epoch 18/20:\n",
      "1874/1875 batch loss: 0.5216989517211914    avg loss: 0.6275319784998894\n",
      "Epoch 19/20:\n",
      "1874/1875 batch loss: 0.318534255027771 7    avg loss: 0.609422858651479\n",
      "Epoch 20/20:\n",
      "1874/1875 batch loss: 0.2152227759361267 7  avg loss: 0.6130135586182276\n"
     ]
    }
   ],
   "source": [
    "clf = clf.fit(train_X, train_y, epochs=20, batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.98      0.92      0.95       980\n",
      "           1       0.97      0.97      0.97      1135\n",
      "           2       0.91      0.87      0.89      1032\n",
      "           3       0.86      0.90      0.88      1010\n",
      "           4       0.90      0.94      0.92       982\n",
      "           5       0.78      0.89      0.83       892\n",
      "           6       0.96      0.90      0.93       958\n",
      "           7       0.97      0.81      0.88      1028\n",
      "           8       0.81      0.86      0.84       974\n",
      "           9       0.84      0.89      0.86      1009\n",
      "\n",
      "   micro avg       0.89      0.89      0.89     10000\n",
      "   macro avg       0.90      0.89      0.89     10000\n",
      "weighted avg       0.90      0.89      0.90     10000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "predicted = clf.predict(test_X)\n",
    "print(classification_report(test_y, predicted))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
