train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleNN, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, 28 * 28)\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('mnist_samples/mnist_label_0.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_1.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_2.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_3.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_4.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_5.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_6.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_7.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_8.pt'),\n",
       " PosixPath('mnist_samples/mnist_label_9.pt')]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset_path = Path(\"mnist_samples\")\n",
    "dataset_path_files = sorted(list(dataset_path.glob(\"*.pt\")))\n",
    "dataset_path_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('pretrained_models/pretrained_mnist_label_0.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_1.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_2.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_3.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_4.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_5.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_6.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_7.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_8.pt'),\n",
       " PosixPath('pretrained_models/pretrained_mnist_label_9.pt')]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saved_models = [\n",
    "    Path(\"pretrained_models\") / (\"pretrained_\" + f.name) for f in dataset_path_files\n",
    "]\n",
    "saved_models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(dataset_file_path: Path, saved_model_path: Path):\n",
    "    model = SimpleNN()  # Initialize train_model\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=0.1)\n",
    "\n",
    "    images, labels = torch.load(dataset_file_path)\n",
    "\n",
    "    # create a tensordataset\n",
    "    dataset = TensorDataset(images, labels)\n",
    "\n",
    "    # create a dataloader for the dataset\n",
    "    train_loader = DataLoader(dataset, batch_size=64, shuffle=True)\n",
    "\n",
    "    # training loop\n",
    "    for epoch in range(1000):\n",
    "        running_loss = 0\n",
    "        for images, labels in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # accumulate loss\n",
    "            running_loss += loss.item()\n",
    "\n",
    "        # Calculate average loss for the epoch\n",
    "        avg_loss = running_loss / len(train_loader)\n",
    "\n",
    "        if (epoch + 1) % 100 == 0:\n",
    "            avg_loss = running_loss / len(train_loader)\n",
    "            print(f\"Epoch {epoch + 1:04d}: Loss = {avg_loss:.6f}\")\n",
    "\n",
    "    # Save the model\n",
    "    torch.save(model.state_dict(), str(saved_model_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2148638/1992216445.py:7: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  images, labels = torch.load(dataset_file_path)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0100: Loss = 0.000087\n",
      "Epoch 0200: Loss = 0.000047\n",
      "Epoch 0300: Loss = 0.000026\n",
      "Epoch 0400: Loss = 0.000020\n",
      "Epoch 0500: Loss = 0.000015\n",
      "Epoch 0600: Loss = 0.000012\n",
      "Epoch 0700: Loss = 0.000011\n",
      "Epoch 0800: Loss = 0.000009\n",
      "Epoch 0900: Loss = 0.000008\n",
      "Epoch 1000: Loss = 0.000007\n",
      "Epoch 0100: Loss = 0.000117\n",
      "Epoch 0200: Loss = 0.000058\n",
      "Epoch 0300: Loss = 0.000035\n",
      "Epoch 0400: Loss = 0.000026\n",
      "Epoch 0500: Loss = 0.000020\n",
      "Epoch 0600: Loss = 0.000017\n",
      "Epoch 0700: Loss = 0.000015\n",
      "Epoch 0800: Loss = 0.000013\n",
      "Epoch 0900: Loss = 0.000011\n",
      "Epoch 1000: Loss = 0.000010\n",
      "Epoch 0100: Loss = 0.000098\n",
      "Epoch 0200: Loss = 0.000046\n",
      "Epoch 0300: Loss = 0.000030\n",
      "Epoch 0400: Loss = 0.000021\n",
      "Epoch 0500: Loss = 0.000017\n",
      "Epoch 0600: Loss = 0.000014\n",
      "Epoch 0700: Loss = 0.000012\n",
      "Epoch 0800: Loss = 0.000010\n",
      "Epoch 0900: Loss = 0.000009\n",
      "Epoch 1000: Loss = 0.000008\n",
      "Epoch 0100: Loss = 0.000090\n",
      "Epoch 0200: Loss = 0.000042\n",
      "Epoch 0300: Loss = 0.000027\n",
      "Epoch 0400: Loss = 0.000020\n",
      "Epoch 0500: Loss = 0.000016\n",
      "Epoch 0600: Loss = 0.000013\n",
      "Epoch 0700: Loss = 0.000011\n",
      "Epoch 0800: Loss = 0.000010\n",
      "Epoch 0900: Loss = 0.000008\n",
      "Epoch 1000: Loss = 0.000008\n",
      "Epoch 0100: Loss = 0.000099\n",
      "Epoch 0200: Loss = 0.000047\n",
      "Epoch 0300: Loss = 0.000030\n",
      "Epoch 0400: Loss = 0.000022\n",
      "Epoch 0500: Loss = 0.000018\n",
      "Epoch 0600: Loss = 0.000014\n",
      "Epoch 0700: Loss = 0.000012\n",
      "Epoch 0800: Loss = 0.000010\n",
      "Epoch 0900: Loss = 0.000009\n",
      "Epoch 1000: Loss = 0.000008\n",
      "Epoch 0100: Loss = 0.000107\n",
      "Epoch 0200: Loss = 0.000052\n",
      "Epoch 0300: Loss = 0.000033\n",
      "Epoch 0400: Loss = 0.000024\n",
      "Epoch 0500: Loss = 0.000019\n",
      "Epoch 0600: Loss = 0.000016\n",
      "Epoch 0700: Loss = 0.000013\n",
      "Epoch 0800: Loss = 0.000011\n",
      "Epoch 0900: Loss = 0.000010\n",
      "Epoch 1000: Loss = 0.000009\n",
      "Epoch 0100: Loss = 0.000098\n",
      "Epoch 0200: Loss = 0.000046\n",
      "Epoch 0300: Loss = 0.000031\n",
      "Epoch 0400: Loss = 0.000023\n",
      "Epoch 0500: Loss = 0.000017\n",
      "Epoch 0600: Loss = 0.000014\n",
      "Epoch 0700: Loss = 0.000012\n",
      "Epoch 0800: Loss = 0.000010\n",
      "Epoch 0900: Loss = 0.000010\n",
      "Epoch 1000: Loss = 0.000008\n",
      "Epoch 0100: Loss = 0.000125\n",
      "Epoch 0200: Loss = 0.000057\n",
      "Epoch 0300: Loss = 0.000036\n",
      "Epoch 0400: Loss = 0.000027\n",
      "Epoch 0500: Loss = 0.000021\n",
      "Epoch 0600: Loss = 0.000017\n",
      "Epoch 0700: Loss = 0.000014\n",
      "Epoch 0800: Loss = 0.000013\n",
      "Epoch 0900: Loss = 0.000011\n",
      "Epoch 1000: Loss = 0.000010\n",
      "Epoch 0100: Loss = 0.000094\n",
      "Epoch 0200: Loss = 0.000046\n",
      "Epoch 0300: Loss = 0.000030\n",
      "Epoch 0400: Loss = 0.000026\n",
      "Epoch 0500: Loss = 0.000017\n",
      "Epoch 0600: Loss = 0.000014\n",
      "Epoch 0700: Loss = 0.000012\n",
      "Epoch 0800: Loss = 0.000011\n",
      "Epoch 0900: Loss = 0.000010\n",
      "Epoch 1000: Loss = 0.000008\n",
      "Epoch 0100: Loss = 0.000096\n",
      "Epoch 0200: Loss = 0.000045\n",
      "Epoch 0300: Loss = 0.000029\n",
      "Epoch 0400: Loss = 0.000021\n",
      "Epoch 0500: Loss = 0.000018\n",
      "Epoch 0600: Loss = 0.000014\n",
      "Epoch 0700: Loss = 0.000012\n",
      "Epoch 0800: Loss = 0.000010\n",
      "Epoch 0900: Loss = 0.000009\n",
      "Epoch 1000: Loss = 0.000008\n"
     ]
    }
   ],
   "source": [
    "for dataset_file_path, saved_model_path in zip(dataset_path_files, saved_models):\n",
    "    train_model(dataset_file_path, saved_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model_path: Path, test_data_path: Path):\n",
    "    # Load the model\n",
    "    model = SimpleNN()\n",
    "    model.load_state_dict(torch.load(str(model_path), weights_only=True))\n",
    "    model.eval()\n",
    "\n",
    "    # Load test data\n",
    "    test_images, test_labels = torch.load(test_data_path, weights_only=True)\n",
    "\n",
    "    # Make predictions\n",
    "    with torch.no_grad():\n",
    "        outputs = model(test_images)\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "\n",
    "    # Calculate accuracy\n",
    "    total = test_labels.size(0)\n",
    "    correct = (predicted == test_labels).sum().item()\n",
    "    accuracy = 100 * correct / total\n",
    "\n",
    "    print(f\"Accuracy on test data: {accuracy:.2f}%\")\n",
    "    return accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n",
      "Accuracy on test data: 100.00%\n"
     ]
    }
   ],
   "source": [
    "for model_path, test_data_path in zip(saved_models, dataset_path_files):\n",
    "    test_model(model_path, test_data_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}