mnist_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ea9e1590-2071-44dc-b49f-1346c789a145",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c50c8ca0-bf6b-4a67-9a76-7ba6a05ac5a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_classes):\n",
    "        super(SimpleNN, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(hidden_size, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d34281ed-a6d1-4911-900e-bd46fbcf57ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\ntrain_dataset = datasets.MNIST(root=\"dataset_mnist/\", train=True, transform=transforms.ToTensor(), download = True)\\ntest_dataset = datasets.MNIST(root=\"dataset_mnist/\", train=False, transform=transforms.ToTensor(), download=True)\\n\\ntrain_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)\\ntest_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)\\n'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load MNIST dataset\n",
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) # Normalize the data\n",
    "train_dataset = datasets.MNIST(root='./dataset_mnist', train=True, transform=transform, download=True)\n",
    "test_dataset = datasets.MNIST(root='./dataset_mnist', train=False, transform=transform)\n",
    "\n",
    "# Data loaders\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=100, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=100, shuffle=False)\n",
    "\n",
    "\"\"\"\n",
    "train_dataset = datasets.MNIST(root=\"dataset_mnist/\", train=True, transform=transforms.ToTensor(), download = True)\n",
    "test_dataset = datasets.MNIST(root=\"dataset_mnist/\", train=False, transform=transforms.ToTensor(), download=True)\n",
    "\n",
    "train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a211c35-02de-4d60-991c-5c6b95753e60",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SimpleNN(\n",
      "  (fc1): Linear(in_features=784, out_features=500, bias=True)\n",
      "  (relu): ReLU()\n",
      "  (fc2): Linear(in_features=500, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Initialize the model, loss function and optimizer\n",
    "model = SimpleNN(input_size=784, hidden_size=500, num_classes=10)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ccfa100d-c711-4110-8a54-b37d2666ddae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "Epoch [1/5], Step [100/600], Loss: 0.8729\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "Epoch [1/5], Step [200/600], Loss: 0.5622\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "Epoch [1/5], Step [300/600], Loss: 0.5184\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n",
      "torch.Size([100, 10])\n",
      "torch.Size([100])\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 15\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[38;5;66;03m# Backward and optimize\u001b[39;00m\n\u001b[1;32m     14\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m---> 15\u001b[0m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     16\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (i\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m) \u001b[38;5;241m%\u001b[39m \u001b[38;5;241m100\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/_tensor.py:525\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    515\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    516\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    517\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    518\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    523\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    524\u001b[0m     )\n\u001b[0;32m--> 525\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    526\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    527\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/autograd/__init__.py:267\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    262\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    264\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    266\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 267\u001b[0m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    268\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    269\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    270\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    271\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    275\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.10/site-packages/torch/autograd/graph.py:744\u001b[0m, in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    742\u001b[0m     unregister_hooks \u001b[38;5;241m=\u001b[39m _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[1;32m    743\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 744\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    745\u001b[0m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    746\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[1;32m    747\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m    748\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "total_steps = len(train_loader)\n",
    "num_epochs = 5\n",
    "for epoch in range(num_epochs):\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        images = images.reshape(-1, 28*28)\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        print(outputs.shape)\n",
    "        print(labels.shape)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        if (i+1) % 100 == 0:\n",
    "            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' \n",
    "                   .format(epoch+1, num_epochs, i+1, total_steps, loss.item()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6601033-c501-43d0-808c-84a2a396fc61",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing the model\n",
    "with torch.no_grad():\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for images, labels in test_loader:\n",
    "        images = images.reshape(-1, 28*28)\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total += labels.size(0)\n",
    "        correct += (predicted == labels).sum().item()\n",
    "\n",
    "    print('Accuracy of the network on the 10000 test images: {} %'.format(100 * correct / total))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697f8a58-d7a6-4327-b4aa-872cb8c98f2f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "508948d4-1ef8-4a7b-9ce9-fbb41650910d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fde2c5ac-4d28-4d5c-9b8f-de711e9279f0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9bb94d8-36d1-4184-9849-603cab12ab4d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8b872f-27bc-4769-bcd0-ac83d0440a41",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffefb07d-5d83-4d56-b932-4ebf4bd4eea2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5beeb5ac-1f6a-43b4-9a20-988fe8d7d2e8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}