Malmantile.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/Users/cho/Downloads/malmantile/Malmantile.ipynb Cell 1\u001b[0m line \u001b[0;36m2\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=290'>291</a>\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(ALL_HANDS_ON_DECK):\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=291'>292</a>\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mIteration \u001b[39m\u001b[39m{\u001b[39;00mi\u001b[39m+\u001b[39m\u001b[39m1\u001b[39m\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n\u001b[0;32m--> <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=292'>293</a>\u001b[0m     model \u001b[39m=\u001b[39m federated_learning(model)\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=294'>295</a>\u001b[0m \u001b[39m# Plot Results\u001b[39;00m\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=295'>296</a>\u001b[0m \u001b[39mfor\u001b[39;00m history, label \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=296'>297</a>\u001b[0m     [tl_history, tr_history, bl_history, br_history, avg_history],\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=297'>298</a>\u001b[0m     [\u001b[39m\"\u001b[39m\u001b[39mTL\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mTR\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mBL\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mBR\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mavg\u001b[39m\u001b[39m\"\u001b[39m],\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=298'>299</a>\u001b[0m ):\n",
      "\u001b[1;32m/Users/cho/Downloads/malmantile/Malmantile.ipynb Cell 1\u001b[0m line \u001b[0;36m2\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=249'>250</a>\u001b[0m \u001b[39mfor\u001b[39;00m model, loader, history, label \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(models, loaders, histories, labels):\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=250'>251</a>\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mTraining \u001b[39m\u001b[39m{\u001b[39;00mlabel\u001b[39m}\u001b[39;00m\u001b[39m...\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m--> <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=251'>252</a>\u001b[0m     train_model(model, loader, optimizers[models\u001b[39m.\u001b[39;49mindex(model)])\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=252'>253</a>\u001b[0m     history\u001b[39m.\u001b[39mappend(test_model(label, model, test_loader))\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=254'>255</a>\u001b[0m avg_state_dict \u001b[39m=\u001b[39m average_weights(\u001b[39m*\u001b[39mmodels)\n",
      "\u001b[1;32m/Users/cho/Downloads/malmantile/Malmantile.ipynb Cell 1\u001b[0m line \u001b[0;36m2\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=226'>227</a>\u001b[0m loss \u001b[39m=\u001b[39m criterion(outputs, labels)\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=228'>229</a>\u001b[0m optimizer\u001b[39m.\u001b[39mzero_grad()\n\u001b[0;32m--> <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=229'>230</a>\u001b[0m loss\u001b[39m.\u001b[39;49mbackward()\n\u001b[1;32m    <a href='vscode-notebook-cell:/Users/cho/Downloads/malmantile/Malmantile.ipynb#W1sZmlsZQ%3D%3D?line=230'>231</a>\u001b[0m optimizer\u001b[39m.\u001b[39mstep()\n",
      "File \u001b[0;32m~/Downloads/malmantile/.conda/lib/python3.10/site-packages/torch/_tensor.py:492\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    482\u001b[0m \u001b[39mif\u001b[39;00m has_torch_function_unary(\u001b[39mself\u001b[39m):\n\u001b[1;32m    483\u001b[0m     \u001b[39mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    484\u001b[0m         Tensor\u001b[39m.\u001b[39mbackward,\n\u001b[1;32m    485\u001b[0m         (\u001b[39mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    490\u001b[0m         inputs\u001b[39m=\u001b[39minputs,\n\u001b[1;32m    491\u001b[0m     )\n\u001b[0;32m--> 492\u001b[0m torch\u001b[39m.\u001b[39;49mautograd\u001b[39m.\u001b[39;49mbackward(\n\u001b[1;32m    493\u001b[0m     \u001b[39mself\u001b[39;49m, gradient, retain_graph, create_graph, inputs\u001b[39m=\u001b[39;49minputs\n\u001b[1;32m    494\u001b[0m )\n",
      "File \u001b[0;32m~/Downloads/malmantile/.conda/lib/python3.10/site-packages/torch/autograd/__init__.py:251\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    246\u001b[0m     retain_graph \u001b[39m=\u001b[39m create_graph\n\u001b[1;32m    248\u001b[0m \u001b[39m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    249\u001b[0m \u001b[39m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    250\u001b[0m \u001b[39m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 251\u001b[0m Variable\u001b[39m.\u001b[39;49m_execution_engine\u001b[39m.\u001b[39;49mrun_backward(  \u001b[39m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    252\u001b[0m     tensors,\n\u001b[1;32m    253\u001b[0m     grad_tensors_,\n\u001b[1;32m    254\u001b[0m     retain_graph,\n\u001b[1;32m    255\u001b[0m     create_graph,\n\u001b[1;32m    256\u001b[0m     inputs,\n\u001b[1;32m    257\u001b[0m     allow_unreachable\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    258\u001b[0m     accumulate_grad\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m,\n\u001b[1;32m    259\u001b[0m )\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import datetime\n",
    "import copy\n",
    "\n",
    "# Constants and global settings\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "LEARNING_RATE = 0.0001\n",
    "BATCH_SIZE = 10\n",
    "ALL_HANDS_ON_DECK = 50\n",
    "\n",
    "\n",
    "# Transformations\n",
    "class QuadrantTransform:\n",
    "    def __init__(self, quadrant):\n",
    "        self.quadrant = quadrant\n",
    "\n",
    "    def __call__(self, img):\n",
    "        tensor = transforms.ToTensor()(img)\n",
    "        blank = torch.zeros_like(tensor)\n",
    "        slices = {\n",
    "            \"tl\": (slice(None), slice(0, 16), slice(0, 16)),\n",
    "            \"tr\": (slice(None), slice(0, 16), slice(16, 32)),\n",
    "            \"bl\": (slice(None), slice(16, 32), slice(0, 16)),\n",
    "            \"br\": (slice(None), slice(16, 32), slice(16, 32)),\n",
    "        }\n",
    "        blank[slices[self.quadrant]] = tensor[slices[self.quadrant]]\n",
    "        return (blank - 0.5) / 0.5\n",
    "\n",
    "\n",
    "def create_transform(quadrant):\n",
    "    return transforms.Compose([QuadrantTransform(quadrant)])\n",
    "\n",
    "\n",
    "def load_dataset(transform, train=True):\n",
    "    return torchvision.datasets.CIFAR10(\n",
    "        root=\"./data\", train=train, download=True, transform=transform\n",
    "    )\n",
    "\n",
    "\n",
    "def create_loader(dataset, batch_size=BATCH_SIZE):\n",
    "    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "\n",
    "def overlay_images(img_tl, img_tr, img_bl, img_br):\n",
    "    # Assuming each image has the same dimensions\n",
    "    height, width, channels = img_tl.shape\n",
    "\n",
    "    # Create an empty canvas of the same dimensions\n",
    "    combined_img = np.zeros((height, width, channels))\n",
    "\n",
    "    # Place the non-zero pixels from each image onto the canvas\n",
    "    # For this to work, the regions where the images overlap should have zeroed out pixels in all but one of the images\n",
    "    combined_img += img_tl\n",
    "    combined_img += img_tr\n",
    "    combined_img += img_bl\n",
    "    combined_img += img_br\n",
    "\n",
    "    return combined_img\n",
    "\n",
    "\n",
    "def show_example_image(tl_loader, tr_loader, bl_loader, br_loader):\n",
    "    classes = (\n",
    "        \"plane\",\n",
    "        \"car\",\n",
    "        \"bird\",\n",
    "        \"cat\",\n",
    "        \"deer\",\n",
    "        \"dog\",\n",
    "        \"frog\",\n",
    "        \"horse\",\n",
    "        \"ship\",\n",
    "        \"truck\",\n",
    "    )\n",
    "\n",
    "    # Get a batch from each loader\n",
    "    tl_images, tl_labels = next(iter(tl_loader))\n",
    "    tr_images, tr_labels = next(iter(tr_loader))\n",
    "    bl_images, bl_labels = next(iter(bl_loader))\n",
    "    br_images, br_labels = next(iter(br_loader))\n",
    "\n",
    "    # Select a random index from the batch\n",
    "    idx = np.random.randint(0, len(tl_images))\n",
    "\n",
    "    # Denormalize the images for display\n",
    "    def denormalize(img):\n",
    "        return img * 0.5 + 0.5\n",
    "\n",
    "    # Select the image from the batch using the random index and denormalize\n",
    "    img_tl = denormalize(tl_images[idx]).numpy().transpose(1, 2, 0)\n",
    "    img_tr = denormalize(tr_images[idx]).numpy().transpose(1, 2, 0)\n",
    "    img_bl = denormalize(bl_images[idx]).numpy().transpose(1, 2, 0)\n",
    "    img_br = denormalize(br_images[idx]).numpy().transpose(1, 2, 0)\n",
    "\n",
    "    combined_img = overlay_images(img_tl, img_tr, img_bl, img_br)\n",
    "\n",
    "    # Extract labels for the selected images (just as a demonstration, we can show multiple labels if needed)\n",
    "    label_combined = f\"{classes[tl_labels[idx].item()]}-{classes[tr_labels[idx].item()]}-{classes[bl_labels[idx].item()]}-{classes[br_labels[idx].item()]}\"\n",
    "\n",
    "    # Plotting\n",
    "    fig, axs = plt.subplots(1, 5, figsize=(20, 4))\n",
    "\n",
    "    axs[0].imshow(img_tl)\n",
    "    axs[0].set_title(f\"Top Left Quadrant\\n{classes[tl_labels[idx].item()]}\")\n",
    "    axs[0].axis(\"off\")\n",
    "\n",
    "    axs[1].imshow(img_tr)\n",
    "    axs[1].set_title(f\"Top Right Quadrant\\n{classes[tr_labels[idx].item()]}\")\n",
    "    axs[1].axis(\"off\")\n",
    "\n",
    "    axs[2].imshow(img_bl)\n",
    "    axs[2].set_title(f\"Bottom Left Quadrant\\n{classes[bl_labels[idx].item()]}\")\n",
    "    axs[2].axis(\"off\")\n",
    "\n",
    "    axs[3].imshow(img_br)\n",
    "    axs[3].set_title(f\"Bottom Right Quadrant\\n{classes[br_labels[idx].item()]}\")\n",
    "    axs[3].axis(\"off\")\n",
    "\n",
    "    axs[4].imshow(combined_img)\n",
    "    axs[4].set_title(f\"Combined\\n{label_combined}\")\n",
    "    axs[4].axis(\"off\")\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def conv3x3(in_channels, out_channels, stride=1):\n",
    "    return nn.Conv2d(\n",
    "        in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False\n",
    "    )\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, stride=1, downsample=None):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = conv3x3(in_channels, out_channels, stride)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv3x3(out_channels, out_channels)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        self.downsample = downsample\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        if self.downsample:\n",
    "            residual = self.downsample(x)\n",
    "        out += residual\n",
    "        out = self.relu(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, layers, num_classes=10):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.in_channels = 64  # Increase from 16 to 64 for CIFAR-10\n",
    "        self.conv = conv3x3(3, 64)\n",
    "        self.bn = nn.BatchNorm2d(64)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.layer1 = self.make_layer(block, 64, layers[0])\n",
    "        self.layer2 = self.make_layer(block, 128, layers[1], 2)\n",
    "        self.layer3 = self.make_layer(block, 256, layers[2], 2)\n",
    "        self.layer4 = self.make_layer(block, 512, layers[3], 2)  # Added another layer\n",
    "        self.avg_pool = nn.AvgPool2d(4)  # Adjust from 8 to 4 for CIFAR-10's 32x32\n",
    "        self.fc = nn.Linear(512, num_classes)\n",
    "\n",
    "    def make_layer(self, block, out_channels, blocks, stride=1):\n",
    "        downsample = None\n",
    "        if (stride != 1) or (self.in_channels != out_channels):\n",
    "            downsample = nn.Sequential(\n",
    "                conv3x3(self.in_channels, out_channels, stride=stride),\n",
    "                nn.BatchNorm2d(out_channels),\n",
    "            )\n",
    "        layers = []\n",
    "        layers.append(block(self.in_channels, out_channels, stride, downsample))\n",
    "        self.in_channels = out_channels\n",
    "        for i in range(1, blocks):\n",
    "            layers.append(block(out_channels, out_channels))\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv(x)\n",
    "        out = self.bn(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.layer4(out)  # Pass through the additional layer\n",
    "        out = self.avg_pool(out)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.fc(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "# Training and testing functions\n",
    "def test_model(debug_string, model, loader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for images, labels in loader:\n",
    "            images, labels = images.to(DEVICE), 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",
    "    accuracy = correct / total\n",
    "    print(debug_string, accuracy)\n",
    "    return accuracy\n",
    "\n",
    "\n",
    "def train_model(model, trainloader, optimizer):\n",
    "    model.train().to(DEVICE)\n",
    "    for i, (images, labels) in enumerate(trainloader):\n",
    "        images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "def average_weights(*models):\n",
    "    avg_dict = {\n",
    "        key: sum([model.state_dict()[key] for model in models]) / len(models)\n",
    "        for key in models[0].state_dict().keys()\n",
    "    }\n",
    "    return avg_dict\n",
    "\n",
    "\n",
    "def federated_learning(original_model):\n",
    "    loaders = [trainloader_tl, trainloader_tr, trainloader_bl, trainloader_br]\n",
    "    histories = [tl_history, tr_history, bl_history, br_history]\n",
    "    labels = [\"TL\", \"TR\", \"BL\", \"BR\"]\n",
    "\n",
    "    models = [copy.deepcopy(original_model).to(DEVICE) for _ in loaders]\n",
    "    optimizers = [optim.Adam(model.parameters(), lr=LEARNING_RATE) for model in models]\n",
    "\n",
    "    for model, loader, history, label in zip(models, loaders, histories, labels):\n",
    "        print(f\"Training {label}...\")\n",
    "        train_model(model, loader, optimizers[models.index(model)])\n",
    "        history.append(test_model(label, model, test_loader))\n",
    "\n",
    "    avg_state_dict = average_weights(*models)\n",
    "    new_model = ResNet(ResidualBlock, [2, 2, 2, 2]).to(DEVICE)\n",
    "    new_model.load_state_dict(avg_state_dict)\n",
    "    avg_history.append(test_model(\"avg\", new_model, test_loader))\n",
    "    return new_model\n",
    "\n",
    "\n",
    "# Main Execution\n",
    "if __name__ == \"__main__\":\n",
    "    # Data Preparation\n",
    "    transform = transforms.Compose(\n",
    "        [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]\n",
    "    )\n",
    "\n",
    "    trainsets = [\n",
    "        load_dataset(create_transform(quadrant))\n",
    "        for quadrant in [\"tl\", \"tr\", \"bl\", \"br\"]\n",
    "    ]\n",
    "    trainloader_tl, trainloader_tr, trainloader_bl, trainloader_br = [\n",
    "        create_loader(trainset) for trainset in trainsets\n",
    "    ]\n",
    "\n",
    "    show_example_image(trainloader_tl, trainloader_tr, trainloader_bl, trainloader_br)\n",
    "\n",
    "    testset = torchvision.datasets.CIFAR10(\n",
    "        root=\"./data\", train=False, download=True, transform=transform\n",
    "    )\n",
    "    test_loader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)\n",
    "\n",
    "    # Initialize Model and Criterion\n",
    "    model = ResNet(ResidualBlock, [2, 2, 2, 2]).to(DEVICE)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Training and Testing\n",
    "    tl_history, tr_history, bl_history, br_history, avg_history = [], [], [], [], []\n",
    "\n",
    "    for i in range(ALL_HANDS_ON_DECK):\n",
    "        print(f\"Iteration {i+1}\")\n",
    "        model = federated_learning(model)\n",
    "\n",
    "    # Plot Results\n",
    "    for history, label in zip(\n",
    "        [tl_history, tr_history, bl_history, br_history, avg_history],\n",
    "        [\"TL\", \"TR\", \"BL\", \"BR\", \"avg\"],\n",
    "    ):\n",
    "        plt.plot(history, label=label)\n",
    "\n",
    "    plt.legend()\n",
    "    plt.xlabel(\"Federations\")\n",
    "    plt.ylabel(\"Accuracy\")\n",
    "    plt.title(\"Project Malmantile\")\n",
    "\n",
    "    current_time = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "    plt.savefig(\"malmantile-\" + current_time + \".png\")\n",
    "    plt.show()\n",
    "\n",
    "    # Save Model\n",
    "    torch.save(model.state_dict(), \"malmantile.ckpt\")\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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}