ff-cifar10-convnet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8c05389e",
   "metadata": {
    "id": "8c05389e"
   },
   "source": [
    "# Implementation of Hinton's Forward-Forward Algorithm on CIFAR-10 using CNN\n",
    "\n",
    "This notebook contains an implementation of Geoffrey Hinton's Forward-Forward Algorithm. Details about this algorithm can be found in Hinton's [paper](https://arxiv.org/abs/2212.13345#:~:text=The%20Forward%2DForward%20algorithm%20replaces,generated%20by%20the%20network%20itself).\n",
    "\n",
    "(Incomplete) Demonstrated below is 4. Experiments with CIFAR-10\n",
    "## Preparation:\n",
    "\n",
    "- **Import Libraries:** Set up the coding environment with the required libraries.\n",
    "\n",
    "## Definition of Necessary Functions and Classes:\n",
    "\n",
    "- **A. Function to Generate Negative Labels:** Generate negative samples for training.\n",
    "\n",
    "- **B. Function to Overlay Label onto the Input Data:** Overlay the correct class onto the input data for positive forward pass.\n",
    "\n",
    "- **C. Custom Net Class:** Definition of the network using the custom forward-forward approach.\n",
    "\n",
    "- **D1. Custom Fully Connected Layer Class:** Definition of a custom fully connected layer class.\n",
    "\n",
    "- **D1. Custom Fully Connected Layer Class:** Definition of a custom fully connected layer class.\n",
    "\n",
    "- **E. Network Hyperparameters and Device Setup:** Set the hyperparameters for the network and configure the computation device.\n",
    "\n",
    "## Implementation Steps:\n",
    "\n",
    "- **1. Load Data:** Load the MNIST dataset for training and testing.\n",
    "\n",
    "- **2. Create Network:** Instantiate the network architecture.\n",
    "\n",
    "- **3. Train the Network:** Train the network using the forward-forward algorithm.\n",
    "\n",
    "- **4. Test the Network:** Test the network's performance on the MNIST test set.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "095c14e9",
   "metadata": {
    "id": "095c14e9"
   },
   "source": [
    "### Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5d660a",
   "metadata": {
    "id": "4a5d660a"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import clear_output\n",
    "import torch.nn as nn\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision.transforms import Compose, ToTensor, Normalize\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.optim import Adam"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ddb1f7a",
   "metadata": {
    "id": "9ddb1f7a"
   },
   "source": [
    "### A. Function to generate negative labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cdefb0a",
   "metadata": {
    "id": "4cdefb0a"
   },
   "outputs": [],
   "source": [
    "# Generates negative labels for the training data, which are required for contrastive divergence training\n",
    "def get_y_neg(y):\n",
    "    y_neg = y.clone()\n",
    "    for idx, y_samp in enumerate(y):\n",
    "        allowed_indices = list(range(10))\n",
    "        allowed_indices.remove(y_samp.item())\n",
    "        y_neg[idx] = torch.tensor(allowed_indices)[torch.randint(len(allowed_indices), size=(1,))].item()\n",
    "    return y_neg.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eadc5f64",
   "metadata": {
    "id": "eadc5f64"
   },
   "source": [
    "### B. Function to overlay label onto the input data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3647e4",
   "metadata": {
    "id": "7e3647e4"
   },
   "outputs": [],
   "source": [
    "# Overlays label info onto the input data, making a certain position in the tensor representative of the label\n",
    "# experiment 1\n",
    "\n",
    "def overlay_y_on_x(x, y, classes=10):\n",
    "    \n",
    "    # trial 6\n",
    "    x_ = x.clone()\n",
    "    x_[range(x.shape[0]), :, 0, :classes] *= 0.0\n",
    "    x_[range(x.shape[0]), :, 0, y] = 1\n",
    "    return x_\n",
    "    \n",
    "    '''\n",
    "    # trial 5\n",
    "    \n",
    "    x_padded = nn.functional.pad(x, (0, 0, 1, 0))\n",
    "    x_ = torch.zeros_like(x_padded)\n",
    "    x_[range(x.shape[0]), :, 0, y] = 1.0  # set the pixel corresponding to the label of each individual image\n",
    "    x_ += x_padded\n",
    "    return x_\n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    # trial 4\n",
    "    \n",
    "    x_padded = nn.functional.pad(x, (0, 0, 1, 0))\n",
    "    x_ = torch.zeros_like(x_padded)\n",
    "    for i in range(x.size(0)):  # iterate over the batch dimension\n",
    "        x_[i, :, 0, y[i]] = 1.0  # set the pixel corresponding to the label of each individual image\n",
    "    x_ += x_padded\n",
    "    return x_\n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    # trial 3\n",
    "    \n",
    "    x_padded = nn.functional.pad(x, (0, 0, 1, 0))\n",
    "    x_ = torch.zeros_like(x_padded)\n",
    "    x_[:, :, 0, y] = 1.0  # all batches, all channels, use 1.0 for white pixel in normalized dataset\n",
    "    x_ += x_padded\n",
    "    return x_\n",
    "    \n",
    "    \n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    # trial 2\n",
    "    \n",
    "    x_padded = nn.functional.pad(x, (0, 0, 1, 0))\n",
    "    x_ = torch.zeros_like(x_padded)\n",
    "    x_[:, 0, 0, y] = x.view(x.shape[0], -1).max(dim=1)[0]\n",
    "    x_ += x_padded\n",
    "    return x_\n",
    "    \n",
    "    '''\n",
    "    \n",
    "    '''\n",
    "    # trial 1\n",
    "    \n",
    "    x_ = x.clone()\n",
    "    x_[:, 0, :classes, :classes] *= 0.0\n",
    "    x_[range(x.shape[0]), 0, y, y] = x.max()\n",
    "    return x_\n",
    "    '''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58738dc6",
   "metadata": {
    "id": "58738dc6"
   },
   "source": [
    "### C. Custom Net Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "879c3966",
   "metadata": {
    "id": "879c3966"
   },
   "outputs": [],
   "source": [
    "# Define a Net class that inherits from torch.nn.Module which is the base class for all neural network modules in PyTorch.\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        '''\n",
    "        \n",
    "        # trial 10\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            Layer(3, 64, kernel_size=3, padding=1),  # Layer 0 \n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(64, 64, kernel_size=3, padding=1),  # Layer 2\n",
    "            nn.BatchNorm2d(64),\n",
    "            nn.ReLU(inplace=True),\n",
    "            \n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 4\n",
    "\n",
    "            Layer(64, 128, kernel_size=3, padding=1),  # Layer 5 \n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(128, 128, kernel_size=3, padding=1),  # Layer 7 \n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            \n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 9\n",
    "\n",
    "            Layer(128, 256, kernel_size=3, padding=1),  # Layer 10\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(256, 256, kernel_size=3, padding=1),  # Layer 12\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(256, 256, kernel_size=3, padding=1),  # Layer 14 \n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 16\n",
    "\n",
    "            Layer(256, 512, kernel_size=3, padding=1),  # Layer 17\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 19\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 21\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 23\n",
    "\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 24\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 26\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 28\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "        '''\n",
    "        \n",
    "    \n",
    "        # trial 9\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            Layer(3, 64, kernel_size=5, padding=2),  # Layer 0\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(64, 128, kernel_size=5, padding=2, stride=2),  # Layer 2\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(128, 256, kernel_size=5, padding=2, stride=2),  # Layer 4\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(256, 512, kernel_size=5, padding=2, stride=2),  # Layer 6\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(512, 1024, kernel_size=5, padding=2, stride=2),  # Layer 8\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(1024, 2048, kernel_size=5, padding=2, stride=2),  # Layer 10\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            Layer(2048, 10, kernel_size=1, stride=1),  # Layer 12\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        \n",
    "        \n",
    " \n",
    "        '''\n",
    "        # trial 8\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            Layer(3, 64, kernel_size=3, padding=1),  # Layer 0 \n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(64, 64, kernel_size=3, padding=1),  # Layer 2\n",
    "            nn.ReLU(inplace=True),\n",
    "            \n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 4\n",
    "\n",
    "            Layer(64, 128, kernel_size=3, padding=1),  # Layer 5 \n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(128, 128, kernel_size=3, padding=1),  # Layer 7 \n",
    "            nn.ReLU(inplace=True),\n",
    "            \n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 9\n",
    "\n",
    "            Layer(128, 256, kernel_size=3, padding=1),  # Layer 10\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(256, 256, kernel_size=3, padding=1),  # Layer 12\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(256, 256, kernel_size=3, padding=1),  # Layer 14 \n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 16\n",
    "\n",
    "            Layer(256, 512, kernel_size=3, padding=1),  # Layer 17\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 19\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 21\n",
    "            nn.ReLU(inplace=True),\n",
    "\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),  # Layer 23\n",
    "\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 24\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 26\n",
    "            nn.ReLU(inplace=True),\n",
    "            Layer(512, 512, kernel_size=3, padding=1),  # Layer 28\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "        \n",
    "        '''      \n",
    "        \n",
    "        '''\n",
    "        # trial 7\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 64, kernel_size=5, padding=2),  # Layer 0 with 5x5 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "\n",
    "        Layer(64, 128, kernel_size=5, padding=2, stride=2),  # Layer 2 with 5x5 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "\n",
    "        Layer(128, 256, kernel_size=5, padding=2, stride=2),  # Layer 4 with 5x5 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "\n",
    "        Layer(256, 10, kernel_size=1, stride=1),  # Layer 6 with 1x1 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "        )\n",
    "        \n",
    "        '''\n",
    "        \n",
    "        '''\n",
    "        # trial 6\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 3, kernel_size=11, padding=5),  # Layer 1 with 11x11 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(3, 3, kernel_size=11, padding=5),  # Layer 2 with 11x11 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(3, 3, kernel_size=11, padding=5),  # Layer 3 with 11x11 receptive field\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.Flatten(),  # Flatten the output\n",
    "        FullyConnectedLayer(3 * 32 * 32, 10),  # Fully connected layer with 3072 inputs and 10 outputs\n",
    "        )\n",
    "        '''\n",
    "        \n",
    "              \n",
    "        '''\n",
    "        # trial 5\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 32, kernel_size=11, padding=1),  # Conv 1 with kernel size 11x11\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(32, 32, kernel_size=9, padding=1),  # Conv 2 with kernel size 11x11\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(32, 64, kernel_size=4, padding=1),  # Conv 3 with kernel size 11x11\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        nn.Flatten(),\n",
    "        FullyConnectedLayer(128 * 3 * 3, 500),  # Adjusted input dimension\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 500),\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 10)\n",
    "        )\n",
    "        '''\n",
    "\n",
    "        '''\n",
    "        # trial 4\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 32, kernel_size=11, padding=1),  # Conv 1\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(32, 32, kernel_size=11, padding=1),  # Conv 2\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(32, 64, kernel_size=5, padding=1),  # Conv 3\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(64, 64, kernel_size=5, padding=1),  # Conv 4\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(64, 128, kernel_size=3, padding=1),  # Conv 5\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(128, 128, kernel_size=3, padding=1),  # Conv 6\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        nn.Flatten(),\n",
    "        FullyConnectedLayer(128 * 4 * 4, 500),  # Note: If the dimensions don't match, adjust the first dimension here\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 500),\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 10)\n",
    "        )\n",
    "        '''\n",
    "        \n",
    "        '''\n",
    "        # trial 3  \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 32, kernel_size=3, padding=1),  # Conv 1\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(32, 32, kernel_size=3, padding=1),  # Conv 2\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(32, 64, kernel_size=3, padding=1),  # Conv 3\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(64, 64, kernel_size=3, padding=1),  # Conv 4\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(64, 128, kernel_size=3, padding=1),  # Conv 5\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(128, 128, kernel_size=3, padding=1),  # Conv 6\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        nn.Flatten(),\n",
    "        FullyConnectedLayer(128 * 4 * 4, 500),  # Note: If the dimensions don't match, adjust the first dimension here\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 500),\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 10)\n",
    "        )\n",
    "        '''\n",
    "        \n",
    "        '''\n",
    "        # trial 2\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 32, kernel_size=3, padding=1),  # Conv 1\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(32, 32, kernel_size=3, padding=1),  # Conv 2\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(32, 64, kernel_size=3, padding=1),  # Conv 3\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(64, 64, kernel_size=3, padding=1),  # Conv 4\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        Layer(64, 128, kernel_size=3, padding=1),  # Conv 5\n",
    "        nn.ReLU(inplace=True),\n",
    "        Layer(128, 128, kernel_size=3, padding=1),  # Conv 6\n",
    "        nn.ReLU(inplace=True),\n",
    "        nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "        nn.Flatten(),\n",
    "        FullyConnectedLayer(128 * 4 * 4, 500),\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 500),\n",
    "        nn.ReLU(inplace=True),\n",
    "        FullyConnectedLayer(500, 10)\n",
    "        )\n",
    "        \n",
    "        '''\n",
    "\n",
    "        '''\n",
    "        # trial 1\n",
    "\n",
    "        self.layers = nn.Sequential(\n",
    "        Layer(3, 32, kernel_size=3, padding=1),\n",
    "        nn.MaxPool2d(2, 2),\n",
    "        Layer(32, 64, kernel_size=3, padding=1),\n",
    "        nn.MaxPool2d(2, 2),\n",
    "        nn.Flatten(),\n",
    "        FullyConnectedLayer(4096, 500),  # use new FullyConnectedLayer\n",
    "        FullyConnectedLayer(500, 500),\n",
    "        FullyConnectedLayer(500, 10),\n",
    "        )\n",
    "        \n",
    "        '''\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layers(x)\n",
    "\n",
    "    def predict(self, x):\n",
    "        goodness_per_label = []\n",
    "        for label in range(10):\n",
    "            h = overlay_y_on_x(x, label)\n",
    "            goodness = []\n",
    "            for layer in self.layers:\n",
    "                h = layer(h)\n",
    "                goodness += [(h.pow(2).sum() / h.numel()).unsqueeze(0)]\n",
    "            goodness_per_label += [torch.sum(torch.stack(goodness)).unsqueeze(0)]\n",
    "        goodness_per_label = torch.cat(goodness_per_label, 0)\n",
    "        return goodness_per_label.argmax(0)\n",
    "\n",
    "    def custom_train(self, x_pos, x_neg):\n",
    "        h_pos, h_neg = x_pos, x_neg\n",
    "        for i, layer in enumerate(self.layers):\n",
    "            print(\"training layer: \", i)\n",
    "            if isinstance(layer, Layer):  # only call custom_train on instances of the Layer class\n",
    "                h_pos, h_neg = layer.custom_train(h_pos, h_neg)\n",
    "            elif isinstance(layer, FullyConnectedLayer):  # only call custom_train on instances of the FullyConnectedLayer class\n",
    "                h_pos, h_neg = layer.custom_train(h_pos, h_neg)\n",
    "            else:  # for other layers, just pass the data through\n",
    "                h_pos = layer(h_pos)\n",
    "                h_neg = layer(h_neg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88668c56",
   "metadata": {
    "id": "88668c56"
   },
   "source": [
    "### D1. Custom Fully Connected Layer Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205738e5",
   "metadata": {
    "id": "205738e5"
   },
   "outputs": [],
   "source": [
    "# Fully Connected layers of the network.\n",
    "class FullyConnectedLayer(nn.Module):\n",
    "    def __init__(self, in_features, out_features, bias=True, final_layer=False):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(in_features, out_features, bias=bias)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.opt = Adam(self.parameters(), lr=args.lr)\n",
    "        self.fcl_threshold = args.fcl_threshold\n",
    "        self.num_epochs = args.epochs\n",
    "        self.final_layer = final_layer\n",
    "\n",
    "    def forward(self, x):\n",
    "        #L2 norm\n",
    "        x_direction = x / (x.norm(2, 1, keepdim=True) + 1e-4)\n",
    "        return self.relu(self.linear(x_direction))\n",
    "\n",
    "    def custom_train(self, x_pos, x_neg):\n",
    "        for i in range(self.num_epochs):\n",
    "            g_pos = self.forward(x_pos).pow(2).mean().unsqueeze(0)  # mean over all dimensions in a sample\n",
    "            g_neg = self.forward(x_neg).pow(2).mean().unsqueeze(0)  # mean over all dimensions in a sample\n",
    "            loss = torch.log(1 + torch.exp(torch.cat([-g_pos + self.fcl_threshold, g_neg - self.fcl_threshold]))).mean()\n",
    "            self.opt.zero_grad()\n",
    "            loss.backward()\n",
    "            self.opt.step()\n",
    "            if i % args.log_interval == 0:\n",
    "                print(\"Loss: \", loss.item())\n",
    "        return self.forward(x_pos).detach(), self.forward(x_neg).detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6bcfaa3",
   "metadata": {
    "id": "c6bcfaa3"
   },
   "source": [
    "### D2. Custom Convolutional Layer Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21e3ace7",
   "metadata": {
    "id": "21e3ace7"
   },
   "outputs": [],
   "source": [
    "# Convolutional layers of the network.\n",
    "class Layer(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=0, bias=True, final_layer=False):\n",
    "        super().__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.opt = Adam(self.parameters(), lr=args.lr)\n",
    "        self.conv_threshold = args.conv_threshold\n",
    "        self.num_epochs = args.epochs\n",
    "        self.final_layer = final_layer\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        '''\n",
    "        # trial 10\n",
    "        # batch norm\n",
    "        x = self.conv(x)\n",
    "        return self.relu(x) \n",
    "        '''\n",
    "        \n",
    "        # L2 norm\n",
    "        x_direction = x / (x.norm(2, 1, keepdim=True) + 1e-4)\n",
    "        return self.relu(self.conv(x_direction))           \n",
    "\n",
    "    def custom_train(self, x_pos, x_neg):\n",
    "        \n",
    "        # initialize lists to hold values\n",
    "        loss_values = []\n",
    "        g_pos_values = []\n",
    "        g_neg_values = []\n",
    "\n",
    "        # initialize figure\n",
    "        fig = plt.figure(figsize=(12,8))\n",
    "        \n",
    "        for i in range(self.num_epochs):\n",
    "            g_pos = self.forward(x_pos).pow(2).mean(dim=(1,2,3)).unsqueeze(0)  # mean over all dimensions\n",
    "            g_neg = self.forward(x_neg).pow(2).mean(dim=(1,2,3)).unsqueeze(0)  # mean over all dimensions\n",
    "            loss = torch.log(1 + torch.exp(torch.cat([-g_pos + self.conv_threshold, g_neg - self.conv_threshold]))).mean()\n",
    "            self.opt.zero_grad()\n",
    "            \n",
    "            #loss.backward()\n",
    "            loss.backward(retain_graph=True)\n",
    "\n",
    "            self.opt.step()\n",
    "            \n",
    "            if i % args.log_interval == 0:\n",
    "                loss_values.append(loss.item())\n",
    "                g_pos_values.append(g_pos.mean().item())  # take mean of all batch values\n",
    "                g_neg_values.append(g_neg.mean().item())  # take mean of all batch values\n",
    "\n",
    "                # plotting\n",
    "                plt.subplot(3,1,1)\n",
    "                plt.plot(loss_values, color='blue')\n",
    "                plt.title(\"Loss during training\")\n",
    "\n",
    "                plt.subplot(3,1,2)\n",
    "                plt.plot(g_pos_values, color='green')\n",
    "                plt.title(\"g_pos during training\")\n",
    "\n",
    "                plt.subplot(3,1,3)\n",
    "                plt.plot(g_neg_values, color='red')\n",
    "                plt.title(\"g_neg during training\")\n",
    "\n",
    "                plt.tight_layout()\n",
    "                clear_output(wait=True)  # this clears the output of the cell, useful for updating the plots\n",
    "                plt.show()\n",
    "\n",
    "            # Print the loss at each step\n",
    "            print(f'Loss at step {i}: {loss.item()}')\n",
    "            \n",
    "        return self.forward(x_pos).detach(), self.forward(x_neg).detach()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae11a7b4",
   "metadata": {
    "id": "ae11a7b4"
   },
   "source": [
    "### E. Network Hyperparameters and device setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d91f7b",
   "metadata": {
    "id": "35d91f7b"
   },
   "outputs": [],
   "source": [
    "class Args:\n",
    "    train_size = 1000 #10000 #50000\n",
    "    test_size = 100 #2000 #10000\n",
    "    epochs = 1000\n",
    "    lr = 0.05\n",
    "    no_cuda = False\n",
    "    no_mps = False\n",
    "    save_model = False\n",
    "    fcl_threshold = 1\n",
    "    conv_threshold = 0.02\n",
    "    seed = 1234\n",
    "    log_interval = 10\n",
    "\n",
    "args = 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",
    "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",
    "train_kwargs = {\"batch_size\": args.train_size}\n",
    "test_kwargs = {\"batch_size\": args.test_size}\n",
    "if use_cuda:\n",
    "    cuda_kwargs = {\"num_workers\": 1, \"pin_memory\": True, \"shuffle\": True}\n",
    "    train_kwargs.update(cuda_kwargs)\n",
    "    test_kwargs.update(cuda_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c16d01f",
   "metadata": {
    "id": "1c16d01f"
   },
   "source": [
    "### 1. Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc415cf",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6cc415cf",
    "outputId": "1a17a060-0694-4afe-c363-16525dcf59a9"
   },
   "outputs": [],
   "source": [
    "# Load Data\n",
    "transform = Compose(\n",
    "    [\n",
    "        ToTensor(),\n",
    "        Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_loader = DataLoader(CIFAR10(\"./data/\", train=True, download=True, transform=transform), **train_kwargs)\n",
    "test_loader = DataLoader(CIFAR10(\"./data/\", train=False, download=True, transform=transform), **test_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c4ce8d3",
   "metadata": {
    "id": "7c4ce8d3"
   },
   "source": [
    "### 2. Create network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4658e5d1",
   "metadata": {
    "id": "4658e5d1"
   },
   "outputs": [],
   "source": [
    "# Create Model\n",
    "net = Net().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "521e97e5",
   "metadata": {
    "id": "521e97e5"
   },
   "source": [
    "### 3. Train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90f02705",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = next(iter(train_loader))\n",
    "x, y = x.to(device), y.to(device)\n",
    "x_pos = overlay_y_on_x(x, y)\n",
    "y_neg = get_y_neg(y)\n",
    "x_neg = overlay_y_on_x(x, y_neg)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a679bd21",
   "metadata": {},
   "source": [
    "### Inspect Tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e0e0bd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "x.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a37f7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x.size(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29170241",
   "metadata": {},
   "outputs": [],
   "source": [
    "y.size()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a7247fb",
   "metadata": {},
   "source": [
    "### Visualize Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d0642ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(5, 3, figsize=(10, 10))\n",
    "\n",
    "# Define a dictionary to map class indices to class names (replace this with your actual classes)\n",
    "class_dict = {i: 'class_' + str(i) for i in range(10)}\n",
    "\n",
    "for i in range(5):\n",
    "    img = x[i].cpu().numpy().transpose(1,2,0)\n",
    "    pos_img = x_pos[i].cpu().numpy().transpose(1,2,0)\n",
    "    neg_img = x_neg[i].cpu().numpy().transpose(1,2,0)\n",
    "\n",
    "    axs[i, 0].imshow(img)\n",
    "    axs[i, 0].set_title('Original: ' + class_dict[int(y[i])] + '\\n Shape: ' + str(img.shape))\n",
    "    \n",
    "    axs[i, 1].imshow(pos_img)\n",
    "    axs[i, 1].set_title('Positive: ' + class_dict[int(y[i])] + '\\n Shape: ' + str(pos_img.shape))\n",
    "    \n",
    "    axs[i, 2].imshow(neg_img)\n",
    "    axs[i, 2].set_title('Negative: ' + class_dict[int(y_neg[i])] + '\\n Shape: ' + str(neg_img.shape))\n",
    "\n",
    "for ax in axs.flat:\n",
    "    ax.axis('off')\n",
    "\n",
    "plt.tight_layout()  # Adjusts subplot params so that subplots are nicely fit in the figure\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f7dd1fc",
   "metadata": {},
   "source": [
    "### Train the Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "991a1862",
   "metadata": {},
   "outputs": [],
   "source": [
    "net.custom_train(x_pos, x_neg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "z_F1ZtN2IIzS",
   "metadata": {
    "id": "z_F1ZtN2IIzS"
   },
   "outputs": [],
   "source": [
    "print(\"Train Accuracy: {:.2f}%\".format(100 * net.predict(x).eq(y).float().mean().item()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d255817d",
   "metadata": {
    "id": "d255817d"
   },
   "source": [
    "### 4. Test the Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389d35f1",
   "metadata": {
    "id": "389d35f1"
   },
   "outputs": [],
   "source": [
    "# Test Model\n",
    "x_te, y_te = next(iter(test_loader))\n",
    "x_te, y_te = x_te.to(device), y_te.to(device)\n",
    "if args.save_model:\n",
    "    torch.save(net.state_dict(), \"cifar10_ff.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68c820e7",
   "metadata": {
    "id": "68c820e7"
   },
   "outputs": [],
   "source": [
    "print(\"Test Accuracy: {:.2f}%\".format(100 * net.predict(x_te).eq(y_te).float().mean().item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b2d84d3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10 (tensorflow)",
   "language": "python",
   "name": "tensorflow"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}