compEfficiency.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/zephyr/anaconda3/envs/modmi/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device is: mps\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from matplotlib.patches import Ellipse, Circle\n",
    "import time\n",
    "from itertools import islice\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm.auto import tqdm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "\n",
    "steps = 10000\n",
    "\n",
    "def cycle(iterable):\n",
    "    while True:\n",
    "        for x in iterable:\n",
    "            yield x\n",
    "\n",
    "device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "print(\"device is:\", device)\n",
    "\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "#torch.set_default_tensor_type(torch.DoubleTensor)\n",
    "\n",
    "\n",
    "class BioLinear2D(nn.Module):\n",
    "\n",
    "    def __init__(self, in_dim, out_dim, in_fold=1, out_fold=1, out_ring=False):\n",
    "        super(BioLinear2D, self).__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.out_dim = out_dim\n",
    "        self.linear = nn.Linear(in_dim, out_dim)\n",
    "        self.in_fold = in_fold\n",
    "        self.out_fold = out_fold\n",
    "        assert in_dim % in_fold == 0\n",
    "        assert out_dim % out_fold == 0\n",
    "\n",
    "        #compute in_cor, shape: (in_dim_sqrt, in_dim_sqrt)\n",
    "        in_dim_fold = int(in_dim/in_fold)\n",
    "        out_dim_fold = int(out_dim/out_fold)\n",
    "        in_dim_sqrt = int(np.sqrt(in_dim_fold))\n",
    "        out_dim_sqrt = int(np.sqrt(out_dim_fold))\n",
    "        x = np.linspace(1/(2*in_dim_sqrt), 1-1/(2*in_dim_sqrt), num=in_dim_sqrt)\n",
    "        X, Y = np.meshgrid(x, x)\n",
    "        self.in_coordinates = torch.tensor(np.transpose(np.array([X.reshape(-1,), Y.reshape(-1,)])), dtype=torch.float)\n",
    "\n",
    "        # compute out_cor, shape: (out_dim_sqrt, out_dim_sqrt)\n",
    "        if out_ring:\n",
    "            thetas = np.linspace(1/(2*out_dim_fold)*2*np.pi, (1-1/(2*out_dim_fold))*2*np.pi, num=out_dim_fold)\n",
    "            self.out_coordinates = 0.5+torch.tensor(np.transpose(np.array([np.cos(thetas), np.sin(thetas)]))/4, dtype=torch.float)\n",
    "        else:\n",
    "            x = np.linspace(1/(2*out_dim_sqrt), 1-1/(2*out_dim_sqrt), num=out_dim_sqrt)\n",
    "            X, Y = np.meshgrid(x, x)\n",
    "            self.out_coordinates = torch.tensor(np.transpose(np.array([X.reshape(-1,), Y.reshape(-1,)])), dtype=torch.float)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "\n",
    "class BioMLP2D(nn.Module):\n",
    "    def __init__(self, in_dim=2, out_dim=2, w=2, depth=2, shp=None, token_embedding=False, embedding_size=None):\n",
    "        super(BioMLP2D, self).__init__()\n",
    "        if shp == None:\n",
    "            shp = [in_dim] + [w]*(depth-1) + [out_dim]\n",
    "            self.in_dim = in_dim\n",
    "            self.out_dim = out_dim\n",
    "            self.depth = depth\n",
    "\n",
    "        else:\n",
    "            self.in_dim = shp[0]\n",
    "            self.out_dim = shp[-1]\n",
    "            self.depth = len(shp) - 1\n",
    "        linear_list = []\n",
    "        for i in range(self.depth):\n",
    "            if i == 0:\n",
    "                # for modular addition\n",
    "                #linear_list.append(BioLinear(shp[i], shp[i+1], in_fold=2))\n",
    "                # for regression\n",
    "                linear_list.append(BioLinear2D(shp[i], shp[i+1], in_fold=1).to(device))\n",
    "            elif i == self.depth - 1:\n",
    "                linear_list.append(BioLinear2D(shp[i], shp[i+1], in_fold=1, out_ring=True).to(device))\n",
    "            else:\n",
    "                linear_list.append(BioLinear2D(shp[i], shp[i+1]).to(device))\n",
    "        self.linears = nn.ModuleList(linear_list)\n",
    "\n",
    "\n",
    "        if token_embedding == True:\n",
    "            # embedding size: number of tokens * embedding dimension\n",
    "            self.embedding = torch.nn.Parameter(torch.normal(0,1,size=embedding_size)).to(device)\n",
    "\n",
    "        self.shp = shp\n",
    "        # parameters for the bio-inspired trick\n",
    "        self.l0 = 0.5 # distance between two nearby layers\n",
    "        self.in_perm = nn.Parameter(torch.tensor(np.arange(int(self.in_dim/self.linears[0].in_fold)), dtype=torch.float))\n",
    "        self.out_perm = nn.Parameter(torch.tensor(np.arange(int(self.out_dim/self.linears[-1].out_fold)), dtype=torch.float))\n",
    "        self.top_k = 30\n",
    "        self.token_embedding = token_embedding\n",
    "\n",
    "    def forward(self, x):\n",
    "        shp = x.shape\n",
    "        x = x.reshape(shp[0],-1)\n",
    "        shp = x.shape\n",
    "        in_fold = self.linears[0].in_fold\n",
    "        x = x.reshape(shp[0], in_fold, int(shp[1]/in_fold))\n",
    "        x = x[:,:,self.in_perm.long()]\n",
    "        x = x.reshape(shp[0], shp[1])\n",
    "        f = torch.nn.SiLU()\n",
    "        for i in range(self.depth-1):\n",
    "            x = f(self.linears[i](x))\n",
    "        x = self.linears[-1](x)\n",
    "\n",
    "        out_perm_inv = torch.zeros(self.out_dim, dtype=torch.long)\n",
    "        out_perm_inv[self.out_perm.long()] = torch.arange(self.out_dim)\n",
    "        x = x[:,out_perm_inv]\n",
    "        #x = x[:,self.out_perm]\n",
    "\n",
    "        return x\n",
    "\n",
    "    def get_linear_layers(self):\n",
    "        return self.linears\n",
    "\n",
    "    def get_cc(self, weight_factor=2.0, bias_penalize=True, no_penalize_last=False):\n",
    "        # compute connection cost\n",
    "        cc = 0\n",
    "        num_linear = len(self.linears)\n",
    "        for i in range(num_linear):\n",
    "            if i == num_linear - 1 and no_penalize_last:\n",
    "                weight_factor = 0.\n",
    "            biolinear = self.linears[i].to(device)\n",
    "            dist = torch.sum(torch.abs(biolinear.out_coordinates.unsqueeze(dim=1) - biolinear.in_coordinates.unsqueeze(dim=0)),dim=2).to(device)\n",
    "            # print(\"yooo\")\n",
    "            # print(biolinear.linear.weight.device)\n",
    "            # print(dist.device)\n",
    "            # print(self.l0)\n",
    "            cc += torch.mean(torch.abs(biolinear.linear.weight)*(weight_factor*dist+self.l0))\n",
    "            if bias_penalize == True:\n",
    "                cc += torch.mean(torch.abs(biolinear.linear.bias)*(self.l0))\n",
    "        if self.token_embedding:\n",
    "            cc += torch.mean(torch.abs(self.embedding)*(self.l0))\n",
    "            #pass\n",
    "        return cc\n",
    "\n",
    "    def swap_weight(self, weights, j, k, swap_type=\"out\"):\n",
    "        with torch.no_grad():\n",
    "            if swap_type == \"in\":\n",
    "                temp = weights[:,j].clone()\n",
    "                weights[:,j] = weights[:,k].clone()\n",
    "                weights[:,k] = temp\n",
    "            elif swap_type == \"out\":\n",
    "                temp = weights[j].clone()\n",
    "                weights[j] = weights[k].clone()\n",
    "                weights[k] = temp\n",
    "            else:\n",
    "                raise Exception(\"Swap type {} is not recognized!\".format(swap_type))\n",
    "\n",
    "    def swap_bias(self, biases, j, k):\n",
    "        with torch.no_grad():\n",
    "            temp = biases[j].clone()\n",
    "            biases[j] = biases[k].clone()\n",
    "            biases[k] = temp\n",
    "\n",
    "    def swap(self, i, j, k):\n",
    "        # in the ith layer (of neurons), swap the jth and the kth neuron.\n",
    "        # Note: n layers of weights means n+1 layers of neurons.\n",
    "        # (incoming, outgoing) * weights + biases are swapped.\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            return\n",
    "            # for images, do not allow input_perm\n",
    "            # input layer, only has outgoing weights; update in_perm\n",
    "            weights = linears[i].linear.weight\n",
    "            infold = linears[i].in_fold\n",
    "            fold_dim = int(weights.shape[1]/infold)\n",
    "            for l in range(infold):\n",
    "                self.swap_weight(weights, j+fold_dim*l, k+fold_dim*l, swap_type=\"in\")\n",
    "            # change input_perm. do not allow input_perm for images\n",
    "            self.swap_bias(self.in_perm, j, k)\n",
    "        elif i == num_linear:\n",
    "            # output layer, only has incoming weights and biases; update out_perm\n",
    "            weights = linears[i-1].linear.weight\n",
    "            biases = linears[i-1].linear.bias\n",
    "            self.swap_weight(weights, j, k, swap_type=\"out\")\n",
    "            self.swap_bias(biases, j, k)\n",
    "            # change output_perm\n",
    "            self.swap_bias(self.out_perm, j, k)\n",
    "        else:\n",
    "            # middle layer : (incoming, outgoing) * weights, and biases\n",
    "            weights_in = linears[i-1].linear.weight\n",
    "            weights_out = linears[i].linear.weight\n",
    "            biases = linears[i-1].linear.bias\n",
    "            self.swap_weight(weights_in, j, k, swap_type=\"out\")\n",
    "            self.swap_weight(weights_out, j, k, swap_type=\"in\")\n",
    "            self.swap_bias(biases, j, k)\n",
    "\n",
    "    def get_top_id(self, i, top_k=20):\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            # input layer\n",
    "            weights = linears[i].linear.weight\n",
    "            score = torch.sum(torch.abs(weights), dim=0)\n",
    "            in_fold = linears[0].in_fold\n",
    "            #print(score.shape)\n",
    "            score = torch.sum(score.reshape(in_fold, int(score.shape[0]/in_fold)), dim=0)\n",
    "        elif i == num_linear:\n",
    "            # output layer\n",
    "            weights = linears[i-1].linear.weight\n",
    "            score = torch.sum(torch.abs(weights), dim=1)\n",
    "        else:\n",
    "            weights_in = linears[i-1].linear.weight\n",
    "            weights_out = linears[i].linear.weight\n",
    "            score = torch.sum(torch.abs(weights_out), dim=0) + torch.sum(torch.abs(weights_in), dim=1)\n",
    "        #print(score.shape)\n",
    "        top_index = torch.flip(torch.argsort(score),[0])[:top_k]\n",
    "        return top_index, score\n",
    "\n",
    "    def relocate_ij(self, i, j):\n",
    "        # In the ith layer (of neurons), relocate the jth neuron\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i < num_linear:\n",
    "            num_neuron = int(linears[i].linear.weight.shape[1]/linears[i].in_fold)\n",
    "        else:\n",
    "            num_neuron = linears[i-1].linear.weight.shape[0]\n",
    "        ccs = []\n",
    "        for k in range(num_neuron):\n",
    "            self.swap(i,j,k)\n",
    "            ccs.append(self.get_cc())\n",
    "            self.swap(i,j,k)\n",
    "        k = torch.argmin(torch.stack(ccs))\n",
    "        self.swap(i,j,k)\n",
    "\n",
    "    def relocate_i(self, i):\n",
    "        # Relocate neurons in the ith layer\n",
    "        top_id = self.get_top_id(i, top_k=self.top_k)\n",
    "        for j in top_id[0]:\n",
    "            self.relocate_ij(i,j)\n",
    "\n",
    "    def relocate(self):\n",
    "        # Relocate neurons in the whole model\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        for i in range(num_linear+1):\n",
    "            self.relocate_i(i)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train = torchvision.datasets.MNIST(root=\"/tmp\", train=True, transform=torchvision.transforms.ToTensor(), download=True)\n",
    "# test = torchvision.datasets.MNIST(root=\"/tmp\", train=False, transform=torchvision.transforms.ToTensor(), download=True)\n",
    "# train_loader = torch.utils.data.DataLoader(train, batch_size=50, shuffle=True)\n",
    "\n",
    "def accuracy(network, dataset, device, N=2000, batch_size=50):\n",
    "    dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    for x, labels in islice(dataset_loader, N // batch_size):\n",
    "        #print(x.shape)\n",
    "        logits = network(x.to(device))\n",
    "        predicted_labels = torch.argmax(logits, dim=1)\n",
    "        correct += torch.sum(predicted_labels == labels.to(device))\n",
    "        total += x.size(0)\n",
    "    return correct / total\n",
    "\n",
    "def loss_f(network, dataset, device, N=2000, batch_size=50):\n",
    "    dataset_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "    loss = 0\n",
    "    total = 0\n",
    "    for x, labels in islice(dataset_loader, N // batch_size):\n",
    "        logits = network(x.to(device))\n",
    "        loss += torch.sum((logits-torch.eye(10,)[labels].to(device))**2)\n",
    "        total += x.size(0)\n",
    "    return loss / total\n",
    "\n",
    "train = torchvision.datasets.MNIST(root=\"/tmp\", train=True, transform=torchvision.transforms.ToTensor(), download=True)\n",
    "test = torchvision.datasets.MNIST(root=\"/tmp\", train=False, transform=torchvision.transforms.ToTensor(), download=True)\n",
    "\n",
    "data_size = 60000\n",
    "train = torch.utils.data.Subset(train, range(data_size))\n",
    "train_loader = torch.utils.data.DataLoader(train, batch_size=100, shuffle=True)\n",
    "\n",
    "def L2(model):\n",
    "    L2_ = 0.\n",
    "    for p in mlp.parameters():\n",
    "        L2_ += torch.sum(p**2)\n",
    "    return L2_\n",
    "\n",
    "def rescale(model, alpha):\n",
    "    for p in mlp.parameters():\n",
    "        p.data = alpha * p.data\n",
    "\n",
    "\n",
    "width = 200\n",
    "mlp = BioMLP2D(shp=(784,100,100,10))\n",
    "mlp.to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.AdamW(mlp.parameters(), lr=1e-3, weight_decay=0.0)\n",
    "\n",
    "one_hots = torch.eye(10, 10).to(device)\n",
    "\n",
    "mlp.eval()\n",
    "print(\"Initial accuracy: {0:.4f}\".format(accuracy(mlp, test, device)))\n",
    "\n",
    "test_accuracies = []\n",
    "train_accuracies = []\n",
    "\n",
    "step = 0\n",
    "mlp.train()\n",
    "pbar = tqdm(islice(cycle(train_loader), steps), total=steps)\n",
    "\n",
    "best_train_loss = 1e4\n",
    "best_test_loss = 1e4\n",
    "best_train_acc = 0.\n",
    "best_test_acc = 0.\n",
    "\n",
    "log = 200\n",
    "# lamb = 0.01\n",
    "# swap_log = 500\n",
    "plot_log = 500\n",
    "\n",
    "train_type = 1\n",
    "\n",
    "#train_type = 1; #no L1\n",
    "# train_type = 2; #L1\n",
    "# train_type = 3: L1 + Local\n",
    "# train_type = 4: L1 + Swap\n",
    "# train_type = 5: L1 + Local + Swap\n",
    "lamb = 0 if train_type==1 else 0.01\n",
    "swap_log = 200 if train_type >= 4 else float('inf')\n",
    "weight_factor = 2. if train_type == 3 or train_type == 5 else 0.\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "for x, label in pbar:\n",
    "\n",
    "    if step == int(steps/4):\n",
    "        lamb *= 10\n",
    "    elif step == int(steps/2):\n",
    "        lamb *= 10\n",
    "\n",
    "    mlp.train()\n",
    "    optimizer.zero_grad()\n",
    "    loss_train = loss_fn(mlp(x.to(device)), one_hots[label])\n",
    "    cc = mlp.get_cc(weight_factor, no_penalize_last=True)\n",
    "    total_loss = loss_train + lamb*cc\n",
    "    total_loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    if step % log == 0:\n",
    "        with torch.no_grad():\n",
    "            mlp.eval()\n",
    "            train_acc = accuracy(mlp, train, device).item()\n",
    "            test_acc = accuracy(mlp, test, device).item()\n",
    "            train_loss = loss_f(mlp, train, device).item()\n",
    "            test_loss = loss_f(mlp, test, device).item()\n",
    "\n",
    "            \n",
    "\n",
    "            if train_acc > best_train_acc:\n",
    "                best_train_acc = train_acc\n",
    "            if test_acc > best_test_acc:\n",
    "                best_test_acc = test_acc\n",
    "            if train_loss < best_train_loss:\n",
    "                best_train_loss = train_loss\n",
    "            if test_loss < best_test_loss:\n",
    "                best_test_loss = test_loss\n",
    "            mlp.train()\n",
    "            pbar.set_description(\"{:3.3f} | {:3.3f} | {:3.3f} | {:3.3f} | {:3.3f} \".format(train_acc, test_acc, train_loss, test_loss, cc))\n",
    "    step += 1\n",
    "\n",
    "    if step % swap_log == 0:\n",
    "        mlp.relocate()\n",
    "    \n",
    "training_time = time.time() - start \n",
    "\n",
    "# mlp.to(\"cpu\")\n",
    "# torch.save(mlp.state_dict(), 'fivemodels/bimt.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 0:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n",
      "Batch 1:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n",
      "Batch 2:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n",
      "Batch 3:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n",
      "Batch 4:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n",
      "Batch 5:\n",
      "Image from Class 3/5: torch.Size([1, 1, 28, 28])\n",
      "Image from Class 6/8: torch.Size([1, 1, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# Load the MNIST dataset\n",
    "train_dataset = torchvision.datasets.MNIST(root=\"/tmp\", train=True, transform=torchvision.transforms.ToTensor(), download=True)\n",
    "\n",
    "# Create datasets for each class\n",
    "class_3_dataset = [(image, label) for image, label in train_dataset if label == 3]\n",
    "class_5_dataset = [(image, label) for image, label in train_dataset if label == 5]\n",
    "class_6_dataset = [(image, label) for image, label in train_dataset if label == 6]\n",
    "class_8_dataset = [(image, label) for image, label in train_dataset if label == 8]\n",
    "\n",
    "# Custom dataset for paired data\n",
    "class PairedDataset(Dataset):\n",
    "    def __init__(self, dataset1, dataset2):\n",
    "        self.dataset1 = dataset1\n",
    "        self.dataset2 = dataset2\n",
    "\n",
    "    def __len__(self):\n",
    "        # Ensure both datasets are of the same length\n",
    "        return min(len(self.dataset1), len(self.dataset2))\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        image1, _ = self.dataset1[idx % len(self.dataset1)]\n",
    "        image2, _ = self.dataset2[idx % len(self.dataset2)]\n",
    "        return image1, image2\n",
    "\n",
    "# Create the paired datasets\n",
    "paired_dataset_3_8 = PairedDataset(class_3_dataset, class_8_dataset)\n",
    "paired_dataset_5_6 = PairedDataset(class_5_dataset, class_6_dataset)\n",
    "\n",
    "# Combine the paired datasets\n",
    "final_dataset = paired_dataset_3_8 + paired_dataset_5_6\n",
    "\n",
    "# Create a DataLoader for the final dataset\n",
    "final_loader = DataLoader(final_dataset, batch_size=1, shuffle=True)\n",
    "loader_3_8 = DataLoader(paired_dataset_3_8, batch_size=1, shuffle=True)\n",
    "\n",
    "# Example: Iterate over the DataLoader\n",
    "for i, (image1, image2) in enumerate(final_loader):\n",
    "    print(f\"Batch {i}:\")\n",
    "    print(\"Image from Class 3/5:\", image1.shape)\n",
    "    print(\"Image from Class 6/8:\", image2.shape)\n",
    "    if i == 5:  # Stop after 6 iterations for demonstration\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'fivemodels/bimt.pt': 0.7264511585235596,\n",
       " 'fivemodels/l1local.pt': 0.7328782081604004,\n",
       " 'fivemodels/l1only.pt': 0.7100377082824707,\n",
       " 'fivemodels/l1swap.pt': 0.7255556583404541,\n",
       " 'fivemodels/fully_dense.pt': 0.7272439002990723}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "models = [\"fivemodels/bimt.pt\", \"fivemodels/l1local.pt\", \"fivemodels/l1only.pt\", \"fivemodels/l1swap.pt\", \"fivemodels/fully_dense.pt\"]\n",
    "\n",
    "model_result = {}\n",
    "avg_inf_time ={}\n",
    "for val, model in enumerate(models):\n",
    "    mlp = BioMLP2D(shp=(784,100,100,10)).to(\"cpu\")\n",
    "    mlp.load_state_dict(torch.load(model))\n",
    "    avg_inf = 0\n",
    "    for i, (image1, image2) in enumerate(final_loader):\n",
    "        clean_tensor = image2\n",
    "        start = time.time()\n",
    "        mlp.eval()\n",
    "        with torch.no_grad():\n",
    "            og_op = mlp(clean_tensor)\n",
    "            avg_inf += time.time()-start\n",
    "    avg_inf_time[model] = avg_inf   #/len(final_loader)\n",
    "\n",
    "avg_inf_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "modmi",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}