{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# For tips on running notebooks in Google Colab, see\n# https://pytorch.org/tutorials/beginner/colab\n%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[Learn the Basics](intro.html) \\|\\| **Quickstart** \\|\\|\n[Tensors](tensorqs_tutorial.html) \\|\\| [Datasets &\nDataLoaders](data_tutorial.html) \\|\\|\n[Transforms](transforms_tutorial.html) \\|\\| [Build\nModel](buildmodel_tutorial.html) \\|\\|\n[Autograd](autogradqs_tutorial.html) \\|\\|\n[Optimization](optimization_tutorial.html) \\|\\| [Save & Load\nModel](saveloadrun_tutorial.html)\n\nQuickstart\n==========\n\nThis section runs through the API for common tasks in machine learning.\nRefer to the links in each section to dive deeper.\n\nWorking with data\n-----------------\n\nPyTorch has two [primitives to work with\ndata](https://pytorch.org/docs/stable/data.html):\n`torch.utils.data.DataLoader` and `torch.utils.data.Dataset`. `Dataset`\nstores the samples and their corresponding labels, and `DataLoader`\nwraps an iterable around the `Dataset`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import torch\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets\nfrom torchvision.transforms import ToTensor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "PyTorch offers domain-specific libraries such as\n[TorchText](https://pytorch.org/text/stable/index.html),\n[TorchVision](https://pytorch.org/vision/stable/index.html), and\n[TorchAudio](https://pytorch.org/audio/stable/index.html), all of which\ninclude datasets. For this tutorial, we will be using a TorchVision\ndataset.\n\nThe `torchvision.datasets` module contains `Dataset` objects for many\nreal-world vision data like CIFAR, COCO ([full list\nhere](https://pytorch.org/vision/stable/datasets.html)). In this\ntutorial, we use the FashionMNIST dataset. Every TorchVision `Dataset`\nincludes two arguments: `transform` and `target_transform` to modify the\nsamples and labels respectively.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Download training data from open datasets.\ntraining_data = datasets.FashionMNIST(\n    root=\"data\",\n    train=True,\n    download=True,\n    transform=ToTensor(),\n)\n\n# Download test data from open datasets.\ntest_data = datasets.FashionMNIST(\n    root=\"data\",\n    train=False,\n    download=True,\n    transform=ToTensor(),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We pass the `Dataset` as an argument to `DataLoader`. This wraps an\niterable over our dataset, and supports automatic batching, sampling,\nshuffling and multiprocess data loading. Here we define a batch size of\n64, i.e. each element in the dataloader iterable will return a batch of\n64 features and labels.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "batch_size = 64\n\n# Create data loaders.\ntrain_dataloader = DataLoader(training_data, batch_size=batch_size)\ntest_dataloader = DataLoader(test_data, batch_size=batch_size)\n\nfor X, y in test_dataloader:\n    print(f\"Shape of X [N, C, H, W]: {X.shape}\")\n    print(f\"Shape of y: {y.shape} {y.dtype}\")\n    break"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read more about [loading data in PyTorch](data_tutorial.html).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Creating Models\n===============\n\nTo define a neural network in PyTorch, we create a class that inherits\nfrom\n[nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html).\nWe define the layers of the network in the `__init__` function and\nspecify how data will pass through the network in the `forward`\nfunction. To accelerate operations in the neural network, we move it to\nthe\n[accelerator](https://pytorch.org/docs/stable/torch.html#accelerators)\nsuch as CUDA, MPS, MTIA, or XPU. If the current accelerator is\navailable, we will use it. Otherwise, we use the CPU.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else \"cpu\"\nprint(f\"Using {device} device\")\n\n# Define model\nclass NeuralNetwork(nn.Module):\n    def __init__(self):\n        super().__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\nmodel = NeuralNetwork().to(device)\nprint(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read more about [building neural networks in\nPyTorch](buildmodel_tutorial.html).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Optimizing the Model Parameters\n===============================\n\nTo train a model, we need a [loss\nfunction](https://pytorch.org/docs/stable/nn.html#loss-functions) and an\n[optimizer](https://pytorch.org/docs/stable/optim.html).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "loss_fn = nn.CrossEntropyLoss()\noptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "In a single training loop, the model makes predictions on the training\ndataset (fed to it in batches), and backpropagates the prediction error\nto adjust the model\\'s parameters.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": 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        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n\n        if batch % 100 == 0:\n            loss, current = loss.item(), (batch + 1) * len(X)\n            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We also check the model\\'s performance against the test dataset to\nensure it is learning.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": 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": {},
      "source": [
        "The training process is conducted over several iterations (*epochs*).\nDuring each epoch, the model learns parameters to make better\npredictions. We print the model\\'s accuracy and loss at each epoch;\nwe\\'d like to see the accuracy increase and the loss decrease with every\nepoch.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "epochs = 5\nfor 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)\nprint(\"Done!\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read more about [Training your model](optimization_tutorial.html).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Saving Models\n=============\n\nA common way to save a model is to serialize the internal state\ndictionary (containing the model parameters).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "torch.save(model.state_dict(), \"model.pth\")\nprint(\"Saved PyTorch Model State to model.pth\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loading Models\n==============\n\nThe process for loading a model includes re-creating the model structure\nand loading the state dictionary into it.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = NeuralNetwork().to(device)\nmodel.load_state_dict(torch.load(\"model.pth\", weights_only=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This model can now be used to make predictions.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "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\nmodel.eval()\nx, y = test_data[0][0], test_data[0][1]\nwith torch.no_grad():\n    x = x.to(device)\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": {},
      "source": [
        "Read more about [Saving & Loading your\nmodel](saveloadrun_tutorial.html).\n"
      ]
    }
  ],
  "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}