{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f666d06-d9f4-4b67-8508-b8f408918776",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import argparse\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.optim.lr_scheduler import StepLR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "149de614-bf51-4bdd-af30-fd3f69bb0c38",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, 1)\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, 1)\n",
    "        self.dropout1 = nn.Dropout(0.25)\n",
    "        self.dropout2 = nn.Dropout(0.5)\n",
    "        self.fc1 = nn.Linear(9216, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.relu(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "\n",
    "def train(args, model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = F.nll_loss(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        if batch_idx % args.log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(data), len(train_loader.dataset),\n",
    "                100. * batch_idx / len(train_loader), loss.item()))\n",
    "            if args.dry_run:\n",
    "                break\n",
    "\n",
    "\n",
    "def test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss\n",
    "            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
    "        test_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0a78a92b-0bbb-4a75-9aeb-adc58d577c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    # Training settings\n",
    "    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')\n",
    "    parser.add_argument('--batch-size', type=int, default=64, metavar='N',\n",
    "                        help='input batch size for training (default: 64)')\n",
    "    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n",
    "                        help='input batch size for testing (default: 1000)')\n",
    "    parser.add_argument('--epochs', type=int, default=14, metavar='N',\n",
    "                        help='number of epochs to train (default: 14)')\n",
    "    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',\n",
    "                        help='learning rate (default: 1.0)')\n",
    "    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',\n",
    "                        help='Learning rate step gamma (default: 0.7)')\n",
    "    parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                        help='disables CUDA training')\n",
    "    parser.add_argument('--no-mps', action='store_true', default=False,\n",
    "                        help='disables macOS GPU training')\n",
    "    parser.add_argument('--dry-run', action='store_true', default=False,\n",
    "                        help='quickly check a single pass')\n",
    "    parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                        help='random seed (default: 1)')\n",
    "    parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
    "                        help='how many batches to wait before logging training status')\n",
    "    parser.add_argument('--save-model', action='store_true', default=False,\n",
    "                        help='For Saving the current Model')\n",
    "    args = parser.parse_args()\n",
    "    use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "    use_mps = not args.no_mps and torch.backends.mps.is_available()\n",
    "\n",
    "    torch.manual_seed(args.seed)\n",
    "\n",
    "    if use_cuda:\n",
    "        device = torch.device(\"cuda\")\n",
    "    elif use_mps:\n",
    "        device = torch.device(\"mps\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "\n",
    "    train_kwargs = {'batch_size': args.batch_size}\n",
    "    test_kwargs = {'batch_size': args.test_batch_size}\n",
    "    if use_cuda:\n",
    "        cuda_kwargs = {'num_workers': 1,\n",
    "                       'pin_memory': True,\n",
    "                       'shuffle': True}\n",
    "        train_kwargs.update(cuda_kwargs)\n",
    "        test_kwargs.update(cuda_kwargs)\n",
    "\n",
    "    transform=transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "        ])\n",
    "    dataset1 = datasets.MNIST('data', train=True, download=True,\n",
    "                       transform=transform)\n",
    "    dataset2 = datasets.MNIST('data', train=False,\n",
    "                       transform=transform)\n",
    "    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)\n",
    "    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)\n",
    "\n",
    "    model = Net().to(device)\n",
    "    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)\n",
    "\n",
    "    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)\n",
    "    for epoch in range(1, args.epochs + 1):\n",
    "        train(args, model, device, train_loader, optimizer, epoch)\n",
    "        test(model, device, test_loader)\n",
    "        scheduler.step()\n",
    "\n",
    "    if args.save_model:\n",
    "        torch.save(model.state_dict(), \"mnist_cnn.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c4296754-a6f2-495d-872d-f52638860eac",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "usage: ipykernel_launcher.py [-h] [--batch-size N] [--test-batch-size N]\n",
      "                             [--epochs N] [--lr LR] [--gamma M] [--no-cuda]\n",
      "                             [--no-mps] [--dry-run] [--seed S]\n",
      "                             [--log-interval N] [--save-model]\n",
      "ipykernel_launcher.py: error: unrecognized arguments: -f /root/.local/share/jupyter/runtime/kernel-862b3d0b-fb43-47b8-ba05-b3d6f3b095d0.json\n"
     ]
    },
    {
     "ename": "SystemExit",
     "evalue": "2",
     "output_type": "error",
     "traceback": [
      "An exception has occurred, use %tb to see the full traceback.\n",
      "\u001b[0;31mSystemExit\u001b[0m\u001b[0;31m:\u001b[0m 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/IPython/core/interactiveshell.py:3556: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.\n",
      "  warn(\"To exit: use 'exit', 'quit', or Ctrl-D.\", stacklevel=1)\n"
     ]
    }
   ],
   "source": [
    "main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "437c4b58-9148-4716-9202-35c56528e8bf",
   "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
}