PyTorch_Meetup.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation\n",
    "\n",
    "\n",
    "1. Install conda. For example in Ubuntu:\n",
    "<br/>\n",
    "\n",
    "```bash\n",
    "curl -O https://repo.anaconda.com/archive/Anaconda3-2018.12-Linux-x86_64.sh\n",
    "sh Anaconda3-2018.12-Linux-x86_64.sh\n",
    "```\n",
    "\n",
    "2. Install conda packages: \n",
    "<br/>\n",
    "\n",
    "```bash\n",
    "conda install pytorch-cpu torchvision-cpu -c pytorch\n",
    "```\n",
    "PyTorch conda package comes automatically with CUDA and cuDNN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create tensor\n",
    "\n",
    "Functions which create tensors in PyTorch are similar to the functions which create ndarray in numpy.  One can also create a torch.tensor from numpy.ndarray and convert torch.tensor to numpy.ndarray easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.empty(2, 3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.rand(2, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.zeros(2, 3, dtype=torch.long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor([5.5, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "a = np.ones((2,3))\n",
    "x = torch.from_numpy(a)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "torch.tensor always copies data. To avoid copy one can use torch.as_tensor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.ones((2,3))\n",
    "b = torch.tensor(x)\n",
    "x[0][0] = 3\n",
    "print(x)\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.ones((2,3))\n",
    "b = torch.as_tensor(x)\n",
    "x[0][0] = 3\n",
    "print(x)\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What is tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([5., 3.], requires_grad=True)\n",
    "z = sum(x + x)\n",
    "z.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (type(z), type(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each tensor consists of many attributes. Some of them are:\n",
    "* Data of the tensor, which is also tensor itself"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Parameter which shows if the tensor needs to compute gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.requires_grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Gradients of the tensor is of the same size as tensor.data or None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.grad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Function of computational graph which computes gradients during backward path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z.grad_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Parameter which shows if the tensor is a leaf of computational graph. The tensor is a leaf if it is created by one of the following methods: \n",
    "\n",
    "    * direct initialization of tensors\n",
    "    * any operations on tensors with require_grad=False\n",
    "    * .detach() function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z.is_leaf, x.is_leaf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([1.0, 2.0], requires_grad=True)\n",
    "y = torch.tensor([3.0, 4.0], requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.data, y.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.grad, y.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = torch.dot(x, y) * x\n",
    "print(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$\\textbf{x}=[x_1, x_2]$\n",
    "<br/>\n",
    "$\\textbf{y}=[y_1, y_2]$\n",
    "<br/>\n",
    "$\\textbf{z} = (x_1*y_1+ x_2*y_2)* [x_1, x_2]$\n",
    "<br/>\n",
    "Both elements are vectors, and gradient of one vector over another is a matrix (Jacobian):\n",
    "<br/>\n",
    "<br/>\n",
    "$\\dfrac{\\partial \\textbf{z}}{\\partial \\textbf{x}} =   \n",
    "\\left[ {\\begin{array}{cc}\n",
    "   \\dfrac{\\partial z_1}{\\partial x_1} & \\dfrac{\\partial z_1}{\\partial x_2} \\\\\n",
    "   \\dfrac{\\partial z_2}{\\partial x_1} & \\dfrac{\\partial z_2}{\\partial x_2}\\\\\n",
    "  \\end{array} } \\right] $\n",
    "<br/>\n",
    "Do we compute Jacobian during backpropagation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    z.backward()\n",
    "except RuntimeError as re:\n",
    "    print(\"RuntimeError:\", re)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The error says that the gradient can be computed only for scalar outputs. Lets make it scalar by summing $z_1$ and $z_2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = torch.dot(x, y) * x\n",
    "z = z.sum()\n",
    "print (z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z.backward()\n",
    "print (x.grad.data, y.grad.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "backward() function accepts vector as an argument. The transverse Jacobian is multiplied by this vector and the loss function became scalar: \n",
    "<br/>\n",
    "<br/>\n",
    "\\begin{equation}\n",
    "\\label{eq:Jacob}\n",
    "\\tag{1}\n",
    "\\dfrac{\\partial \\textbf{z}}{\\partial \\textbf{x}} =   \n",
    "\\left[ {\\begin{array}{cc}\n",
    "   \\dfrac{\\partial z_1}{\\partial x_1} & \\dfrac{\\partial z_1}{\\partial x_2} \\\\\n",
    "   \\dfrac{\\partial z_2}{\\partial x_1} & \\dfrac{\\partial z_2}{\\partial x_2}\\\\\n",
    "  \\end{array} } \\right]^T \n",
    "  \\times\n",
    "  \\left[ {\\begin{array}{c}\n",
    "   \\dfrac{\\partial l}{\\partial z_1} \\\\\n",
    "   \\dfrac{\\partial l}{\\partial z_2} \\\\\n",
    "  \\end{array} } \\right] = \n",
    "  \\left[ {\\begin{array}{c}\n",
    "   \\dfrac{\\partial l}{\\partial x_1} \\\\\n",
    "   \\dfrac{\\partial l}{\\partial x_2} \\\\\n",
    "  \\end{array} } \\right]\n",
    " \\end{equation}\n",
    "<br/>\n",
    "The following expression should give the same result as the former:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = torch.dot(x, y) * x\n",
    "v = torch.tensor([1., 1.])\n",
    "z.backward(v)\n",
    "print (x.grad.data, y.grad.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gradients are not the same! The reason is that gradients are always accumulated. We need to set gradients to zero before the next cast of backward() function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x.grad.zero_()\n",
    "y.grad.zero_()\n",
    "z = torch.dot(x, y) * x\n",
    "v = torch.tensor([1., 1.])\n",
    "z.backward(v)\n",
    "print (x.grad.data, y.grad.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LEARNING_RATE = 1e-2\n",
    "\n",
    "x = torch.tensor([1, 2, 3, 4, 5],\n",
    "                     dtype=torch.float32)\n",
    "y = torch.tensor([5, 6, 8, 9, 10],\n",
    "                     dtype=torch.float32)\n",
    "\n",
    "w = torch.tensor([1], requires_grad=True,\n",
    "                     dtype=torch.float32)\n",
    "b = torch.tensor([1], requires_grad=True,\n",
    "                     dtype=torch.float32)\n",
    "\n",
    "for i in range(1000):\n",
    "    y_pred = w*x + b\n",
    "    z = sum((y_pred - y)**2)\n",
    "    z.backward()\n",
    "    w.data -= LEARNING_RATE * w.grad.data\n",
    "    b.data -= LEARNING_RATE * b.grad.data\n",
    "    w.grad.data.zero_()\n",
    "    b.grad.data.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x.numpy(), y.numpy(), 'o')\n",
    "plt.plot(x.numpy(), y_pred.detach().numpy(), '-')\n",
    "plt.show()\n",
    "print(\"w =\", w.data.numpy()[0],\"; b =\", b.data.numpy()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use optimizer\n",
    "\n",
    "One can use optimizers from torch.optim instead of updating weights manually. Optimizers accepts weights with gradients as parameter and then make step() which is updating weights of model based on the gradients. In this case the training will look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor([1, 2, 3, 4, 5],\n",
    "                     dtype=torch.float32)\n",
    "y = torch.tensor([5, 6, 8, 9, 10],\n",
    "                     dtype=torch.float32)\n",
    "\n",
    "w = torch.tensor([1], requires_grad=True,\n",
    "                     dtype=torch.float32)\n",
    "b = torch.tensor([1], requires_grad=True,\n",
    "                     dtype=torch.float32)\n",
    "\n",
    "optimizer = torch.optim.SGD([w,b],\n",
    "                            lr=1e-2)\n",
    "\n",
    "for i in range(1000):\n",
    "    y_pred = w*x + b\n",
    "    z = y_pred - y\n",
    "    z.backward(gradient=2*z)\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, that here we do not compute mean square. Instead we use vector $v = [2 z_1, 2 z_2]$ to reduce Jacobian to vector according to equation (1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x.numpy(), y.numpy(), 'o')\n",
    "plt.plot(x.numpy(), y_pred.detach().numpy(), '-')\n",
    "plt.show()\n",
    "\n",
    "print(\"w =\", w.data.numpy()[0],\"; b =\", b.data.numpy()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All optimizers accept forward and backward pass function as an argument. Though, SGD and Adam optimizers will work without it, for some of optimizers the closure() function is required. The creation of computational graph as well as backward pass should be given in step() function as a parameter. For example, the former construction will not work with LBFGS optimizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    optimizer = torch.optim.LBFGS([w,b], lr=LEARNING_RATE)\n",
    "\n",
    "    for i in range(1000):\n",
    "        y_pred = w*x + b\n",
    "        z = y_pred - y\n",
    "        z.backward(gradient=2*z) # << d sum((dy)**2)\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "except TypeError as te:\n",
    "    print (\"TypeError:\", te)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So one needs to pass the closure to optimizer inside the loop. We will also use tqdm tool to track the progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tt(*x): return torch.tensor(x, dtype=torch.float32)\n",
    "def tw(*x): return torch.tensor(x, requires_grad=True, dtype=torch.float32)\n",
    "\n",
    "x = tt(1, 2, 3, 4, 5)\n",
    "y = tt(5, 6, 8, 9, 10)\n",
    "w = tw(1)\n",
    "b = tw(1)\n",
    "\n",
    "optimizer = torch.optim.LBFGS([w,b], lr=LEARNING_RATE)\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from time import sleep\n",
    "def closure():\n",
    "    optimizer.zero_grad()\n",
    "    y_pred = w*x + b\n",
    "    z = sum((y_pred - y)**2)\n",
    "    z.backward()\n",
    "    return z\n",
    "\n",
    "#t = trange(50, leave=False)\n",
    "t = tqdm(range(50), desc='Training Loss', leave=True)\n",
    "for i in t:\n",
    "    z = optimizer.step(closure)\n",
    "    t.set_description('Training Loss: %.2g' % z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(x.numpy(), y.numpy(), 'o')\n",
    "plt.plot(x.numpy(), y_pred.detach().numpy(), '-')\n",
    "plt.show()\n",
    "\n",
    "print(\"w =\", w.data.numpy()[0],\"; b =\", b.data.numpy()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quality control of Tic-Tac pills production"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameters\n",
    "VALIDATION_SPLIT = 0.2\n",
    "INPUT_SIZE = 1\n",
    "NUM_EPOCHS = 2\n",
    "BATCH_SIZE = 128\n",
    "LEARNING_RATE = 1e-4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a model\n",
    "\n",
    "We will use pre-trained neural network, AlexNet, which was developed to classify images in ImageNet contest. The torchvision library provides pretrained models which are easy to use. The model weights will be downloaded automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alexnet = torchvision.models.alexnet(pretrained=True)\n",
    "print(alexnet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model consists of two parts: \"features\" , which extracts features from image and \"classifier\" which classify images based on the extracted features. The classifier is created to classify between 1000 different classes. To adopt AlexNet to our dataset, we firstly fix parameters of \"features\" part, and secondly will modify classifier for classification between two classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do not compute gradients for features parameters\n",
    "for param in alexnet.features.parameters():\n",
    "    param.require_grad = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace the model classifier\n",
    "alexnet.classifier = nn.Sequential(*[nn.Dropout(p=0.5),\n",
    "                                     nn.Linear(9216, 1000),\n",
    "                                     nn.ReLU(),\n",
    "                                     nn.Linear(1000, 1),\n",
    "                                     nn.Sigmoid()])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset consists of about 40000 images 160x160 of valid Tic Tac pills and 1000 images of broken Tic-Tac pills. The dataset can be downloaded from https://goo.gl/CWmLWD .\n",
    "\n",
    "![alt text](pills.png)\n",
    "\n",
    "We first create transformations and then load dataset using predefinded transformations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create transformations\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "    normalize])\n",
    "\n",
    "# load dataset\n",
    "dataset = torchvision.datasets.ImageFolder('tictac_dataset/',\n",
    "                                            transform=data_transform)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The AlexNet is trained on the colored images which were normalized by each of 3 image chanels to have certain mean and standard deviation. We normalize pills that they will have the same values. Note that pills images are grayscale, though we still load them having 3 channels.\n",
    "<br/>\n",
    "<br/>\n",
    "On this stage, no images loaded in RAM. Instead, dataset stores classes ids and pathes to images:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset.class_to_idx)\n",
    "print (dataset.imgs[0:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To split dataset to training and test, we can use PyTorch utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = len(dataset)\n",
    "m = int(VALIDATION_SPLIT*n)\n",
    "Train, Test = torch.utils.data.random_split(dataset, [n - m, m])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also create function which return WeightedRandomSampler which balances classes of our dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_balance_sampler(dataset):\n",
    "    l = [t[1] for t in dataset]\n",
    "    num_defect = sum(l)\n",
    "    frac = num_defect/len(l)\n",
    "    weight = np.array([frac, 1-frac])\n",
    "    samples_weight = np.array([weight[t] for t in l])\n",
    "    return torch.utils.data.sampler.WeightedRandomSampler(samples_weight, len(samples_weight))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to create generators of batches. These generators can be iterated in the loop and will return BATCH_SIZE number of images with labels 0-correct 1-defect. Note that after this step no images are loaded in RAM, the loading of images happens batch-wise durich iteration of the generators. pin_memory argument allows to allocate memory on GPU to fasten transfer between RAM and GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_sampler = make_balance_sampler(Train)\n",
    "train_loader = torch.utils.data.DataLoader(Train, batch_size=BATCH_SIZE, \n",
    "                                           sampler=train_sampler, pin_memory=False)\n",
    "test_sampler = make_balance_sampler(Test)\n",
    "test_loader = torch.utils.data.DataLoader(Test, batch_size=BATCH_SIZE,\n",
    "                                          sampler=test_sampler, pin_memory=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CUDA in PyTorch "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTorch behaviour differs from  Keras. CUDA is not used by default, one needs to specify what is running on CUDA. Though, the way to do it is very simple. One needs to inialize tensor on CUDA, then all operations on this tensor will be performed on GPU. In the following string we specify how to use CUDA automatically if available:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda:0\" if use_cuda else \"cpu\")\n",
    "if use_cuda:\n",
    "    alexnet.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training of neural network\n",
    "We use Adam optimizer with initial small learning rate. We pass only alexnet.classifier.parameters(), as this is the part we want to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize weights and optimizer\n",
    "optimizer = torch.optim.Adam(alexnet.classifier.parameters(), lr=LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AlexNet is an instance of nn.Module. All models which are created as class by inhereting from nn.Module has two modes: \"train\" and \"eval\". The default is train=True. This parameters are needed for some layers, such as BatchNormalization or Dropout, which behave differently during evaluation and training. This way we set alexnet.classifier to training mode:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alexnet.classifier = alexnet.classifier.train(True)\n",
    "alexnet.features = alexnet.features.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training loop will look similar to linear regression. The only difference is that we need to push torch tensors to GPU. The key non_blocking is added for assynchronization of GPU, which makes training a little bit faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(alexnet.classifier.parameters(), lr=LEARNING_RATE)\n",
    "for epoch in tqdm(range(NUM_EPOCHS), desc='Epochs'): \n",
    "    loss_mini_batch = torch.tensor(0., device=device)\n",
    "    count = torch.tensor(0., device=device)\n",
    "    t = tqdm(enumerate(train_loader))\n",
    "    for i, (images, labels) in t:\n",
    "        if use_cuda:\n",
    "            labels = labels.view(-1,1).to(device, dtype=torch.float32,\n",
    "                                          non_blocking=True)\n",
    "            images = images.to(device, non_blocking=True)\n",
    "        else:\n",
    "            labels = labels.view(-1,1).to(dtype=torch.float32)\n",
    "        outputs = alexnet(images)\n",
    "        loss = nn.functional.binary_cross_entropy(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step() \n",
    "        optimizer.zero_grad()\n",
    "        loss_mini_batch += loss.data\n",
    "        count += torch.tensor(1., device=device)\n",
    "        t.set_description('Training Loss: %.2g' % (loss_mini_batch/count).data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Speed up CUDA computations\n",
    "In order to speed up training of the model, one can save images to the hard drive as cuda-tensors. That will improve speed of training in almost two times. One can use the following script to transfer images:\n",
    "\n",
    "```python\n",
    "from PIL import Image\n",
    "import os\n",
    "import multiprocessing as mp\n",
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "data_transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.ToTensor(),\n",
    "    normalize])\n",
    "\n",
    "\n",
    "def convert_jpg2bmp(x):\n",
    "    infname, outdir = x[0], x[1]\n",
    "    outname = os.path.join(os.path.basename(os.path.dirname(infname)),\n",
    "                           os.path.basename(infname))\n",
    "    outname = os.path.join(outdir, outname)\n",
    "    outname = os.path.splitext(outname)[0] + '.pth'\n",
    "    img = Image.open(infname)\n",
    "    tens = data_transform(img)\n",
    "    tens = tens.to(device='cuda')\n",
    "    torch.save(tens, outname)\n",
    "\n",
    "outdir = 'tictac_cuda'\n",
    "\n",
    "# get list of images\n",
    "lifiles = []\n",
    "liclasses =  [name for name in os.listdir(\"tictac_dataset\")]\n",
    "for cli in liclasses:\n",
    "    indir = os.path.join('tictac_dataset', cli)\n",
    "    lii = [os.path.join(indir, name) for name in os.listdir(indir) if name.endswith(\".jpg\")]\n",
    "    lifiles = lifiles + lii\n",
    "\n",
    "# create directories\n",
    "os.makedirs(outdir, exist_ok=True)\n",
    "for cli in liclasses:\n",
    "    diri = os.path.join(outdir, cli)\n",
    "    os.makedirs(diri, exist_ok=True)\n",
    "\n",
    "# transfer images to pickled tensors\n",
    "pool = mp.Pool(processes=6)\n",
    "liparams = [[i, outdir] for i in lifiles]\n",
    "pool.map(convert_jpg2bmp, liparams)\n",
    "pool.close()\n",
    "```\n",
    "\n",
    "Note that it will take about half an hour on 6 CPU to transform all images. Afterward, one can create DataLoaders again with new images and train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataset\n",
    "dataset = torchvision.datasets.DatasetFolder('tictac_cuda/', loader=torch.load, extensions='.pth')\n",
    "\n",
    "n = len(dataset)\n",
    "m = int(VALIDATION_SPLIT*n)\n",
    "Train, Test = torch.utils.data.random_split(dataset, [n - m, m])\n",
    "\n",
    "train_sampler = make_balance_sampler(Train)\n",
    "test_sampler = make_balance_sampler(Test)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(Train, batch_size=BATCH_SIZE, \n",
    "                                           sampler=train_sampler)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(Test, batch_size=BATCH_SIZE,\n",
    "                                          sampler=test_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm_notebook as tqdm\n",
    "\n",
    "for epoch in tqdm(range(NUM_EPOCHS), desc='Epochs'):\n",
    "    \n",
    "    loss_mini_batch = torch.tensor(0., device=device)\n",
    "    count = torch.tensor(0., device=device)\n",
    "    \n",
    "    t = tqdm(enumerate(train_loader))\n",
    "    for i, (images, labels) in t:\n",
    "        labels = labels.view(-1,1).to(device, dtype=torch.float32, non_blocking=True)\n",
    "        images = torch.cat([images]*3, dim=1)\n",
    "        outputs = alexnet(images)\n",
    "        loss = nn.functional.binary_cross_entropy(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step() \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        loss_mini_batch += loss.data\n",
    "        count += torch.tensor(1., device=device)\n",
    "        t.set_description('Training Loss: %.2f' % (loss_mini_batch/count).data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save and load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save({\n",
    "            'model_state_dict': alexnet.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            }, \"alexnetstate_1.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is worth to check if the state of optimizer is saved properly by loading it after saving:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# need to load first model and optimizer again\n",
    "alexnet = torchvision.models.alexnet(pretrained=True)\n",
    "alexnet.classifier = alexnet.classifier.train(True)\n",
    "alexnet.features = alexnet.features.eval()\n",
    "for param in alexnet.features.parameters():\n",
    "    param.require_grad = False\n",
    "alexnet.classifier = nn.Sequential(*[nn.Dropout(p=0.5),\n",
    "                                     nn.Linear(9216, 1000),\n",
    "                                     nn.ReLU(),\n",
    "                                     nn.Linear(1000, 1),\n",
    "                                     nn.Sigmoid()])\n",
    "if use_cuda:\n",
    "    alexnet.cuda()\n",
    "optimizer = torch.optim.Adam(alexnet.classifier.parameters())\n",
    "\n",
    "# then load stat of model and optimizer\n",
    "checkpoint = torch.load(\"alexnetstate_1.pth\")\n",
    "alexnet.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the Model\n",
    "li_proba = []\n",
    "li_labels = []\n",
    "alexnet = alexnet.eval()\n",
    "for images, labels in test_loader:\n",
    "    labels = labels.to(device, dtype=torch.float32, non_blocking=True)\n",
    "    images = torch.cat([images]*3, dim=1)\n",
    "    outputs = alexnet(images)\n",
    "    proba = outputs.data.cpu().numpy()\n",
    "    li_proba.extend(proba)\n",
    "    li_labels.extend(labels.data.cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "from sklearn.metrics import precision_recall_curve\n",
    "import pylab as plt\n",
    "# confusion matrix\n",
    "li_pred = np.where(np.array(li_proba)>0.6, 1, 0)\n",
    "confusion_matrix(li_labels, li_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw precision-recall curve\n",
    "precision, recall, thr = precision_recall_curve(li_labels, li_proba)\n",
    "\n",
    "plt.plot(recall, precision)\n",
    "plt.xlabel('recall')\n",
    "plt.ylabel('precision')\n",
    "plt.show()"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}