runs/Mar25_20-32-24_nkd5rwp4qf/1648240347.1094823/quick_start_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "a4090294-3349-4815-96f4-98010b657359",
     "kernelId": ""
    }
   },
   "source": [
    "# Paperspace Gradient: PyTorch Quick Start\n",
    "Last modified: Nov 18th 2021"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "4936c59a-8535-43cf-a527-e9323b2b658e",
     "kernelId": ""
    }
   },
   "source": [
    "## Purpose and intended audience\n",
    "\n",
    "This Quick Start tutorial demonstrates PyTorch usage in a Gradient Notebook. It is aimed at users who are relatviely new to PyTorch, although you will need to be familiar with Python to understand PyTorch code.\n",
    "\n",
    "We use PyTorch to\n",
    "\n",
    "- Build a neural network that classifies FashionMNIST images\n",
    "- Train and evaluate the network\n",
    "- Save the model\n",
    "- Perform predictions\n",
    "\n",
    "followed by some next steps that you can take to proceed with using Gradient.\n",
    "\n",
    "The material is based on the original [PyTorch Quick Start](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html).\n",
    "\n",
    "See the end of the notebook for the original copyright notice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "a55c3131-9437-483d-9c19-a165fbf8b6d4",
     "kernelId": ""
    }
   },
   "source": [
    "## Check that you are on a GPU instance\n",
    "\n",
    "The notebook is designed to run on a Gradient GPU instance (as opposed to a CPU-only instance). The instance type, e.g., A4000, can be seen by clicking on the instance icon on the left-hand navigation bar in the Gradient Notebook interface. It will say if it is CPU or GPU.\n",
    "\n",
    "![instance_type.png](instance_type.png)\n",
    "\n",
    "The *Creating models* section below also determines whether or not a GPU is available for us to use.\n",
    "\n",
    "If the instance type is CPU, you can change it by clicking *Stop Instance*, then the instance type displayed to get a drop-down list. Select a GPU instance and start up the Notebook again.\n",
    "\n",
    "For help with instances, see the Gradient documentation on [instance types](https://docs.paperspace.com/gradient/more/instance-types) or [starting a Gradient Notebook](https://docs.paperspace.com/gradient/explore-train-deploy/notebooks)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "cd28b5e4-862f-4fc5-b02d-2335345647fa",
     "kernelId": ""
    }
   },
   "source": [
    "## Add ipywidgets\n",
    "This is temporary to enable PyTorch to work on Gradient notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 1,
     "id": "86ef45c8-089d-4d76-b919-99bccbd7edbb",
     "kernelId": "",
     "source_hidden": false
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Collecting ipywidgets\n",
      "  Downloading ipywidgets-7.6.5-py2.py3-none-any.whl (121 kB)\n",
      "\u001b[K     |████████████████████████████████| 121 kB 26.7 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: ipykernel>=4.5.1 in /opt/conda/lib/python3.8/site-packages (from ipywidgets) (6.4.1)\n",
      "Requirement already satisfied: nbformat>=4.2.0 in /opt/conda/lib/python3.8/site-packages (from ipywidgets) (5.1.3)\n",
      "Collecting jupyterlab-widgets>=1.0.0\n",
      "  Downloading jupyterlab_widgets-1.0.2-py3-none-any.whl (243 kB)\n",
      "\u001b[K     |████████████████████████████████| 243 kB 26.2 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: traitlets>=4.3.1 in /opt/conda/lib/python3.8/site-packages (from ipywidgets) (5.1.0)\n",
      "Requirement already satisfied: ipython-genutils~=0.2.0 in /opt/conda/lib/python3.8/site-packages (from ipywidgets) (0.2.0)\n",
      "Requirement already satisfied: ipython>=4.0.0 in /opt/conda/lib/python3.8/site-packages (from ipywidgets) (7.28.0)\n",
      "Collecting widgetsnbextension~=3.5.0\n",
      "  Downloading widgetsnbextension-3.5.2-py2.py3-none-any.whl (1.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 1.6 MB 27.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: jupyter-client<8.0 in /opt/conda/lib/python3.8/site-packages (from ipykernel>=4.5.1->ipywidgets) (7.0.6)\n",
      "Requirement already satisfied: debugpy<2.0,>=1.0.0 in /opt/conda/lib/python3.8/site-packages (from ipykernel>=4.5.1->ipywidgets) (1.5.0)\n",
      "Requirement already satisfied: matplotlib-inline<0.2.0,>=0.1.0 in /opt/conda/lib/python3.8/site-packages (from ipykernel>=4.5.1->ipywidgets) (0.1.3)\n",
      "Requirement already satisfied: tornado<7.0,>=4.2 in /opt/conda/lib/python3.8/site-packages (from ipykernel>=4.5.1->ipywidgets) (6.1)\n",
      "Requirement already satisfied: setuptools>=18.5 in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (58.2.0)\n",
      "Requirement already satisfied: pickleshare in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (0.7.5)\n",
      "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (3.0.20)\n",
      "Requirement already satisfied: decorator in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (5.1.0)\n",
      "Requirement already satisfied: pexpect>4.3 in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (4.8.0)\n",
      "Requirement already satisfied: pygments in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (2.10.0)\n",
      "Requirement already satisfied: jedi>=0.16 in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (0.18.0)\n",
      "Requirement already satisfied: backcall in /opt/conda/lib/python3.8/site-packages (from ipython>=4.0.0->ipywidgets) (0.2.0)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /opt/conda/lib/python3.8/site-packages (from jedi>=0.16->ipython>=4.0.0->ipywidgets) (0.8.2)\n",
      "Requirement already satisfied: nest-asyncio>=1.5 in /opt/conda/lib/python3.8/site-packages (from jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (1.5.1)\n",
      "Requirement already satisfied: jupyter-core>=4.6.0 in /opt/conda/lib/python3.8/site-packages (from jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (4.8.1)\n",
      "Requirement already satisfied: entrypoints in /opt/conda/lib/python3.8/site-packages (from jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (0.3)\n",
      "Requirement already satisfied: pyzmq>=13 in /opt/conda/lib/python3.8/site-packages (from jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (22.3.0)\n",
      "Requirement already satisfied: python-dateutil>=2.1 in /opt/conda/lib/python3.8/site-packages (from jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (2.8.2)\n",
      "Requirement already satisfied: jsonschema!=2.5.0,>=2.4 in /opt/conda/lib/python3.8/site-packages (from nbformat>=4.2.0->ipywidgets) (4.0.1)\n",
      "Requirement already satisfied: pyrsistent!=0.17.0,!=0.17.1,!=0.17.2,>=0.14.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets) (0.18.0)\n",
      "Requirement already satisfied: attrs>=17.4.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets) (21.2.0)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /opt/conda/lib/python3.8/site-packages (from pexpect>4.3->ipython>=4.0.0->ipywidgets) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /opt/conda/lib/python3.8/site-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=4.0.0->ipywidgets) (0.2.5)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/lib/python3.8/site-packages (from python-dateutil>=2.1->jupyter-client<8.0->ipykernel>=4.5.1->ipywidgets) (1.16.0)\n",
      "Requirement already satisfied: notebook>=4.4.1 in /opt/conda/lib/python3.8/site-packages (from widgetsnbextension~=3.5.0->ipywidgets) (6.4.1)\n",
      "Requirement already satisfied: argon2-cffi in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (21.1.0)\n",
      "Requirement already satisfied: nbconvert in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (6.2.0)\n",
      "Requirement already satisfied: terminado>=0.8.3 in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.12.1)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (3.0.1)\n",
      "Requirement already satisfied: prometheus-client in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.11.0)\n",
      "Requirement already satisfied: Send2Trash>=1.5.0 in /opt/conda/lib/python3.8/site-packages (from notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (1.8.0)\n",
      "Requirement already satisfied: cffi>=1.0.0 in /opt/conda/lib/python3.8/site-packages (from argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (1.14.6)\n",
      "Requirement already satisfied: pycparser in /opt/conda/lib/python3.8/site-packages (from cffi>=1.0.0->argon2-cffi->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (2.20)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.8/site-packages (from jinja2->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (2.0.1)\n",
      "Requirement already satisfied: nbclient<0.6.0,>=0.5.0 in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.5.4)\n",
      "Requirement already satisfied: bleach in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (4.1.0)\n",
      "Requirement already satisfied: pandocfilters>=1.4.1 in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (1.5.0)\n",
      "Requirement already satisfied: testpath in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.5.0)\n",
      "Requirement already satisfied: mistune<2,>=0.8.1 in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.8.4)\n",
      "Requirement already satisfied: defusedxml in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.7.1)\n",
      "Requirement already satisfied: jupyterlab-pygments in /opt/conda/lib/python3.8/site-packages (from nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.1.2)\n",
      "Requirement already satisfied: webencodings in /opt/conda/lib/python3.8/site-packages (from bleach->nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (0.5.1)\n",
      "Requirement already satisfied: packaging in /opt/conda/lib/python3.8/site-packages (from bleach->nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (21.0)\n",
      "Requirement already satisfied: pyparsing>=2.0.2 in /opt/conda/lib/python3.8/site-packages (from packaging->bleach->nbconvert->notebook>=4.4.1->widgetsnbextension~=3.5.0->ipywidgets) (2.4.7)\n",
      "Installing collected packages: widgetsnbextension, jupyterlab-widgets, ipywidgets\n",
      "Successfully installed ipywidgets-7.6.5 jupyterlab-widgets-1.0.2 widgetsnbextension-3.5.2\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install ipywidgets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "28402a66-a8c4-4672-9592-cc530b58d439",
     "kernelId": ""
    }
   },
   "source": [
    "## Working with data\n",
    "\n",
    "PyTorch has two [primitives to work with data](https://pytorch.org/docs/stable/data.html):\n",
    "``torch.utils.data.DataLoader`` and ``torch.utils.data.Dataset``.\n",
    "``Dataset`` stores the samples and their corresponding labels, and ``DataLoader`` wraps an iterable around\n",
    "the ``Dataset``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 2,
     "id": "2bab3caa-e156-4635-bc21-53031ebea60d",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor, Lambda, Compose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "0dfb0116-56cd-4795-bc5e-79baad627726",
     "kernelId": ""
    }
   },
   "source": [
    "PyTorch offers domain-specific libraries such as [TorchText](https://pytorch.org/text/stable/index.html),\n",
    "[TorchVision](https://pytorch.org/vision/stable/index.html), and [TorchAudio](https://pytorch.org/audio/stable/index.html),\n",
    "all of which include datasets. For this tutorial, we will be using a TorchVision dataset.\n",
    "\n",
    "The ``torchvision.datasets`` module contains ``Dataset`` objects for many real-world vision data like\n",
    "CIFAR, COCO ([full list here](https://pytorch.org/vision/stable/datasets.html)). In this tutorial, we\n",
    "use the FashionMNIST dataset. Every TorchVision ``Dataset`` includes two arguments: ``transform`` and\n",
    "``target_transform`` to modify the samples and labels respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 3,
     "id": "631deddf-30f0-45f1-84ab-e5f4c510c500",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Download training data from open datasets\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")\n",
    "\n",
    "# Download test data from open datasets\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "0ace6ebf-b493-4b75-9bfa-dc48bc676b21",
     "kernelId": ""
    }
   },
   "source": [
    "We pass the ``Dataset`` as an argument to ``DataLoader``. This wraps an iterable over our dataset, and supports\n",
    "automatic batching, sampling, shuffling and multiprocess data loading. Here we define a batch size of 64, i.e., each element\n",
    "in the dataloader iterable will return a batch of 64 features and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 4,
     "id": "8e65f970-dce8-460c-b5f2-9cbee0c14900",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of X [N, C, H, W]:  torch.Size([64, 1, 28, 28])\n",
      "Shape of y:  torch.Size([64]) torch.int64\n"
     ]
    }
   ],
   "source": [
    "batch_size = 64\n",
    "\n",
    "# Create data loaders\n",
    "train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
    "\n",
    "for X, y in test_dataloader:\n",
    "    print(\"Shape of X [N, C, H, W]: \", X.shape)\n",
    "    print(\"Shape of y: \", y.shape, y.dtype)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "f9d1b1f7-0850-4676-93b6-902f78be237d",
     "kernelId": ""
    }
   },
   "source": [
    "Read more about [loading data in PyTorch](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "d9cc95fe-194b-4a6f-b01d-91510dfcfb00",
     "kernelId": ""
    }
   },
   "source": [
    "## Creating models, including GPU\n",
    "\n",
    "To define a neural network in PyTorch, we create a class that inherits\n",
    "from [nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html). We define the layers of the network\n",
    "in the ``__init__`` function and specify how data will pass through the network in the ``forward`` function. To accelerate\n",
    "operations in the neural network, we move it to the GPU if available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 5,
     "id": "d58d5484-8ca0-4400-91c5-d0e71cf89c12",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device\n",
      "NeuralNetwork(\n",
      "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
      "  (linear_relu_stack): Sequential(\n",
      "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=512, out_features=512, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=512, out_features=10, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Get cpu or gpu device for training\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(\"Using {} device\".format(device))\n",
    "\n",
    "# Define model\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 10)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "7ee591d8-e529-481b-8107-e84454893bd2",
     "kernelId": ""
    }
   },
   "source": [
    "Read more about [building neural networks in PyTorch](https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "b6db5b4f-80b9-4f9e-8feb-76d0ef1e346f",
     "kernelId": ""
    }
   },
   "source": [
    "## Optimizing the model parameters\n",
    "\n",
    "To train a model, we need a [loss function](https://pytorch.org/docs/stable/nn.html#loss-functions)\n",
    "and an [optimizer](https://pytorch.org/docs/stable/optim.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 6,
     "id": "8c22a532-16e0-440d-888e-d879e5f53c7c",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "5efe3473-ecf7-411c-a13b-ba54f5c257a6",
     "kernelId": ""
    }
   },
   "source": [
    "In a single training loop, the model makes predictions on the training dataset (fed to it in batches), and\n",
    "backpropagates the prediction error to adjust the model's parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 7,
     "id": "3d1af6c1-299b-4572-902a-c5e52ce0a7d2",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def train(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "\n",
    "        # Compute prediction error\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "f86e28f0-bb94-4443-a673-f6d3461d4e94",
     "kernelId": ""
    }
   },
   "source": [
    "We also check the model's performance against the test dataset to ensure it is learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 8,
     "id": "112d81e3-cdf8-4b1e-afca-6344be54f5e5",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def test(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    test_loss, correct = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "4e366ecc-735f-42dd-b04e-a94816b94fd8",
     "kernelId": ""
    }
   },
   "source": [
    "The training process is conducted over several iterations (*epochs*). During each epoch, the model learns\n",
    "parameters to make better predictions. We print the model's accuracy and loss at each epoch; we'd like to see the\n",
    "accuracy increase and the loss decrease with every epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 9,
     "id": "50bf09d9-1318-43ef-92aa-6ee308fcafa1",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 2.303235  [    0/60000]\n",
      "loss: 2.289679  [ 6400/60000]\n",
      "loss: 2.273108  [12800/60000]\n",
      "loss: 2.267172  [19200/60000]\n",
      "loss: 2.248831  [25600/60000]\n",
      "loss: 2.225987  [32000/60000]\n",
      "loss: 2.227034  [38400/60000]\n",
      "loss: 2.194261  [44800/60000]\n",
      "loss: 2.190697  [51200/60000]\n",
      "loss: 2.161292  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 53.8%, Avg loss: 2.155593 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 2.169532  [    0/60000]\n",
      "loss: 2.153734  [ 6400/60000]\n",
      "loss: 2.097200  [12800/60000]\n",
      "loss: 2.113983  [19200/60000]\n",
      "loss: 2.057467  [25600/60000]\n",
      "loss: 2.015557  [32000/60000]\n",
      "loss: 2.031434  [38400/60000]\n",
      "loss: 1.952968  [44800/60000]\n",
      "loss: 1.957087  [51200/60000]\n",
      "loss: 1.897905  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 60.1%, Avg loss: 1.885614 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 1.924514  [    0/60000]\n",
      "loss: 1.886686  [ 6400/60000]\n",
      "loss: 1.767823  [12800/60000]\n",
      "loss: 1.810671  [19200/60000]\n",
      "loss: 1.700105  [25600/60000]\n",
      "loss: 1.668604  [32000/60000]\n",
      "loss: 1.677238  [38400/60000]\n",
      "loss: 1.577084  [44800/60000]\n",
      "loss: 1.603734  [51200/60000]\n",
      "loss: 1.514089  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 60.3%, Avg loss: 1.522196 \n",
      "\n",
      "Epoch 4\n",
      "-------------------------------\n",
      "loss: 1.592778  [    0/60000]\n",
      "loss: 1.553160  [ 6400/60000]\n",
      "loss: 1.404765  [12800/60000]\n",
      "loss: 1.476303  [19200/60000]\n",
      "loss: 1.357471  [25600/60000]\n",
      "loss: 1.362992  [32000/60000]\n",
      "loss: 1.364555  [38400/60000]\n",
      "loss: 1.289281  [44800/60000]\n",
      "loss: 1.328217  [51200/60000]\n",
      "loss: 1.238191  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 62.5%, Avg loss: 1.260456 \n",
      "\n",
      "Epoch 5\n",
      "-------------------------------\n",
      "loss: 1.338341  [    0/60000]\n",
      "loss: 1.316752  [ 6400/60000]\n",
      "loss: 1.157560  [12800/60000]\n",
      "loss: 1.258749  [19200/60000]\n",
      "loss: 1.131236  [25600/60000]\n",
      "loss: 1.164936  [32000/60000]\n",
      "loss: 1.173478  [38400/60000]\n",
      "loss: 1.111497  [44800/60000]\n",
      "loss: 1.156012  [51200/60000]\n",
      "loss: 1.079641  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 64.0%, Avg loss: 1.098095 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "epochs = 5\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train(train_dataloader, model, loss_fn, optimizer)\n",
    "    test(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "7bfc0721-ce35-4380-9d90-0f3f17bae210",
     "kernelId": ""
    }
   },
   "source": [
    "Read more about [Training your model](optimization_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "88e2d48b-f1c2-43b0-956d-673d31e777cc",
     "kernelId": ""
    }
   },
   "source": [
    "## Saving models\n",
    "\n",
    "A common way to save a model is to serialize the internal state dictionary (containing the model parameters)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 10,
     "id": "5674fda2-6f1d-447c-ac05-d21934c7fe6f",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saved PyTorch Model State to model.pth\n"
     ]
    }
   ],
   "source": [
    "torch.save(model.state_dict(), \"model.pth\")\n",
    "print(\"Saved PyTorch Model State to model.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "b1e15431-85cf-4788-aa7f-5c12d77f4ac3",
     "kernelId": ""
    }
   },
   "source": [
    "## Loading models\n",
    "\n",
    "The process for loading a model includes re-creating the model structure and loading\n",
    "the state dictionary into it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": false,
     "execution_count": 11,
     "id": "ee2271cf-5092-43ad-afed-b64d2e6aea2c",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = NeuralNetwork()\n",
    "model.load_state_dict(torch.load(\"model.pth\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "83cc12b8-fca2-4ea0-91f6-cdd8065d6164",
     "kernelId": ""
    }
   },
   "source": [
    "This model can now be used to make predictions.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "gradient": {
     "editing": true,
     "execution_count": 12,
     "id": "efed4977-824f-4816-91c0-05f4e10d8b54",
     "kernelId": ""
    },
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predicted: \"Ankle boot\", Actual: \"Ankle boot\"\n"
     ]
    }
   ],
   "source": [
    "classes = [\n",
    "    \"T-shirt/top\",\n",
    "    \"Trouser\",\n",
    "    \"Pullover\",\n",
    "    \"Dress\",\n",
    "    \"Coat\",\n",
    "    \"Sandal\",\n",
    "    \"Shirt\",\n",
    "    \"Sneaker\",\n",
    "    \"Bag\",\n",
    "    \"Ankle boot\",\n",
    "]\n",
    "\n",
    "model.eval()\n",
    "x, y = test_data[0][0], test_data[0][1]\n",
    "with torch.no_grad():\n",
    "    pred = model(x)\n",
    "    predicted, actual = classes[pred[0].argmax(0)], classes[y]\n",
    "    print(f'Predicted: \"{predicted}\", Actual: \"{actual}\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "0b064ce8-bacb-45c2-8ef3-3a45ff7ecd5a",
     "kernelId": ""
    }
   },
   "source": [
    "Read more about [Saving & Loading your model](saveloadrun_tutorial.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "379b3389-034a-4c17-a742-dd7c6a8281ce",
     "kernelId": ""
    }
   },
   "source": [
    "## Next steps\n",
    "\n",
    "To proceed with PyTorch in Gradient, you can:\n",
    "    \n",
    " - Look at other Gradient material, such as the [tutorials](https://docs.paperspace.com/gradient/get-started/tutorials-list), [ML Showcase](https://ml-showcase.paperspace.com), [blog](https://blog.paperspace.com), or [community](https://community.paperspace.com)\n",
    " - Try out further [PyTorch tutorials](https://pytorch.org/tutorials/beginner/basics/intro.html)\n",
    " - Start writing your own projects, using our [documentation](https://docs.paperspace.com/gradient) when needed\n",
    " \n",
    "If you get stuck or need help, [contact support](https://support.paperspace.com), and we will be happy to assist.\n",
    "\n",
    "Good luck!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "gradient": {
     "editing": false,
     "id": "a4d2e55f-6c65-48fe-a9e7-165931791ff2",
     "kernelId": ""
    }
   },
   "source": [
    "## Original PyTorch copyright notice\n",
    "\n",
    "© Copyright 2021, PyTorch."
   ]
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}