{
  "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) \\|\\|\n[Quickstart](quickstart_tutorial.html) \\|\\|\n[Tensors](tensorqs_tutorial.html) \\|\\| [Datasets &\nDataLoaders](data_tutorial.html) \\|\\|\n[Transforms](transforms_tutorial.html) \\|\\| **Build Model** \\|\\|\n[Autograd](autogradqs_tutorial.html) \\|\\|\n[Optimization](optimization_tutorial.html) \\|\\| [Save & Load\nModel](saveloadrun_tutorial.html)\n\nBuild the Neural Network\n========================\n\nNeural networks comprise of layers/modules that perform operations on\ndata. The [torch.nn](https://pytorch.org/docs/stable/nn.html) namespace\nprovides all the building blocks you need to build your own neural\nnetwork. Every module in PyTorch subclasses the\n[nn.Module](https://pytorch.org/docs/stable/generated/torch.nn.Module.html).\nA neural network is a module itself that consists of other modules\n(layers). This nested structure allows for building and managing complex\narchitectures easily.\n\nIn the following sections, we\\'ll build a neural network to classify\nimages in the FashionMNIST dataset.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport torch\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets, transforms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get Device for Training\n=======================\n\nWe want to be able to train our model on an\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\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the Class\n================\n\nWe define our neural network by subclassing `nn.Module`, and initialize\nthe neural network layers in `__init__`. Every `nn.Module` subclass\nimplements the operations on input data in the `forward` method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class 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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We create an instance of `NeuralNetwork`, and move it to the `device`,\nand print its structure.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = NeuralNetwork().to(device)\nprint(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To use the model, we pass it the input data. This executes the model\\'s\n`forward`, along with some [background\noperations](https://github.com/pytorch/pytorch/blob/270111b7b611d174967ed204776985cefca9c144/torch/nn/modules/module.py#L866).\nDo not call `model.forward()` directly!\n\nCalling the model on the input returns a 2-dimensional tensor with dim=0\ncorresponding to each output of 10 raw predicted values for each class,\nand dim=1 corresponding to the individual values of each output. We get\nthe prediction probabilities by passing it through an instance of the\n`nn.Softmax` module.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "X = torch.rand(1, 28, 28, device=device)\nlogits = model(X)\npred_probab = nn.Softmax(dim=1)(logits)\ny_pred = pred_probab.argmax(1)\nprint(f\"Predicted class: {y_pred}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Model Layers\n============\n\nLet\\'s break down the layers in the FashionMNIST model. To illustrate\nit, we will take a sample minibatch of 3 images of size 28x28 and see\nwhat happens to it as we pass it through the network.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "input_image = torch.rand(3,28,28)\nprint(input_image.size())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "nn.Flatten\n==========\n\nWe initialize the\n[nn.Flatten](https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html)\nlayer to convert each 2D 28x28 image into a contiguous array of 784\npixel values ( the minibatch dimension (at dim=0) is maintained).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "flatten = nn.Flatten()\nflat_image = flatten(input_image)\nprint(flat_image.size())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "nn.Linear\n=========\n\nThe [linear\nlayer](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html)\nis a module that applies a linear transformation on the input using its\nstored weights and biases.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "layer1 = nn.Linear(in_features=28*28, out_features=20)\nhidden1 = layer1(flat_image)\nprint(hidden1.size())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "nn.ReLU\n=======\n\nNon-linear activations are what create the complex mappings between the\nmodel\\'s inputs and outputs. They are applied after linear\ntransformations to introduce *nonlinearity*, helping neural networks\nlearn a wide variety of phenomena.\n\nIn this model, we use\n[nn.ReLU](https://pytorch.org/docs/stable/generated/torch.nn.ReLU.html)\nbetween our linear layers, but there\\'s other activations to introduce\nnon-linearity in your model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(f\"Before ReLU: {hidden1}\\n\\n\")\nhidden1 = nn.ReLU()(hidden1)\nprint(f\"After ReLU: {hidden1}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "nn.Sequential\n=============\n\n[nn.Sequential](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html)\nis an ordered container of modules. The data is passed through all the\nmodules in the same order as defined. You can use sequential containers\nto put together a quick network like `seq_modules`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "seq_modules = nn.Sequential(\n    flatten,\n    layer1,\n    nn.ReLU(),\n    nn.Linear(20, 10)\n)\ninput_image = torch.rand(3,28,28)\nlogits = seq_modules(input_image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "nn.Softmax\n==========\n\nThe last linear layer of the neural network returns [logits]{.title-ref}\n- raw values in \\[-infty, infty\\] - which are passed to the\n[nn.Softmax](https://pytorch.org/docs/stable/generated/torch.nn.Softmax.html)\nmodule. The logits are scaled to values \\[0, 1\\] representing the\nmodel\\'s predicted probabilities for each class. `dim` parameter\nindicates the dimension along which the values must sum to 1.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "softmax = nn.Softmax(dim=1)\npred_probab = softmax(logits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Model Parameters\n================\n\nMany layers inside a neural network are *parameterized*, i.e. have\nassociated weights and biases that are optimized during training.\nSubclassing `nn.Module` automatically tracks all fields defined inside\nyour model object, and makes all parameters accessible using your\nmodel\\'s `parameters()` or `named_parameters()` methods.\n\nIn this example, we iterate over each parameter, and print its size and\na preview of its values.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(f\"Model structure: {model}\\n\\n\")\n\nfor name, param in model.named_parameters():\n    print(f\"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "------------------------------------------------------------------------\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Further Reading\n===============\n\n-   [torch.nn API](https://pytorch.org/docs/stable/nn.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
}