feedforward_neural_network.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "feedforward-neural-network.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-vtLxPA05dNh"
      },
      "source": [
        "# Feedforward neural network\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rvhod0sGx5KW",
        "outputId": "c823d2c8-52f3-45bc-bf2d-57ca5d5c5f41"
      },
      "source": [
        "# code by Jinshun Lin\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "# Device configuration\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "# Heper-parameter\n",
        "input_size = 784\n",
        "hidden_size = 500\n",
        "num_classes = 10\n",
        "num_epochs = 5\n",
        "batch_size = 100\n",
        "learning_rate = 0.001\n",
        "\n",
        "# MNIST dataset\n",
        "train_dataset = torchvision.datasets.MNIST(root='../../data',\n",
        "                                           train=True,\n",
        "                                           transform=transforms.ToTensor(),\n",
        "                                           download=True)\n",
        "\n",
        "test_dataset = torchvision.datasets.MNIST(root='../../data',\n",
        "                                          train=False,\n",
        "                                          transform=transforms.ToTensor(),\n",
        "                                          )\n",
        "\n",
        "# Data loader\n",
        "train_loader = torch.utils.data.DataLoader(dataset=train_dataset,\n",
        "                                           batch_size=batch_size,\n",
        "                                           shuffle=True)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(dataset=test_dataset,\n",
        "                                          batch_size=batch_size,\n",
        "                                          shuffle=False)\n",
        "\n",
        "# Fully connected neural network with one hidden layer\n",
        "class NeuralNet(nn.Module):\n",
        "    def __init__(self, input_size, hidden_size, num_classes):\n",
        "        super(NeuralNet, 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",
        "        out = self.fc1(x)\n",
        "        out = self.relu(out)\n",
        "        out = self.fc2(out)\n",
        "        return out\n",
        "\n",
        "model = NeuralNet(input_size, hidden_size, num_classes).to(device)\n",
        "\n",
        "# Loss and optimizer\n",
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate)\n",
        "\n",
        "# Train the model\n",
        "total_step = len(train_loader)\n",
        "for epoch in range(num_epochs):\n",
        "    for i, (images, labels) in enumerate(train_loader):\n",
        "        # Move tensors to the configured device\n",
        "        images = images.reshape(-1, 28*28).to(device)\n",
        "        labels = labels.to(device)\n",
        "\n",
        "        # Forward pass\n",
        "        outputs = model(images)\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}'.format(epoch+1, num_epochs, epoch+1, total_step,loss.item()))\n",
        "\n",
        "# Test the model\n",
        "# In test phase, we don't need to compute gradients (for memory efficiency)\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).to(device)\n",
        "        labels = labels.to(device)\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))\n",
        "\n",
        "# Save the model checkpoint\n",
        "torch.save(model.state_dict(), 'model.ckpt')"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: [1/5], Step [1/600], Loss: 0.2596\n",
            "Epoch: [1/5], Step [1/600], Loss: 0.1978\n",
            "Epoch: [1/5], Step [1/600], Loss: 0.4138\n",
            "Epoch: [1/5], Step [1/600], Loss: 0.3031\n",
            "Epoch: [1/5], Step [1/600], Loss: 0.1812\n",
            "Epoch: [1/5], Step [1/600], Loss: 0.1183\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.1074\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.1194\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.1377\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.0855\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.1017\n",
            "Epoch: [2/5], Step [2/600], Loss: 0.1068\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0513\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0549\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0504\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0916\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0474\n",
            "Epoch: [3/5], Step [3/600], Loss: 0.0284\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0809\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0293\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0601\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0249\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0497\n",
            "Epoch: [4/5], Step [4/600], Loss: 0.0169\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0423\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0117\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0307\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0917\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0535\n",
            "Epoch: [5/5], Step [5/600], Loss: 0.0275\n",
            "Accuracy of the network on the 10000 test images: 97.71 %\n"
          ]
        }
      ]
    }
  ]
}