Mnist - Test with Truncated FFF and SparseComp - Test best models - Different Sizes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "464058eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Repository initialized!\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Repository initialized!\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import dagshub\n",
    "dagshub.init(repo_owner='leocus4', repo_name='TinyFFF', mlflow=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5ad620fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on cuda\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "print(f\"Training on {DEVICE}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3bbe9647",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision.transforms import ToTensor\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.utils.data import random_split\n",
    "from sklearn.metrics import confusion_matrix, classification_report, ConfusionMatrixDisplay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "07f4adae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_size_fc(model, list_of_fc_layers, list_of_fc_sparsity, verbose=False):\n",
    "    '''\n",
    "    model has to be sublass of nn.Module\n",
    "        check the subclass with: issubclass(sub, sup), return true if sub is sublcass of sup\n",
    "                                 isinstance(sub_instance, sup), return true if is sub_instance is subclass of sup\n",
    "    list_of_fc_layers: list of fully connected layer OF THE MODEL (should be a pointer to layer of model)\n",
    "    list_of_fc_sparsity: list of the sparsity for each fully connected layer\n",
    "    '''\n",
    "    assert isinstance(model, nn.Module), \"The model is not a subclass of torch.nn.Module\"\n",
    "    assert len(list_of_fc_layers) == len(list_of_fc_sparsity), \"The lists should be of the same length\"\n",
    "    kb = 1000\n",
    "    verbose and print(\"-------------------------------------------------------------------------------------------\")\n",
    "    model_size_no_sparsity = 0\n",
    "    for param in model.parameters():\n",
    "        model_size_no_sparsity += param.nelement() * param.element_size()\n",
    "    for buffer in model.buffers():\n",
    "        model_size_no_sparsity += buffer.nelement() * buffer.element_size()\n",
    "    \n",
    "    total_size_no_sparsity = 0\n",
    "    total_size_with_sparsity = 0\n",
    "    total_size_with_sparsity_CSC = 0\n",
    "    \n",
    "    size_layer_list = []\n",
    "    num = 0\n",
    "    for fc_layer, sparsity in zip(list_of_fc_layers, list_of_fc_sparsity):\n",
    "        num += 1\n",
    "        # get size\n",
    "        verbose and print(\"Layer \" + str(num), fc_layer)\n",
    "        weight = fc_layer.nelement() * fc_layer.element_size()\n",
    "        \n",
    "        # save in no sparsity\n",
    "        total_size_no_sparsity += weight\n",
    "        \n",
    "        # set sparsity\n",
    "        weight = min(1, 2 * sparsity) * weight\n",
    "        \n",
    "        # FROM Representation\n",
    "        if (sparsity <= 0.5): # Dipende dall'analisi che vuoi fare\n",
    "            if (len(list(fc_layer.shape)) == 3):\n",
    "                verbose and print(\"Layer require additional\", fc_layer.shape[0], \"variables, total size with 4 bytes:\", fc_layer.shape[0]*4 / kb)\n",
    "                total_size_with_sparsity_CSC += (fc_layer.shape[0]*4) # number of filter\n",
    "            elif (len(list(fc_layer.shape)) == 2):\n",
    "                total_size_with_sparsity_CSC += (fc_layer.shape[1] + 1)*4 # number of column\n",
    "                verbose and print(\"Layer require additional\", fc_layer.shape[1]+1, \"variables, total size with 4 bytes:\", (fc_layer.shape[1]+1)*4 / kb)\n",
    "            \n",
    "        total_size_with_sparsity_CSC += weight\n",
    "        \n",
    "        # save in with sparsity\n",
    "        total_size_with_sparsity += weight\n",
    "        \n",
    "        size_layer_list.append(weight)\n",
    "        \n",
    "        # print total - print weight - print bias\n",
    "        verbose and print(\"Layer \"+str(num)+\":\\t\\t\", (weight) / kb,\n",
    "              \"KB, \\tweight:\\t\", weight / kb, \"KB\")\n",
    "    \n",
    "    # print total no sparisty\n",
    "    verbose and print(\"Size FC Layer (no sparsity):\\t\", total_size_no_sparsity / kb,\"KB\")\n",
    "    \n",
    "    # print total with sparsity\n",
    "    verbose and print(\"Size FC Layer (with sparsity):\\t\", total_size_with_sparsity / kb,\"KB\")\n",
    "    \n",
    "    # print model total - total no sparsity\n",
    "    verbose and print(\"Total Size no sparsity:\\t\\t\", model_size_no_sparsity / kb ,\"KB\")\n",
    "    \n",
    "    # print model total - total no sparisty + total with sparsity\n",
    "    model_size_with_sparsity = model_size_no_sparsity - total_size_no_sparsity + total_size_with_sparsity\n",
    "    verbose and print(\"Total Size with sparsity:\\t\", model_size_with_sparsity / kb,\"KB\")\n",
    "    \n",
    "    # print model total - total no sparisty + total with sparsity and CSC\n",
    "    model_size_with_sparsity_CSC = model_size_no_sparsity - total_size_no_sparsity + total_size_with_sparsity_CSC\n",
    "    verbose and print(\"Total Size with sparsity and CSC representation:\\t\", model_size_with_sparsity_CSC / kb,\"KB\")\n",
    "    \n",
    "    verbose and print(\"-------------------------------------------------------------------------------------------\")\n",
    "    \n",
    "    return model_size_with_sparsity, model_size_with_sparsity_CSC, size_layer_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3266534d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def perform_compression(model, list_of_fc_layers, list_of_fc_sparsity, learning_rate, num_epochs, train_loader,\n",
    "                        test_loader,model_device,val_loader=None, model_name=None, given_criterion=None,\n",
    "                        calculate_inputs=None,calculate_outputs=None, history=False, regularizerParam = 0):\n",
    "    '''\n",
    "    model has to be sublass of nn.Module\n",
    "        check the subclass with: issubclass(sub, sup), return true if sub is sublcass of sup\n",
    "                                 isinstance(sub_instance, sup), return true if is sub_instance is subclass of sup\n",
    "    list_of_fc_layers: list of fully connected layer OF THE MODEL (should be a pointer to layer of model)\n",
    "    list_of_fc_sparsity: list of the sparsity for each fully connected layer\n",
    "    NOTE - Sparsity applied only to weight of FC, not on bias\n",
    "    NOTE - The list are modified during execution, so are copied with list.copy() to avoid changing the original list\n",
    "    '''\n",
    "    assert isinstance(model, nn.Module), \"The model is not a subclass of torch.nn.Module\"\n",
    "    assert len(list_of_fc_layers) == len(list_of_fc_sparsity), \"The lists should be of the same length\"\n",
    "    # asset sparsity between 0 and 1\n",
    "    valid_sparsity = True\n",
    "    for sparsity in list_of_fc_sparsity:\n",
    "        if (sparsity > 1) or (sparsity < 0):\n",
    "            valid_sparsity = False\n",
    "    assert valid_sparsity, \"The sparsity value must be between 0 and 1\"\n",
    "    list_of_fc_layers = list_of_fc_layers.copy()\n",
    "    list_of_fc_sparsity = list_of_fc_sparsity.copy()\n",
    "    # The idea is get the model, set all parameter to not require gradient, set fully connected layer to require gradient,\n",
    "    # perform training\n",
    "    \n",
    "    # disabling parameters\n",
    "    for name, param in model.named_parameters():\n",
    "        print(\"Disabling:\", name)\n",
    "        param.requires_grad = False\n",
    "    \n",
    "    # activating fully connected layers only if its sparsity is > 0\n",
    "    # if a layer has sparsity equal to zero we can override with 0\n",
    "    # if all sparsity is set to 1, compression is not requested\n",
    "    sparseTraining = False\n",
    "    for fc_layer, sparsity in zip(list_of_fc_layers, list_of_fc_sparsity):\n",
    "        if (sparsity == 1):\n",
    "            #if (sparseTraining):\n",
    "            print(\"Activating:\", fc_layer.shape)\n",
    "            fc_layer.requires_grad = True\n",
    "        elif (sparsity > 0):\n",
    "            print(\"Activating:\", fc_layer.shape)\n",
    "            fc_layer.requires_grad = True\n",
    "        else:\n",
    "            fc_layer.weight = torch.nn.Parameter(torch.zeros_like(fc_layer.weight), requires_grad=False)\n",
    "            # delete from the list (since no need to update them)\n",
    "            list_of_fc_layers.remove(fc_layer)\n",
    "            list_of_fc_sparsity.remove(sparsity)\n",
    "            \n",
    "        if (sparsity < 1):\n",
    "            sparseTraining = True\n",
    "    \n",
    "    acc = 0\n",
    "    # TEST - compute accuracy\n",
    "    accuracyHistory = []\n",
    "    lastCorrect = 0\n",
    "    totalPredictions = 0\n",
    "    numberOfUpdates = len(test_loader)\n",
    "        \n",
    "    if not (sparseTraining):\n",
    "        print(\"No need to perform compression, all layers's sparsity is set to 1\")\n",
    "    else: # PERFORM TRAINING - COMPRESSION\n",
    "        \n",
    "        # set up\n",
    "        criterion = nn.NLLLoss()\n",
    "        if given_criterion:\n",
    "            criterion = given_criterion\n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "        #optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)\n",
    "        n_total_steps = len(train_loader)\n",
    "        \n",
    "        # to save best results\n",
    "        best_val_epoch, best_val_loss, best_val_acc, best_acc_epoch = 0, 1e6, 0, 0\n",
    "        \n",
    "        for epoch in range(num_epochs):\n",
    "            \n",
    "            model.train()\n",
    "            for i, (inputs, labels) in enumerate(train_loader):\n",
    "                \n",
    "                inputs = inputs.to(model_device)\n",
    "                labels = labels.to(model_device)\n",
    "                                \n",
    "                # Forward pass\n",
    "                \n",
    "                # preforward\n",
    "                if calculate_inputs:\n",
    "                    inputs = calculate_inputs(inputs)\n",
    "                \n",
    "                # forward\n",
    "                if calculate_outputs:\n",
    "                    outputs = calculate_outputs(inputs)\n",
    "                else:\n",
    "                    outputs = model.forward(inputs)\n",
    "                \n",
    "                # Regularization\n",
    "                regularizer = 0\n",
    "                if (regularizerParam != 0):\n",
    "                    for layer in list_of_fc_layers:\n",
    "                        regularizer += (torch.norm(layer.weight)**2)\n",
    "                # Loss\n",
    "                loss = criterion(outputs, labels) + (regularizer * regularizerParam)\n",
    "                \n",
    "                # Backward and optimize\n",
    "                optimizer.zero_grad()\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                \n",
    "                # apply hardthreshold - in the list we have only layer with require_grad = True\n",
    "                for fc_layer, sparsity in zip(list_of_fc_layers, list_of_fc_sparsity):\n",
    "                    layer = fc_layer.data\n",
    "                    new_layer = hardThreshold(layer, sparsity)\n",
    "                    with torch.no_grad():\n",
    "                        fc_layer.data = torch.FloatTensor(new_layer).to(model_device)\n",
    "                \n",
    "                # print Accuracy\n",
    "                if (i+1) % 100 == 0:\n",
    "                    print (f'Epoch [{epoch+1}/{num_epochs}], Step[{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')\n",
    "            \n",
    "            print (f'Epoch [{epoch+1}/{num_epochs}], Step[{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')\n",
    "            \n",
    "            # Use Validation Set at each epochs to pick the most \n",
    "            if (val_loader and model_name):\n",
    "                model.eval()\n",
    "                with torch.no_grad():\n",
    "                    v_loss = 0\n",
    "                    n_correct = 0\n",
    "                    n_samples = 0\n",
    "                    n_iterations = 0\n",
    "                    for inputs, labels in test_loader:\n",
    "                        inputs = inputs.to(model_device)\n",
    "                        labels = labels.to(model_device)\n",
    "                        # Forward pass\n",
    "                \n",
    "                        # preforward\n",
    "                        if calculate_inputs:\n",
    "                            inputs = calculate_inputs(inputs)\n",
    "                        outputs = 0 \n",
    "                        # forward\n",
    "                        if calculate_outputs:\n",
    "                            outputs = calculate_outputs(inputs)\n",
    "                        else:\n",
    "                            outputs = model.forward(inputs)\n",
    "                        \n",
    "                        # for calculating v_loss\n",
    "                        loss = criterion(outputs, labels)                       \n",
    "                        v_loss += loss.item()\n",
    "                        n_iterations += 1\n",
    "                        \n",
    "                        # max returns (value, index)\n",
    "                        _, predicted = torch.max(outputs.data, 1)\n",
    "                        n_samples += labels.size(0)\n",
    "                        n_correct += (predicted == labels).sum().item()\n",
    "                    \n",
    "                    # Val test completed, now checking the results\n",
    "                    v_loss = v_loss/(n_iterations)\n",
    "                    v_loss = round(v_loss, 5)\n",
    "                    v_acc = round(100*(n_correct / n_samples), 5)\n",
    "                    \n",
    "                    if v_acc >= best_val_acc:\n",
    "                        torch.save(model.state_dict(), model_name+\"_acc.h5\")\n",
    "                        best_acc_epoch = epoch + 1\n",
    "                        best_val_acc = v_acc\n",
    "                    if v_loss <= best_val_loss:\n",
    "                        torch.save(model.state_dict(), model_name+\".h5\")\n",
    "                        best_val_epoch = epoch + 1\n",
    "                        best_val_loss = v_loss\n",
    "                    #print(f'Epoch[{epoch+1}]: t_loss: {t_loss} t_acc: {t_acc} v_loss: {v_loss} v_acc: {v_acc}')\n",
    "                    print(f'Epoch[{epoch+1}]: v_loss: {v_loss} v_acc: {v_acc}')\n",
    "        \n",
    "        \n",
    "        # Use Validation Set at each epochs to pick the most \n",
    "        if (val_loader and model_name):\n",
    "            model.load_state_dict(torch.load(model_name+\".h5\", map_location='cpu'))\n",
    "            print('Best model saved at epoch: ', best_val_epoch)\n",
    "            print('Best acc model saved at epoch: ', best_acc_epoch)\n",
    "        \n",
    "        # USING TEST SET TO CHECK ACCURACY\n",
    "        #model.eval()\n",
    "        with torch.no_grad():\n",
    "            n_correct = 0\n",
    "            n_samples = 0\n",
    "            for inputs, labels in test_loader:\n",
    "                inputs = inputs.to(model_device)\n",
    "                labels = labels.to(model_device)\n",
    "                   # Forward pass\n",
    "                \n",
    "                # preforward\n",
    "                if calculate_inputs:\n",
    "                    inputs = calculate_inputs(inputs)\n",
    "                outputs = 0 \n",
    "                # forward\n",
    "                if calculate_outputs:\n",
    "                    outputs = calculate_outputs(inputs)\n",
    "                else:\n",
    "                    outputs = model.forward(inputs)\n",
    "                # max returns (value, index)\n",
    "                \n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                n_samples += labels.size(0)\n",
    "                n_correct += (predicted == labels).sum().item()                \n",
    "            acc = 100.0 * n_correct / n_samples\n",
    "            totalPredictions = n_samples\n",
    "            print(f'Accuracy of the network on the 10000 test images: {acc} %')\n",
    "\n",
    "        \n",
    "    result = {\n",
    "        'correctPredictions': lastCorrect,\n",
    "        'totalPredictions': totalPredictions,\n",
    "        'accuracyThroughEpochs': accuracyHistory,\n",
    "        'numberOfUpdate': numberOfUpdates,\n",
    "    }\n",
    "    \n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e36e4acf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision.datasets import CIFAR10\n",
    "from torchvision.transforms import ToTensor\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.utils.data import random_split\n",
    "from sklearn.metrics import confusion_matrix, classification_report, ConfusionMatrixDisplay\n",
    "\n",
    "\n",
    "def print_full_model(model):\n",
    "    assert isinstance(model, nn.Module), \"The model is not a subclass of torch.nn.Module\"\n",
    "    kb = 1000\n",
    "    model_size = 0\n",
    "    for name, param in model.named_parameters():\n",
    "        layer_size = param.nelement() * param.element_size()\n",
    "        model_size += layer_size\n",
    "        print(name,\"\\t\", param.nelement(), \"\\t\", param.element_size(),\"\\t\", layer_size / kb, \"KB\")\n",
    "        \n",
    "    for name, buffer in model.named_buffers():\n",
    "        layer_size = buffer.nelement() * buffer.element_size()\n",
    "        model_size += layer_size\n",
    "        print(name,\"\\t\", layer_size / kb, \"KB\")\n",
    "    print(\"Model Size:\", model_size / kb, \"KB\")\n",
    "\n",
    "def hardThreshold(A: torch.Tensor, sparsity):\n",
    "    '''\n",
    "    Given a Tensor A and the correponding sparsity, returns a copy in the\n",
    "    format of numpy array with the constraint applied\n",
    "    '''\n",
    "    matrix_A = A.data.cpu().detach().numpy().ravel()    \n",
    "    if len(matrix_A) > 0:\n",
    "        threshold = np.percentile(np.abs(matrix_A), (1 - sparsity) * 100.0, method='higher')\n",
    "        matrix_A[np.abs(matrix_A) < threshold] = 0.0\n",
    "    matrix_A = matrix_A.reshape(A.shape)\n",
    "    return matrix_A\n",
    "\n",
    "def get_layers(model):\n",
    "    \"\"\"Recursively get all layers in a PyTorch model.\"\"\"\n",
    "    list_layers = []\n",
    "    # for name, module in model.named_children():\n",
    "    #     # check type of module\n",
    "    #     is_conv1d = isinstance(module, torch.nn.Conv1d)\n",
    "    #     is_conv2d = isinstance(module, torch.nn.Conv2d)\n",
    "    #     is_linear = isinstance(module, torch.nn.Linear)\n",
    "    #     is_sequential = isinstance(module, torch.nn.Sequential)\n",
    "    #     if (is_conv1d or is_conv2d or is_linear):\n",
    "    #         list_layers.append(module)\n",
    "    #     if (is_sequential):\n",
    "    #         for sub_name, sub_module in module.named_children():\n",
    "    #             print(sub_name)\n",
    "    #             # check type of module\n",
    "    #             is_conv1d = isinstance(sub_module, torch.nn.Conv1d)\n",
    "    #             is_conv2d = isinstance(sub_module, torch.nn.Conv2d)\n",
    "    #             is_linear = isinstance(sub_module, torch.nn.Linear)\n",
    "    #             if (is_conv1d or is_conv2d or is_linear):\n",
    "    #                 list_layers.append(sub_module)\n",
    "    for layer in model.children():\n",
    "        if isinstance(layer, nn.Sequential):\n",
    "            # If it's a sequential container, recursively get its layers\n",
    "            list_layers.extend(get_layers(layer))\n",
    "        else:\n",
    "            # If it's a single layer, add it to the list\n",
    "            if (isinstance(layer, torch.nn.Conv1d) or isinstance(layer, torch.nn.Conv2d) or isinstance(layer, torch.nn.Linear)):\n",
    "                list_layers.append(layer)\n",
    "    return list_layers\n",
    "\n",
    "def apply_sparsity(model, list_of_fc_layers, list_of_fc_sparsity, model_device):\n",
    "    assert isinstance(model, nn.Module), \"The model is not a subclass of torch.nn.Module\"\n",
    "    assert len(list_of_fc_layers) == len(list_of_fc_sparsity), \"The lists should be of the same length\"\n",
    "    # asset sparsity between 0 and 1\n",
    "    valid_sparsity = True\n",
    "    for sparsity in list_of_fc_sparsity:\n",
    "        if (sparsity > 1) or (sparsity < 0):\n",
    "            valid_sparsity = False\n",
    "    assert valid_sparsity, \"The sparsity value must be between 0 and 1\"\n",
    "    \n",
    "    list_of_fc_layers = list_of_fc_layers.copy()\n",
    "    list_of_fc_sparsity = list_of_fc_sparsity.copy()\n",
    "    \n",
    "    # apply hardthreshold - in the list we have only layer with require_grad = True\n",
    "    for fc_layer, sparsity in zip(list_of_fc_layers, list_of_fc_sparsity):\n",
    "        layer = fc_layer.weight.data\n",
    "        new_layer = hardThreshold(layer, sparsity)\n",
    "        with torch.no_grad():\n",
    "            fc_layer.weight.data = torch.FloatTensor(new_layer).to(model_device)\n",
    "    \n",
    "def calculate_accuracy(model, train_loader, test_loader, model_device, calculate_inputs=None, calculate_outputs=None):\n",
    "    assert isinstance(model, nn.Module), \"The model is not a subclass of torch.nn.Module\"\n",
    "\n",
    "    acc = 0\n",
    "\n",
    "    # TEST - compute accuracy\n",
    "    with torch.no_grad():\n",
    "        n_correct = 0\n",
    "        n_samples = 0\n",
    "        for inputs, labels in train_loader:\n",
    "            inputs = inputs.to(model_device)\n",
    "            labels = labels.to(model_device)\n",
    "            # preforward\n",
    "            if calculate_inputs:\n",
    "                inputs = calculate_inputs(inputs)\n",
    "            outputs = 0 \n",
    "            # forward\n",
    "            if calculate_outputs:\n",
    "                outputs = calculate_outputs(inputs)\n",
    "            else:\n",
    "                outputs = model.forward(inputs)\n",
    "                    \n",
    "            # max returns (value, index)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            n_samples += labels.size(0)\n",
    "            n_correct += (predicted == labels).sum().item()\n",
    "\n",
    "        acc = 100.0 * n_correct / n_samples\n",
    "        print(f'Accuracy of the network on the train images: {acc} %')\n",
    "\n",
    "        n_correct = 0\n",
    "        n_samples = 0\n",
    "        for inputs, labels in test_loader:\n",
    "            inputs = inputs.to(model_device)\n",
    "            labels = labels.to(model_device)\n",
    "            # preforward\n",
    "            if calculate_inputs:\n",
    "                inputs = calculate_inputs(inputs)\n",
    "            outputs = 0 \n",
    "            # forward\n",
    "            if calculate_outputs:\n",
    "                outputs = calculate_outputs(inputs)\n",
    "            else:\n",
    "                outputs = model.forward(inputs)\n",
    "            \n",
    "            # max returns (value, index)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            n_samples += labels.size(0)\n",
    "            n_correct += (predicted == labels).sum().item()\n",
    "\n",
    "        acc = 100.0 * n_correct / n_samples\n",
    "        print(f'Accuracy of the network on the 10000 test images: {acc} %')\n",
    "\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e3ae4b4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_sparsity_for_layers(layer_list):\n",
    "    \"\"\"Compute sparsity for each layer in a list of layers.\"\"\"\n",
    "    sparsity_info = []\n",
    "\n",
    "    for layer in layer_list:\n",
    "        weight = layer.data\n",
    "        total_elements = weight.numel()\n",
    "        zero_elements = (weight == 0).sum().item()\n",
    "        sparsity = zero_elements / total_elements\n",
    "        sparsity_info.append((layer.__class__.__name__, sparsity, total_elements, zero_elements))\n",
    "    \n",
    "    # Print the sparsity information for each layer\n",
    "    for layer, sparsity, total_elements, zero_elements in sparsity_info:\n",
    "        print(f'Layer: {layer}, Sparsity: {1-sparsity:.4f}, Total Elements: {total_elements}, Zero Elements: {zero_elements}')\n",
    "\n",
    "    return sparsity_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "06b901a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastfeedforward import FFF\n",
    "\n",
    "def train(net, trainloader, epochs, norm_weight=0.0):\n",
    "    \"\"\"Train the network on the training set.\"\"\"\n",
    "    # Define loss and optimizer\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)\n",
    "\n",
    "    # Train the network for the given number of epochs\n",
    "    for _ in range(epochs):\n",
    "        # Iterate over data\n",
    "        for images, labels in trainloader:\n",
    "            images, labels = images.to(DEVICE), labels.to(DEVICE)\n",
    "            optimizer.zero_grad()\n",
    "            loss = criterion(net(images), labels)\n",
    "            if norm_weight != 0:\n",
    "                loss += norm_weight * net.fff.w1s.pow(2).sum()\n",
    "                loss += norm_weight * net.fff.w2s.pow(2).sum()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "\n",
    "def test(net, testloader):\n",
    "    \"\"\"Validate the network on the entire test set.\"\"\"\n",
    "    # Define loss and optimizer\n",
    "    criterion = torch.nn.CrossEntropyLoss()\n",
    "    correct, total, loss = 0, 0, 0.0\n",
    "    # Train the network for the given number of epochs\n",
    "    with torch.no_grad():\n",
    "        # Iterate over data\n",
    "        for data in testloader:\n",
    "            images, labels = data[0].to(DEVICE), data[1].to(DEVICE)\n",
    "            outputs = net(images)\n",
    "            loss += criterion(outputs, labels).item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    accuracy = correct / total\n",
    "    return loss, accuracy\n",
    "\n",
    "\n",
    "class Net(torch.nn.Module):\n",
    "    def __init__(self, input_width, leaf_width, output_width, depth, dropout, region_leak):\n",
    "        super(Net, self).__init__()\n",
    "        self.fff = FFF(input_width, leaf_width, output_width, depth, torch.nn.ReLU(), dropout, train_hardened=True, region_leak=region_leak)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(len(x), -1)\n",
    "        x = self.fff(x)\n",
    "        x = torch.nn.functional.softmax(x, -1)\n",
    "        return x\n",
    "\n",
    "    def parameters(self):\n",
    "        return self.fff.parameters()\n",
    "\n",
    "\n",
    "class FF(torch.nn.Module):\n",
    "    def __init__(self, input_width, layer_width, output_width):\n",
    "        super(FF, self).__init__()\n",
    "        self.fc1 = torch.nn.Linear(input_width, layer_width)\n",
    "        self.fc2 = torch.nn.Linear(layer_width, output_width)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(len(x), -1)\n",
    "        x = torch.nn.functional.relu(self.fc1(x))\n",
    "        x = torch.nn.functional.softmax(self.fc2(x), -1)\n",
    "        return x\n",
    "\n",
    "    def parameters(self):\n",
    "        return [*self.fc1.parameters(), *self.fc2.parameters()]\n",
    "\n",
    "\n",
    "def compute_n_params(input_width: int, l_w: int, depth: int, output_width: int):\n",
    "    fff = Net(input_width, l_w, output_width, depth, 0, 0)\n",
    "    ff = FF(input_width, l_w, output_width)\n",
    "\n",
    "    n_ff = 0\n",
    "    n_fff = 0\n",
    "    for p in ff.parameters():\n",
    "        n_ff += p.numel()\n",
    "    for i, p in enumerate(fff.parameters()):\n",
    "        print(f\"[{i}-th layer]: {p.shape}\")\n",
    "        n_fff += p.numel()\n",
    "\n",
    "    print(f\"FFF: {n_fff}\\nFF: {n_ff}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f3e192bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import mlflow\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from time import time\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def get_dist(net, testloader):\n",
    "    \"\"\"\n",
    "    Returns the distribution of samples throughout the tree.\n",
    "    \"\"\"\n",
    "\n",
    "    y = []\n",
    "    l = []\n",
    "    with torch.no_grad():\n",
    "        # Iterate over data\n",
    "        for data in testloader:\n",
    "            images, labels = data[0].to(DEVICE), data[1].to(DEVICE)\n",
    "            outputs, leaves = net.forward(images, return_nodes=True)\n",
    "            y.append(labels)\n",
    "            l.append(leaves)\n",
    "    y = torch.concat(y, 0)\n",
    "    l = torch.concat(l, 0)\n",
    "    return y, l\n",
    "\n",
    "\n",
    "class FFFWrapper(torch.nn.Module):\n",
    "    def __init__(self, fff):\n",
    "        super(FFFWrapper, self).__init__()\n",
    "        self._fff = fff\n",
    "        self._fastinference = [None for i in range(2 ** (self._fff.fff.depth.item()))]\n",
    "\n",
    "    def forward(self, x, return_nodes=False):\n",
    "        \"\"\"\n",
    "        Override the forward method in order to log the data distribution.\n",
    "        \"\"\"\n",
    "        x = x.view(len(x), -1)\n",
    "        original_shape = x.shape\n",
    "        batch_size = x.shape[0]\n",
    "        last_node = torch.zeros(len(x))\n",
    "\n",
    "        current_nodes = torch.zeros((batch_size,), dtype=torch.long, device=x.device)\n",
    "        for i in range(self._fff.fff.depth.item()):\n",
    "            plane_coeffs = self._fff.fff.node_weights.index_select(dim=0, index=current_nodes)\n",
    "            plane_offsets = self._fff.fff.node_biases.index_select(dim=0, index=current_nodes)\n",
    "            plane_coeff_score = torch.bmm(x.unsqueeze(1), plane_coeffs.unsqueeze(-1))\n",
    "            plane_score = plane_coeff_score.squeeze(-1) + plane_offsets\n",
    "            plane_choices = (plane_score.squeeze(-1) >= 0).long()\n",
    "\n",
    "            platform = torch.tensor(2 ** i - 1, dtype=torch.long, device=x.device)\n",
    "            next_platform = torch.tensor(2 ** (i+1) - 1, dtype=torch.long, device=x.device)\n",
    "            current_nodes = (current_nodes - platform) * 2 + plane_choices + next_platform\n",
    "\n",
    "        leaves = current_nodes - next_platform\n",
    "        new_logits = torch.empty((batch_size, self._fff.fff.output_width), dtype=torch.float, device=x.device)\n",
    "        last_node = leaves\n",
    "\n",
    "        for i in range(leaves.shape[0]):\n",
    "            leaf_index = leaves[i]\n",
    "            if self._fastinference[leaf_index] is not None:\n",
    "                new_logits[i] = self._fastinference[leaf_index]\n",
    "            else:\n",
    "                logits = torch.matmul( x[i].unsqueeze(0), self._fff.fff.w1s[leaf_index])\n",
    "                logits += self._fff.fff.b1s[leaf_index].unsqueeze(-2)\n",
    "                activations = self._fff.fff.activation(logits)\n",
    "                new_logits[i] = torch.matmul( activations, self._fff.fff.w2s[leaf_index]).squeeze(-2)\n",
    "\n",
    "        if return_nodes:\n",
    "            return new_logits.view(*original_shape[:-1], self._fff.fff.output_width), last_node\n",
    "        return new_logits.view(*original_shape[:-1], self._fff.fff.output_width)\n",
    "\n",
    "\n",
    "    def simplify_leaves(self, trainloader):\n",
    "        y, leaves = (get_dist(self, trainloader))\n",
    "        y = y.cpu().detach().numpy()\n",
    "        outputs = y.max() + 1\n",
    "        leaves = leaves.cpu().detach().numpy()\n",
    "\n",
    "        n_simplifications = 0\n",
    "        ratios = {}\n",
    "        for l in np.unique(leaves):\n",
    "            ratios[l] = torch.zeros(outputs)\n",
    "            indices = leaves == l\n",
    "\n",
    "            for i in range(outputs):\n",
    "                ratios[l][i] = (np.sum(y[indices] == i) / np.sum(indices))\n",
    "\n",
    "            argmax = np.argmax(ratios[l])\n",
    "            if ratios[l][argmax] > 0.7:\n",
    "                output = torch.zeros(outputs)\n",
    "                output[argmax] = 1\n",
    "                self._fastinference[l] = output\n",
    "                n_simplifications += 1\n",
    "                print(f\"Leaf {l} has been replaced with {argmax}\")\n",
    "        print(self._fastinference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "94b9bfb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'trainset': 60000, 'testset': 10000}\n"
     ]
    }
   ],
   "source": [
    "from tqdm import trange\n",
    "from torchvision.datasets import MNIST\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\"\"\"Load CIFAR-10 (training and test set).\"\"\"\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,)),\n",
    "    ]\n",
    ")\n",
    "trainset = MNIST(\"../data\", train=True,  download=True, transform=transform)\n",
    "testset = MNIST(\"../data\",  train=False, download=True, transform=transform)\n",
    "\n",
    "# Select class to keep \n",
    "trainloader = DataLoader(trainset, batch_size=1024, shuffle=True)\n",
    "testloader = DataLoader(testset, batch_size=1024)\n",
    "\n",
    "num_examples = {\"trainset\" : len(trainset), \"testset\" : len(testset)}\n",
    "\n",
    "print(num_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a766c4fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "list_of_run = [\n",
    "'543fe9acb34441a3a82b09ca2ef6046c',\n",
    "'4ea2391b3d014e2fafff3accfb352d2c',\n",
    "'612d1573468c4ffaa89bf54d60ce4508',\n",
    "'24879e69bd174ef6a9973359ea1b9b6c',\n",
    "'7c66090a514e4080b639b6f261d5a134',\n",
    "'a258b245549043ef84a9beff50896872',\n",
    "'57f3c5cd82dc480c94e516ab34620331',\n",
    "'cef7d9cd3ea548b783a400984a7145fc',\n",
    "'706ca101f6a2446db81d58abdf6e2815',\n",
    "'2a116553f512425a9409bd968b9fe8ef',\n",
    "'90b23d76307e425ea4ba7d647c7cb7f6',\n",
    "'57e917cd181c456bbf3c365b5018a2d6',\n",
    "'715fba635c6145b185c29e6aa6b6bcbf',\n",
    "'652dc0e8fe0042309b990da5fc377d60',\n",
    "'580862314acb40cbb6411391f6def1e1',\n",
    "'0c4b701ba2c2434b99a83f0b771b3945',\n",
    "'755af7caf9e74fbdbc2dee292dd8b3d1',\n",
    "'e6c2200cd3a942b081e77a4fcbb21df6',\n",
    "'9ac931f6215644bf9d22e6fcfc7f179f',\n",
    "'69d271a25d744404ad63c43b575192a6',\n",
    "'27f4eafb191340f592dfab6992d3700d',\n",
    "'9d555e630619470c8e4a6d075ba0e65f',\n",
    "'89e999f7bcce4b238cd3bca960ec27da',\n",
    "'15d383cd1c044c8cb7cf5b5de6955b13',\n",
    "]\n",
    "\n",
    "# mlflow.artifacts.download_artifacts(run_id=run_id, dst_path=\".\")\n",
    "# wrapped_model = pickle.load(open(\"./truncated_model.pkl\", \"rb\"))\n",
    "# wrapped_model._fff.fff.depth.item()\n",
    "# wrapped_model._fff.fff.input_width\n",
    "# wrapped_model._fff.fff.leaf_width\n",
    "# wrapped_model._fff.fff.output_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "959d2c7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training variable\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "learning_rate = 0.001\n",
    "num_epochs = 7\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "# Dataset\n",
    "batch_size = 1024\n",
    "val_size = 5000\n",
    "train_size = len(trainset) - val_size\n",
    "train_ds, val_ds = random_split(trainset, [train_size, val_size])\n",
    "train_loader = DataLoader(train_ds, batch_size, shuffle=True, num_workers=4)\n",
    "val_loader = DataLoader(val_ds, batch_size, num_workers=4)\n",
    "test_loader = DataLoader(testset, batch_size, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2ee8c3cd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training on cuda\n",
      "Run:\t 543fe9acb34441a3a82b09ca2ef6046c\n",
      "Depth:\t 4\n",
      "Input:\t 784\n",
      "Output:\t 10\n",
      "Leaf:\t 32\n",
      "Buffer:\t 90\n",
      "1 iteration -  Size: 1515408.8 [1, 1, 0.45, 1, 1, 1]\n",
      "2 iteration -  Size: 1370901.9200000002 [1, 1, 0.405, 1, 1, 1]\n",
      "3 iteration -  Size: 1240845.7280000001 [1, 1, 0.36450000000000005, 1, 1, 1]\n",
      "4 iteration -  Size: 1123795.1552000002 [1, 1, 0.32805000000000006, 1, 1, 1]\n",
      "5 iteration -  Size: 1018449.6396800001 [1, 1, 0.29524500000000004, 1, 1, 1]\n",
      "6 iteration -  Size: 923638.6757120001 [1, 1, 0.2657205, 1, 1, 1]\n",
      "7 iteration -  Size: 838308.8081408 [1, 1, 0.23914845, 1, 1, 1]\n",
      "8 iteration -  Size: 761511.92732672 [1, 1, 0.215233605, 1, 1, 1]\n",
      "9 iteration -  Size: 692394.734594048 [1, 1, 0.1937102445, 1, 1, 1]\n",
      "10 iteration -  Size: 630189.2611346432 [1, 1, 0.17433922005, 1, 1, 1]\n",
      "11 iteration -  Size: 574204.3350211789 [1, 1, 0.156905298045, 1, 1, 1]\n",
      "12 iteration -  Size: 523817.901519061 [1, 1, 0.1412147682405, 1, 1, 1]\n",
      "13 iteration -  Size: 478470.1113671549 [1, 1, 0.12709329141645, 1, 1, 1]\n",
      "14 iteration -  Size: 437657.1002304394 [1, 1, 0.114383962274805, 1, 1, 1]\n",
      "15 iteration -  Size: 400925.3902073955 [1, 1, 0.10294556604732451, 1, 1, 1]\n",
      "16 iteration -  Size: 367866.8511866559 [1, 1, 0.09265100944259205, 1, 1, 1]\n",
      "17 iteration -  Size: 338114.1660679903 [1, 1, 0.08338590849833284, 1, 1, 1]\n",
      "18 iteration -  Size: 311336.7494611913 [1, 1, 0.07504731764849956, 1, 1, 1]\n",
      "19 iteration -  Size: 287237.07451507216 [1, 1, 0.0675425858836496, 1, 1, 1]\n",
      "20 iteration -  Size: 265547.36706356495 [1, 1, 0.06078832729528464, 1, 1, 1]\n",
      "21 iteration -  Size: 246026.63035720843 [1, 1, 0.054709494565756175, 1, 1, 1]\n",
      "22 iteration -  Size: 228457.96732148758 [1, 1, 0.049238545109180555, 1, 1, 1]\n",
      "23 iteration -  Size: 212646.1705893388 [1, 1, 0.0443146905982625, 1, 1, 1]\n",
      "24 iteration -  Size: 198415.55353040493 [1, 1, 0.03988322153843625, 1, 1, 1]\n",
      "25 iteration -  Size: 185607.99817736444 [1, 1, 0.03589489938459262, 1, 1, 1]\n",
      "26 iteration -  Size: 174081.198359628 [1, 1, 0.03230540944613336, 1, 1, 1]\n",
      "27 iteration -  Size: 163707.0785236652 [1, 1, 0.029074868501520024, 1, 1, 1]\n",
      "28 iteration -  Size: 154370.37067129868 [1, 1, 0.026167381651368022, 1, 1, 1]\n",
      "29 iteration -  Size: 145967.3336041688 [1, 1, 0.023550643486231218, 1, 1, 1]\n",
      "30 iteration -  Size: 138404.60024375195 [1, 1, 0.021195579137608098, 1, 1, 1]\n",
      "31 iteration -  Size: 131598.14021937674 [1, 1, 0.019076021223847286, 1, 1, 1]\n",
      "32 iteration -  Size: 125472.32619743906 [1, 1, 0.017168419101462558, 1, 1, 1]\n",
      "33 iteration -  Size: 119959.09357769515 [1, 1, 0.015451577191316302, 1, 1, 1]\n",
      "34 iteration -  Size: 114997.18421992564 [1, 1, 0.013906419472184673, 1, 1, 1]\n",
      "35 iteration -  Size: 113433.18421992564 [0.45, 1, 0.013906419472184673, 1, 1, 1]\n",
      "36 iteration -  Size: 108967.46579793308 [0.45, 1, 0.012515777524966205, 1, 1, 1]\n",
      "37 iteration -  Size: 104733.86579793307 [0.405, 1, 0.012515777524966205, 1, 1, 1]\n",
      "38 iteration -  Size: 100714.71921813977 [0.405, 1, 0.011264199772469584, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.6933\n",
      "Epoch[1]: v_loss: 0.7184 v_acc: 83.66\n",
      "Epoch [2/7], Step[54/54], Loss: 0.5918\n",
      "Epoch[2]: v_loss: 0.57483 v_acc: 85.51\n",
      "Epoch [3/7], Step[54/54], Loss: 0.5116\n",
      "Epoch[3]: v_loss: 0.49398 v_acc: 87.0\n",
      "Epoch [4/7], Step[54/54], Loss: 0.4398\n",
      "Epoch[4]: v_loss: 0.44195 v_acc: 87.71\n",
      "Epoch [5/7], Step[54/54], Loss: 0.3803\n",
      "Epoch[5]: v_loss: 0.40607 v_acc: 88.55\n",
      "Epoch [6/7], Step[54/54], Loss: 0.3876\n",
      "Epoch[6]: v_loss: 0.37988 v_acc: 89.2\n",
      "Epoch [7/7], Step[54/54], Loss: 0.4048\n",
      "Epoch[7]: v_loss: 0.36111 v_acc: 89.63\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 89.63 %\n",
      "39 iteration -  Size: 96904.47921813978 [0.36450000000000005, 1, 0.011264199772469584, 1, 1, 1]\n",
      "40 iteration -  Size: 93287.24729632579 [0.36450000000000005, 1, 0.010137779795222625, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.3574\n",
      "Epoch[1]: v_loss: 0.35623 v_acc: 89.7\n",
      "Epoch [2/7], Step[54/54], Loss: 0.3431\n",
      "Epoch[2]: v_loss: 0.33983 v_acc: 90.16\n",
      "Epoch [3/7], Step[54/54], Loss: 0.3574\n",
      "Epoch[3]: v_loss: 0.32749 v_acc: 90.46\n",
      "Epoch [4/7], Step[54/54], Loss: 0.3348\n",
      "Epoch[4]: v_loss: 0.31912 v_acc: 90.76\n",
      "Epoch [5/7], Step[54/54], Loss: 0.3458\n",
      "Epoch[5]: v_loss: 0.31115 v_acc: 91.01\n",
      "Epoch [6/7], Step[54/54], Loss: 0.3031\n",
      "Epoch[6]: v_loss: 0.30534 v_acc: 91.12\n",
      "Epoch [7/7], Step[54/54], Loss: 0.2615\n",
      "Epoch[7]: v_loss: 0.30065 v_acc: 91.21\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 91.21 %\n",
      "41 iteration -  Size: 89858.03129632579 [0.32805000000000006, 1, 0.010137779795222625, 1, 1, 1]\n",
      "42 iteration -  Size: 86602.52256669321 [0.32805000000000006, 1, 0.009124001815700363, 1, 1, 1]\n",
      "43 iteration -  Size: 83516.22816669321 [0.29524500000000004, 1, 0.009124001815700363, 1, 1, 1]\n",
      "44 iteration -  Size: 80586.27031002389 [0.29524500000000004, 1, 0.008211601634130327, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.2967\n",
      "Epoch[1]: v_loss: 0.3426 v_acc: 90.11\n",
      "Epoch [2/7], Step[54/54], Loss: 0.2954\n",
      "Epoch[2]: v_loss: 0.33295 v_acc: 90.33\n",
      "Epoch [3/7], Step[54/54], Loss: 0.3869\n",
      "Epoch[3]: v_loss: 0.32564 v_acc: 90.44\n",
      "Epoch [4/7], Step[54/54], Loss: 0.2932\n",
      "Epoch[4]: v_loss: 0.32106 v_acc: 90.73\n",
      "Epoch [5/7], Step[54/54], Loss: 0.3327\n",
      "Epoch[5]: v_loss: 0.31648 v_acc: 90.72\n",
      "Epoch [6/7], Step[54/54], Loss: 0.2948\n",
      "Epoch[6]: v_loss: 0.31311 v_acc: 90.89\n",
      "Epoch [7/7], Step[54/54], Loss: 0.3171\n",
      "Epoch[7]: v_loss: 0.31072 v_acc: 90.95\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 90.95 %\n",
      "45 iteration -  Size: 77808.60535002389 [0.2657205, 1, 0.008211601634130327, 1, 1, 1]\n",
      "46 iteration -  Size: 75171.6432790215 [0.2657205, 1, 0.007390441470717294, 1, 1, 1]\n",
      "47 iteration -  Size: 72671.7448150215 [0.23914845, 1, 0.007390441470717294, 1, 1, 1]\n",
      "48 iteration -  Size: 70298.47895111935 [0.23914845, 1, 0.006651397323645565, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.3247\n",
      "Epoch[1]: v_loss: 0.36717 v_acc: 88.93\n",
      "Epoch [2/7], Step[54/54], Loss: 0.3412\n",
      "Epoch[2]: v_loss: 0.35319 v_acc: 89.23\n",
      "Epoch [3/7], Step[54/54], Loss: 0.3324\n",
      "Epoch[3]: v_loss: 0.34347 v_acc: 89.44\n",
      "Epoch [4/7], Step[54/54], Loss: 0.3422\n",
      "Epoch[4]: v_loss: 0.3365 v_acc: 89.73\n",
      "Epoch [5/7], Step[54/54], Loss: 0.3657\n",
      "Epoch[5]: v_loss: 0.33066 v_acc: 89.92\n",
      "Epoch [6/7], Step[54/54], Loss: 0.3433\n",
      "Epoch[6]: v_loss: 0.32603 v_acc: 89.93\n",
      "Epoch [7/7], Step[54/54], Loss: 0.3340\n",
      "Epoch[7]: v_loss: 0.32227 v_acc: 90.13\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 90.13 %\n",
      "49 iteration -  Size: 68048.57033351935 [0.215233605, 1, 0.006651397323645565, 1, 1, 1]\n",
      "50 iteration -  Size: 65912.63105600742 [0.215233605, 1, 0.005986257591281009, 1, 1, 1]\n",
      "51 iteration -  Size: 63928.63105600742 [0.215233605, 1, 0.005986257591281009, 1, 0.45, 1]\n",
      "52 iteration -  Size: 61903.713300167416 [0.1937102445, 1, 0.005986257591281009, 1, 0.45, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/7], Step[54/54], Loss: 0.3888\n",
      "Epoch[1]: v_loss: 0.37928 v_acc: 88.75\n",
      "Epoch [2/7], Step[54/54], Loss: 0.3642\n",
      "Epoch[2]: v_loss: 0.3628 v_acc: 89.15\n",
      "Epoch [3/7], Step[54/54], Loss: 0.3520\n",
      "Epoch[3]: v_loss: 0.35404 v_acc: 89.46\n",
      "Epoch [4/7], Step[54/54], Loss: 0.3451\n",
      "Epoch[4]: v_loss: 0.34731 v_acc: 89.72\n",
      "Epoch [5/7], Step[54/54], Loss: 0.4075\n",
      "Epoch[5]: v_loss: 0.34253 v_acc: 89.83\n",
      "Epoch [6/7], Step[54/54], Loss: 0.2905\n",
      "Epoch[6]: v_loss: 0.33877 v_acc: 90.08\n",
      "Epoch [7/7], Step[54/54], Loss: 0.2909\n",
      "Epoch[7]: v_loss: 0.33598 v_acc: 90.05\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  6\n",
      "Accuracy of the network on the 10000 test images: 90.05 %\n",
      "53 iteration -  Size: 59981.367950406675 [0.1937102445, 1, 0.005387631832152908, 1, 0.45, 1]\n",
      "54 iteration -  Size: 58138.16795040668 [0.1937102445, 1, 0.005387631832152908, 1, 0.405, 1]\n",
      "55 iteration -  Size: 56315.74197015068 [0.17433922005, 1, 0.005387631832152908, 1, 0.405, 1]\n",
      "56 iteration -  Size: 54585.63115536601 [0.17433922005, 1, 0.004848868648937617, 1, 0.405, 1]\n",
      "57 iteration -  Size: 52926.751155366015 [0.17433922005, 1, 0.004848868648937617, 1, 0.36450000000000005, 1]\n",
      "58 iteration -  Size: 51286.56777313561 [0.156905298045, 1, 0.004848868648937617, 1, 0.36450000000000005, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.4360\n",
      "Epoch[1]: v_loss: 0.46494 v_acc: 85.97\n",
      "Epoch [2/7], Step[54/54], Loss: 0.4276\n",
      "Epoch[2]: v_loss: 0.42813 v_acc: 87.26\n",
      "Epoch [3/7], Step[54/54], Loss: 0.4393\n",
      "Epoch[3]: v_loss: 0.41301 v_acc: 87.71\n",
      "Epoch [4/7], Step[54/54], Loss: 0.3604\n",
      "Epoch[4]: v_loss: 0.40337 v_acc: 88.04\n",
      "Epoch [5/7], Step[54/54], Loss: 0.3447\n",
      "Epoch[5]: v_loss: 0.39602 v_acc: 88.27\n",
      "Epoch [6/7], Step[54/54], Loss: 0.4229\n",
      "Epoch[6]: v_loss: 0.39035 v_acc: 88.4\n",
      "Epoch [7/7], Step[54/54], Loss: 0.4061\n",
      "Epoch[7]: v_loss: 0.38584 v_acc: 88.52\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.52 %\n",
      "59 iteration -  Size: 49729.468039829415 [0.156905298045, 1, 0.004363981784043855, 1, 0.36450000000000005, 1]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run:\t 4ea2391b3d014e2fafff3accfb352d2c\n",
      "Depth:\t 4\n",
      "Input:\t 784\n",
      "Output:\t 10\n",
      "Leaf:\t 32\n",
      "Buffer:\t 90\n",
      "1 iteration -  Size: 1515408.8 [1, 1, 0.45, 1, 1, 1]\n",
      "2 iteration -  Size: 1370901.9200000002 [1, 1, 0.405, 1, 1, 1]\n",
      "3 iteration -  Size: 1240845.7280000001 [1, 1, 0.36450000000000005, 1, 1, 1]\n",
      "4 iteration -  Size: 1123795.1552000002 [1, 1, 0.32805000000000006, 1, 1, 1]\n",
      "5 iteration -  Size: 1018449.6396800001 [1, 1, 0.29524500000000004, 1, 1, 1]\n",
      "6 iteration -  Size: 923638.6757120001 [1, 1, 0.2657205, 1, 1, 1]\n",
      "7 iteration -  Size: 838308.8081408 [1, 1, 0.23914845, 1, 1, 1]\n",
      "8 iteration -  Size: 761511.92732672 [1, 1, 0.215233605, 1, 1, 1]\n",
      "9 iteration -  Size: 692394.734594048 [1, 1, 0.1937102445, 1, 1, 1]\n",
      "10 iteration -  Size: 630189.2611346432 [1, 1, 0.17433922005, 1, 1, 1]\n",
      "11 iteration -  Size: 574204.3350211789 [1, 1, 0.156905298045, 1, 1, 1]\n",
      "12 iteration -  Size: 523817.901519061 [1, 1, 0.1412147682405, 1, 1, 1]\n",
      "13 iteration -  Size: 478470.1113671549 [1, 1, 0.12709329141645, 1, 1, 1]\n",
      "14 iteration -  Size: 437657.1002304394 [1, 1, 0.114383962274805, 1, 1, 1]\n",
      "15 iteration -  Size: 400925.3902073955 [1, 1, 0.10294556604732451, 1, 1, 1]\n",
      "16 iteration -  Size: 367866.8511866559 [1, 1, 0.09265100944259205, 1, 1, 1]\n",
      "17 iteration -  Size: 338114.1660679903 [1, 1, 0.08338590849833284, 1, 1, 1]\n",
      "18 iteration -  Size: 311336.7494611913 [1, 1, 0.07504731764849956, 1, 1, 1]\n",
      "19 iteration -  Size: 287237.07451507216 [1, 1, 0.0675425858836496, 1, 1, 1]\n",
      "20 iteration -  Size: 265547.36706356495 [1, 1, 0.06078832729528464, 1, 1, 1]\n",
      "21 iteration -  Size: 246026.63035720843 [1, 1, 0.054709494565756175, 1, 1, 1]\n",
      "22 iteration -  Size: 228457.96732148758 [1, 1, 0.049238545109180555, 1, 1, 1]\n",
      "23 iteration -  Size: 212646.1705893388 [1, 1, 0.0443146905982625, 1, 1, 1]\n",
      "24 iteration -  Size: 198415.55353040493 [1, 1, 0.03988322153843625, 1, 1, 1]\n",
      "25 iteration -  Size: 185607.99817736444 [1, 1, 0.03589489938459262, 1, 1, 1]\n",
      "26 iteration -  Size: 174081.198359628 [1, 1, 0.03230540944613336, 1, 1, 1]\n",
      "27 iteration -  Size: 163707.0785236652 [1, 1, 0.029074868501520024, 1, 1, 1]\n",
      "28 iteration -  Size: 154370.37067129868 [1, 1, 0.026167381651368022, 1, 1, 1]\n",
      "29 iteration -  Size: 145967.3336041688 [1, 1, 0.023550643486231218, 1, 1, 1]\n",
      "30 iteration -  Size: 138404.60024375195 [1, 1, 0.021195579137608098, 1, 1, 1]\n",
      "31 iteration -  Size: 131598.14021937674 [1, 1, 0.019076021223847286, 1, 1, 1]\n",
      "32 iteration -  Size: 125472.32619743906 [1, 1, 0.017168419101462558, 1, 1, 1]\n",
      "33 iteration -  Size: 119959.09357769515 [1, 1, 0.015451577191316302, 1, 1, 1]\n",
      "34 iteration -  Size: 114997.18421992564 [1, 1, 0.013906419472184673, 1, 1, 1]\n",
      "35 iteration -  Size: 113433.18421992564 [0.45, 1, 0.013906419472184673, 1, 1, 1]\n",
      "36 iteration -  Size: 108967.46579793308 [0.45, 1, 0.012515777524966205, 1, 1, 1]\n",
      "37 iteration -  Size: 104733.86579793307 [0.405, 1, 0.012515777524966205, 1, 1, 1]\n",
      "38 iteration -  Size: 100714.71921813977 [0.405, 1, 0.011264199772469584, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 1.0144\n",
      "Epoch[1]: v_loss: 1.03347 v_acc: 81.28\n",
      "Epoch [2/7], Step[54/54], Loss: 0.9682\n",
      "Epoch[2]: v_loss: 0.95601 v_acc: 85.78\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8660\n",
      "Epoch[3]: v_loss: 0.90615 v_acc: 86.69\n",
      "Epoch [4/7], Step[54/54], Loss: 0.9224\n",
      "Epoch[4]: v_loss: 0.88301 v_acc: 87.47\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8522\n",
      "Epoch[5]: v_loss: 0.8634 v_acc: 87.98\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8613\n",
      "Epoch[6]: v_loss: 0.84833 v_acc: 88.49\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8535\n",
      "Epoch[7]: v_loss: 0.83846 v_acc: 88.72\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.72 %\n",
      "39 iteration -  Size: 96904.47921813978 [0.36450000000000005, 1, 0.011264199772469584, 1, 1, 1]\n",
      "40 iteration -  Size: 93287.24729632579 [0.36450000000000005, 1, 0.010137779795222625, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.8793\n",
      "Epoch[1]: v_loss: 0.83576 v_acc: 88.52\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8202\n",
      "Epoch[2]: v_loss: 0.82654 v_acc: 88.71\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8546\n",
      "Epoch[3]: v_loss: 0.82085 v_acc: 88.86\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8004\n",
      "Epoch[4]: v_loss: 0.81388 v_acc: 88.96\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8129\n",
      "Epoch[5]: v_loss: 0.80664 v_acc: 89.27\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8246\n",
      "Epoch[6]: v_loss: 0.80128 v_acc: 89.4\n",
      "Epoch [7/7], Step[54/54], Loss: 0.7927\n",
      "Epoch[7]: v_loss: 0.79639 v_acc: 89.57\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 89.57 %\n",
      "41 iteration -  Size: 89858.03129632579 [0.32805000000000006, 1, 0.010137779795222625, 1, 1, 1]\n",
      "42 iteration -  Size: 86602.52256669321 [0.32805000000000006, 1, 0.009124001815700363, 1, 1, 1]\n",
      "43 iteration -  Size: 83516.22816669321 [0.29524500000000004, 1, 0.009124001815700363, 1, 1, 1]\n",
      "44 iteration -  Size: 80586.27031002389 [0.29524500000000004, 1, 0.008211601634130327, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.7885\n",
      "Epoch[1]: v_loss: 0.81531 v_acc: 88.57\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8278\n",
      "Epoch[2]: v_loss: 0.81232 v_acc: 88.68\n",
      "Epoch [3/7], Step[54/54], Loss: 0.7697\n",
      "Epoch[3]: v_loss: 0.8085 v_acc: 88.78\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8164\n",
      "Epoch[4]: v_loss: 0.80635 v_acc: 88.88\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8040\n",
      "Epoch[5]: v_loss: 0.80413 v_acc: 88.84\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8118\n",
      "Epoch[6]: v_loss: 0.80257 v_acc: 88.94\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8073\n",
      "Epoch[7]: v_loss: 0.80071 v_acc: 88.92\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  6\n",
      "Accuracy of the network on the 10000 test images: 88.92 %\n",
      "45 iteration -  Size: 77808.60535002389 [0.2657205, 1, 0.008211601634130327, 1, 1, 1]\n",
      "46 iteration -  Size: 75171.6432790215 [0.2657205, 1, 0.007390441470717294, 1, 1, 1]\n",
      "47 iteration -  Size: 72671.7448150215 [0.23914845, 1, 0.007390441470717294, 1, 1, 1]\n",
      "48 iteration -  Size: 70298.47895111935 [0.23914845, 1, 0.006651397323645565, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.8076\n",
      "Epoch[1]: v_loss: 0.80051 v_acc: 88.63\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8113\n",
      "Epoch[2]: v_loss: 0.79899 v_acc: 88.66\n",
      "Epoch [3/7], Step[54/54], Loss: 0.7835\n",
      "Epoch[3]: v_loss: 0.79754 v_acc: 88.76\n",
      "Epoch [4/7], Step[54/54], Loss: 0.7607\n",
      "Epoch[4]: v_loss: 0.79633 v_acc: 88.76\n",
      "Epoch [5/7], Step[54/54], Loss: 0.7788\n",
      "Epoch[5]: v_loss: 0.79554 v_acc: 88.81\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8180\n",
      "Epoch[6]: v_loss: 0.79379 v_acc: 88.9\n",
      "Epoch [7/7], Step[54/54], Loss: 0.7839\n",
      "Epoch[7]: v_loss: 0.79219 v_acc: 88.97\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.97 %\n",
      "49 iteration -  Size: 68048.57033351935 [0.215233605, 1, 0.006651397323645565, 1, 1, 1]\n",
      "50 iteration -  Size: 65912.63105600742 [0.215233605, 1, 0.005986257591281009, 1, 1, 1]\n",
      "51 iteration -  Size: 63928.63105600742 [0.215233605, 1, 0.005986257591281009, 1, 0.45, 1]\n",
      "52 iteration -  Size: 61903.713300167416 [0.1937102445, 1, 0.005986257591281009, 1, 0.45, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/7], Step[54/54], Loss: 0.7905\n",
      "Epoch[1]: v_loss: 0.79924 v_acc: 88.44\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8346\n",
      "Epoch[2]: v_loss: 0.79902 v_acc: 88.44\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8017\n",
      "Epoch[3]: v_loss: 0.7976 v_acc: 88.51\n",
      "Epoch [4/7], Step[54/54], Loss: 0.7902\n",
      "Epoch[4]: v_loss: 0.79778 v_acc: 88.52\n",
      "Epoch [5/7], Step[54/54], Loss: 0.7801\n",
      "Epoch[5]: v_loss: 0.79663 v_acc: 88.48\n",
      "Epoch [6/7], Step[54/54], Loss: 0.7725\n",
      "Epoch[6]: v_loss: 0.79588 v_acc: 88.57\n",
      "Epoch [7/7], Step[54/54], Loss: 0.7838\n",
      "Epoch[7]: v_loss: 0.79486 v_acc: 88.61\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.61 %\n",
      "53 iteration -  Size: 59981.367950406675 [0.1937102445, 1, 0.005387631832152908, 1, 0.45, 1]\n",
      "54 iteration -  Size: 58138.16795040668 [0.1937102445, 1, 0.005387631832152908, 1, 0.405, 1]\n",
      "55 iteration -  Size: 56315.74197015068 [0.17433922005, 1, 0.005387631832152908, 1, 0.405, 1]\n",
      "56 iteration -  Size: 54585.63115536601 [0.17433922005, 1, 0.004848868648937617, 1, 0.405, 1]\n",
      "57 iteration -  Size: 52926.751155366015 [0.17433922005, 1, 0.004848868648937617, 1, 0.36450000000000005, 1]\n",
      "58 iteration -  Size: 51286.56777313561 [0.156905298045, 1, 0.004848868648937617, 1, 0.36450000000000005, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([15, 784])\n",
      "Activating: torch.Size([15, 1])\n",
      "Activating: torch.Size([16, 784, 32])\n",
      "Activating: torch.Size([16, 32])\n",
      "Activating: torch.Size([16, 32, 10])\n",
      "Activating: torch.Size([16, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.7568\n",
      "Epoch[1]: v_loss: 0.77504 v_acc: 89.2\n",
      "Epoch [2/7], Step[54/54], Loss: 0.7706\n",
      "Epoch[2]: v_loss: 0.77395 v_acc: 89.28\n",
      "Epoch [3/7], Step[54/54], Loss: 0.7409\n",
      "Epoch[3]: v_loss: 0.77337 v_acc: 89.26\n",
      "Epoch [4/7], Step[54/54], Loss: 0.7435\n",
      "Epoch[4]: v_loss: 0.77315 v_acc: 89.37\n",
      "Epoch [5/7], Step[54/54], Loss: 0.7872\n",
      "Epoch[5]: v_loss: 0.77255 v_acc: 89.26\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8192\n",
      "Epoch[6]: v_loss: 0.77225 v_acc: 89.31\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8029\n",
      "Epoch[7]: v_loss: 0.77166 v_acc: 89.31\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  4\n",
      "Accuracy of the network on the 10000 test images: 89.31 %\n",
      "59 iteration -  Size: 49729.468039829415 [0.156905298045, 1, 0.004363981784043855, 1, 0.36450000000000005, 1]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run:\t 612d1573468c4ffaa89bf54d60ce4508\n",
      "Depth:\t 3\n",
      "Input:\t 784\n",
      "Output:\t 10\n",
      "Leaf:\t 32\n",
      "Buffer:\t 90\n",
      "1 iteration -  Size: 756138.4 [1, 1, 0.45, 1, 1, 1]\n",
      "2 iteration -  Size: 683884.9600000001 [1, 1, 0.405, 1, 1, 1]\n",
      "3 iteration -  Size: 618856.8640000001 [1, 1, 0.36450000000000005, 1, 1, 1]\n",
      "4 iteration -  Size: 560331.5776000001 [1, 1, 0.32805000000000006, 1, 1, 1]\n",
      "5 iteration -  Size: 507658.81984000007 [1, 1, 0.29524500000000004, 1, 1, 1]\n",
      "6 iteration -  Size: 460253.33785600006 [1, 1, 0.2657205, 1, 1, 1]\n",
      "7 iteration -  Size: 417588.4040704 [1, 1, 0.23914845, 1, 1, 1]\n",
      "8 iteration -  Size: 379189.96366336 [1, 1, 0.215233605, 1, 1, 1]\n",
      "9 iteration -  Size: 344631.367297024 [1, 1, 0.1937102445, 1, 1, 1]\n",
      "10 iteration -  Size: 313528.6305673216 [1, 1, 0.17433922005, 1, 1, 1]\n",
      "11 iteration -  Size: 285536.16751058947 [1, 1, 0.156905298045, 1, 1, 1]\n",
      "12 iteration -  Size: 260342.9507595305 [1, 1, 0.1412147682405, 1, 1, 1]\n",
      "13 iteration -  Size: 237669.05568357746 [1, 1, 0.12709329141645, 1, 1, 1]\n",
      "14 iteration -  Size: 217262.5501152197 [1, 1, 0.114383962274805, 1, 1, 1]\n",
      "15 iteration -  Size: 198896.69510369774 [1, 1, 0.10294556604732451, 1, 1, 1]\n",
      "16 iteration -  Size: 182367.42559332796 [1, 1, 0.09265100944259205, 1, 1, 1]\n",
      "17 iteration -  Size: 167491.08303399515 [1, 1, 0.08338590849833284, 1, 1, 1]\n",
      "18 iteration -  Size: 154102.37473059565 [1, 1, 0.07504731764849956, 1, 1, 1]\n",
      "19 iteration -  Size: 142052.53725753608 [1, 1, 0.0675425858836496, 1, 1, 1]\n",
      "20 iteration -  Size: 131207.68353178247 [1, 1, 0.06078832729528464, 1, 1, 1]\n",
      "21 iteration -  Size: 121447.31517860421 [1, 1, 0.054709494565756175, 1, 1, 1]\n",
      "22 iteration -  Size: 112662.98366074379 [1, 1, 0.049238545109180555, 1, 1, 1]\n",
      "23 iteration -  Size: 104757.0852946694 [1, 1, 0.0443146905982625, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.3105\n",
      "Epoch[1]: v_loss: 0.28455 v_acc: 92.26\n",
      "Epoch [2/7], Step[54/54], Loss: 0.2219\n",
      "Epoch[2]: v_loss: 0.25702 v_acc: 92.9\n",
      "Epoch [3/7], Step[54/54], Loss: 0.2315\n",
      "Epoch[3]: v_loss: 0.24271 v_acc: 93.24\n",
      "Epoch [4/7], Step[54/54], Loss: 0.2393\n",
      "Epoch[4]: v_loss: 0.23258 v_acc: 93.45\n",
      "Epoch [5/7], Step[54/54], Loss: 0.2098\n",
      "Epoch[5]: v_loss: 0.22616 v_acc: 93.66\n",
      "Epoch [6/7], Step[54/54], Loss: 0.2323\n",
      "Epoch[6]: v_loss: 0.22112 v_acc: 93.8\n",
      "Epoch [7/7], Step[54/54], Loss: 0.1956\n",
      "Epoch[7]: v_loss: 0.21541 v_acc: 93.91\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 93.91 %\n",
      "24 iteration -  Size: 97641.77676520246 [1, 1, 0.03988322153843625, 1, 1, 1]\n",
      "25 iteration -  Size: 91237.99908868222 [1, 1, 0.03589489938459262, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.1910\n",
      "Epoch[1]: v_loss: 0.21484 v_acc: 93.96\n",
      "Epoch [2/7], Step[54/54], Loss: 0.1990\n",
      "Epoch[2]: v_loss: 0.21059 v_acc: 94.03\n",
      "Epoch [3/7], Step[54/54], Loss: 0.1677\n",
      "Epoch[3]: v_loss: 0.20678 v_acc: 94.13\n",
      "Epoch [4/7], Step[54/54], Loss: 0.1720\n",
      "Epoch[4]: v_loss: 0.20452 v_acc: 94.22\n",
      "Epoch [5/7], Step[54/54], Loss: 0.1901\n",
      "Epoch[5]: v_loss: 0.20209 v_acc: 94.22\n",
      "Epoch [6/7], Step[54/54], Loss: 0.1669\n",
      "Epoch[6]: v_loss: 0.20059 v_acc: 94.3\n",
      "Epoch [7/7], Step[54/54], Loss: 0.1738\n",
      "Epoch[7]: v_loss: 0.19906 v_acc: 94.34\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 94.34 %\n",
      "26 iteration -  Size: 85474.599179814 [1, 1, 0.03230540944613336, 1, 1, 1]\n",
      "27 iteration -  Size: 80287.5392618326 [1, 1, 0.029074868501520024, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.1903\n",
      "Epoch[1]: v_loss: 0.20105 v_acc: 94.19\n",
      "Epoch [2/7], Step[54/54], Loss: 0.1553\n",
      "Epoch[2]: v_loss: 0.19644 v_acc: 94.22\n",
      "Epoch [3/7], Step[54/54], Loss: 0.1653\n",
      "Epoch[3]: v_loss: 0.19553 v_acc: 94.3\n",
      "Epoch [4/7], Step[54/54], Loss: 0.1394\n",
      "Epoch[4]: v_loss: 0.19358 v_acc: 94.36\n",
      "Epoch [5/7], Step[54/54], Loss: 0.1380\n",
      "Epoch[5]: v_loss: 0.19327 v_acc: 94.35\n",
      "Epoch [6/7], Step[54/54], Loss: 0.1583\n",
      "Epoch[6]: v_loss: 0.19292 v_acc: 94.33\n",
      "Epoch [7/7], Step[54/54], Loss: 0.1317\n",
      "Epoch[7]: v_loss: 0.19212 v_acc: 94.45\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 94.45 %\n",
      "28 iteration -  Size: 75619.18533564934 [1, 1, 0.026167381651368022, 1, 1, 1]\n",
      "29 iteration -  Size: 71417.6668020844 [1, 1, 0.023550643486231218, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.1499\n",
      "Epoch[1]: v_loss: 0.19854 v_acc: 94.2\n",
      "Epoch [2/7], Step[54/54], Loss: 0.1847\n",
      "Epoch[2]: v_loss: 0.19505 v_acc: 94.39\n",
      "Epoch [3/7], Step[54/54], Loss: 0.1331\n",
      "Epoch[3]: v_loss: 0.19432 v_acc: 94.54\n",
      "Epoch [4/7], Step[54/54], Loss: 0.1676\n",
      "Epoch[4]: v_loss: 0.19238 v_acc: 94.52\n",
      "Epoch [5/7], Step[54/54], Loss: 0.1485\n",
      "Epoch[5]: v_loss: 0.19204 v_acc: 94.53\n",
      "Epoch [6/7], Step[54/54], Loss: 0.1414\n",
      "Epoch[6]: v_loss: 0.19271 v_acc: 94.51\n",
      "Epoch [7/7], Step[54/54], Loss: 0.1326\n",
      "Epoch[7]: v_loss: 0.19203 v_acc: 94.54\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 94.54 %\n",
      "30 iteration -  Size: 67636.30012187597 [1, 1, 0.021195579137608098, 1, 1, 1]\n",
      "31 iteration -  Size: 64233.07010968837 [1, 1, 0.019076021223847286, 1, 1, 1]\n",
      "32 iteration -  Size: 61170.16309871953 [1, 1, 0.017168419101462558, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.2016\n",
      "Epoch[1]: v_loss: 0.23059 v_acc: 93.05\n",
      "Epoch [2/7], Step[54/54], Loss: 0.1817\n",
      "Epoch[2]: v_loss: 0.21803 v_acc: 93.51\n",
      "Epoch [3/7], Step[54/54], Loss: 0.1562\n",
      "Epoch[3]: v_loss: 0.21205 v_acc: 93.5\n",
      "Epoch [4/7], Step[54/54], Loss: 0.1957\n",
      "Epoch[4]: v_loss: 0.20811 v_acc: 93.77\n",
      "Epoch [5/7], Step[54/54], Loss: 0.1463\n",
      "Epoch[5]: v_loss: 0.20688 v_acc: 93.73\n",
      "Epoch [6/7], Step[54/54], Loss: 0.1765\n",
      "Epoch[6]: v_loss: 0.20535 v_acc: 93.8\n",
      "Epoch [7/7], Step[54/54], Loss: 0.1437\n",
      "Epoch[7]: v_loss: 0.20391 v_acc: 93.92\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 93.92 %\n",
      "33 iteration -  Size: 58413.546788847576 [1, 1, 0.015451577191316302, 1, 1, 1]\n",
      "34 iteration -  Size: 55932.59210996282 [1, 1, 0.013906419472184673, 1, 1, 1]\n",
      "35 iteration -  Size: 53699.73289896654 [1, 1, 0.012515777524966205, 1, 1, 1]\n",
      "36 iteration -  Size: 54644.532898966536 [0.45, 1, 0.012515777524966205, 1, 1, 1]\n",
      "37 iteration -  Size: 52634.95960906988 [0.45, 1, 0.011264199772469584, 1, 1, 1]\n",
      "38 iteration -  Size: 50659.27960906988 [0.405, 1, 0.011264199772469584, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/7], Step[54/54], Loss: 0.2853\n",
      "Epoch[1]: v_loss: 0.29964 v_acc: 91.34\n",
      "Epoch [2/7], Step[54/54], Loss: 0.3000\n",
      "Epoch[2]: v_loss: 0.28139 v_acc: 92.01\n",
      "Epoch [3/7], Step[54/54], Loss: 0.2119\n",
      "Epoch[3]: v_loss: 0.26964 v_acc: 92.36\n",
      "Epoch [4/7], Step[54/54], Loss: 0.2580\n",
      "Epoch[4]: v_loss: 0.26152 v_acc: 92.61\n",
      "Epoch [5/7], Step[54/54], Loss: 0.2733\n",
      "Epoch[5]: v_loss: 0.25582 v_acc: 92.73\n",
      "Epoch [6/7], Step[54/54], Loss: 0.2480\n",
      "Epoch[6]: v_loss: 0.25107 v_acc: 92.88\n",
      "Epoch [7/7], Step[54/54], Loss: 0.2622\n",
      "Epoch[7]: v_loss: 0.24909 v_acc: 92.88\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 92.88 %\n",
      "39 iteration -  Size: 48850.663648162896 [0.405, 1, 0.010137779795222625, 1, 1, 1]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run:\t 24879e69bd174ef6a9973359ea1b9b6c\n",
      "Depth:\t 3\n",
      "Input:\t 784\n",
      "Output:\t 10\n",
      "Leaf:\t 32\n",
      "Buffer:\t 90\n",
      "1 iteration -  Size: 756138.4 [1, 1, 0.45, 1, 1, 1]\n",
      "2 iteration -  Size: 683884.9600000001 [1, 1, 0.405, 1, 1, 1]\n",
      "3 iteration -  Size: 618856.8640000001 [1, 1, 0.36450000000000005, 1, 1, 1]\n",
      "4 iteration -  Size: 560331.5776000001 [1, 1, 0.32805000000000006, 1, 1, 1]\n",
      "5 iteration -  Size: 507658.81984000007 [1, 1, 0.29524500000000004, 1, 1, 1]\n",
      "6 iteration -  Size: 460253.33785600006 [1, 1, 0.2657205, 1, 1, 1]\n",
      "7 iteration -  Size: 417588.4040704 [1, 1, 0.23914845, 1, 1, 1]\n",
      "8 iteration -  Size: 379189.96366336 [1, 1, 0.215233605, 1, 1, 1]\n",
      "9 iteration -  Size: 344631.367297024 [1, 1, 0.1937102445, 1, 1, 1]\n",
      "10 iteration -  Size: 313528.6305673216 [1, 1, 0.17433922005, 1, 1, 1]\n",
      "11 iteration -  Size: 285536.16751058947 [1, 1, 0.156905298045, 1, 1, 1]\n",
      "12 iteration -  Size: 260342.9507595305 [1, 1, 0.1412147682405, 1, 1, 1]\n",
      "13 iteration -  Size: 237669.05568357746 [1, 1, 0.12709329141645, 1, 1, 1]\n",
      "14 iteration -  Size: 217262.5501152197 [1, 1, 0.114383962274805, 1, 1, 1]\n",
      "15 iteration -  Size: 198896.69510369774 [1, 1, 0.10294556604732451, 1, 1, 1]\n",
      "16 iteration -  Size: 182367.42559332796 [1, 1, 0.09265100944259205, 1, 1, 1]\n",
      "17 iteration -  Size: 167491.08303399515 [1, 1, 0.08338590849833284, 1, 1, 1]\n",
      "18 iteration -  Size: 154102.37473059565 [1, 1, 0.07504731764849956, 1, 1, 1]\n",
      "19 iteration -  Size: 142052.53725753608 [1, 1, 0.0675425858836496, 1, 1, 1]\n",
      "20 iteration -  Size: 131207.68353178247 [1, 1, 0.06078832729528464, 1, 1, 1]\n",
      "21 iteration -  Size: 121447.31517860421 [1, 1, 0.054709494565756175, 1, 1, 1]\n",
      "22 iteration -  Size: 112662.98366074379 [1, 1, 0.049238545109180555, 1, 1, 1]\n",
      "23 iteration -  Size: 104757.0852946694 [1, 1, 0.0443146905982625, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 1.1652\n",
      "Epoch[1]: v_loss: 1.1739 v_acc: 77.39\n",
      "Epoch [2/7], Step[54/54], Loss: 1.0776\n",
      "Epoch[2]: v_loss: 1.07727 v_acc: 81.93\n",
      "Epoch [3/7], Step[54/54], Loss: 1.0187\n",
      "Epoch[3]: v_loss: 0.99246 v_acc: 85.37\n",
      "Epoch [4/7], Step[54/54], Loss: 0.9446\n",
      "Epoch[4]: v_loss: 0.95495 v_acc: 86.46\n",
      "Epoch [5/7], Step[54/54], Loss: 0.9409\n",
      "Epoch[5]: v_loss: 0.93399 v_acc: 86.9\n",
      "Epoch [6/7], Step[54/54], Loss: 0.9092\n",
      "Epoch[6]: v_loss: 0.92141 v_acc: 87.15\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8989\n",
      "Epoch[7]: v_loss: 0.91086 v_acc: 87.28\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 87.28 %\n",
      "24 iteration -  Size: 97641.77676520246 [1, 1, 0.03988322153843625, 1, 1, 1]\n",
      "25 iteration -  Size: 91237.99908868222 [1, 1, 0.03589489938459262, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.8937\n",
      "Epoch[1]: v_loss: 0.90257 v_acc: 87.46\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8849\n",
      "Epoch[2]: v_loss: 0.8953 v_acc: 87.55\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8881\n",
      "Epoch[3]: v_loss: 0.88937 v_acc: 87.76\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8858\n",
      "Epoch[4]: v_loss: 0.88247 v_acc: 87.88\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8514\n",
      "Epoch[5]: v_loss: 0.87816 v_acc: 88.1\n",
      "Epoch [6/7], Step[54/54], Loss: 0.7956\n",
      "Epoch[6]: v_loss: 0.87267 v_acc: 88.18\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8799\n",
      "Epoch[7]: v_loss: 0.86888 v_acc: 88.26\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.26 %\n",
      "26 iteration -  Size: 85474.599179814 [1, 1, 0.03230540944613336, 1, 1, 1]\n",
      "27 iteration -  Size: 80287.5392618326 [1, 1, 0.029074868501520024, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.8360\n",
      "Epoch[1]: v_loss: 0.86572 v_acc: 88.33\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8122\n",
      "Epoch[2]: v_loss: 0.86247 v_acc: 88.43\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8336\n",
      "Epoch[3]: v_loss: 0.86039 v_acc: 88.55\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8037\n",
      "Epoch[4]: v_loss: 0.85846 v_acc: 88.57\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8698\n",
      "Epoch[5]: v_loss: 0.85615 v_acc: 88.61\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8156\n",
      "Epoch[6]: v_loss: 0.854 v_acc: 88.67\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8320\n",
      "Epoch[7]: v_loss: 0.85304 v_acc: 88.68\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 88.68 %\n",
      "28 iteration -  Size: 75619.18533564934 [1, 1, 0.026167381651368022, 1, 1, 1]\n",
      "29 iteration -  Size: 71417.6668020844 [1, 1, 0.023550643486231218, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.7866\n",
      "Epoch[1]: v_loss: 0.85171 v_acc: 88.72\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8170\n",
      "Epoch[2]: v_loss: 0.85073 v_acc: 88.74\n",
      "Epoch [3/7], Step[54/54], Loss: 0.7474\n",
      "Epoch[3]: v_loss: 0.84909 v_acc: 88.83\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8767\n",
      "Epoch[4]: v_loss: 0.84859 v_acc: 88.85\n",
      "Epoch [5/7], Step[54/54], Loss: 0.7606\n",
      "Epoch[5]: v_loss: 0.8477 v_acc: 88.88\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8204\n",
      "Epoch[6]: v_loss: 0.84714 v_acc: 88.85\n",
      "Epoch [7/7], Step[54/54], Loss: 0.7726\n",
      "Epoch[7]: v_loss: 0.84579 v_acc: 89.0\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 89.0 %\n",
      "30 iteration -  Size: 67636.30012187597 [1, 1, 0.021195579137608098, 1, 1, 1]\n",
      "31 iteration -  Size: 64233.07010968837 [1, 1, 0.019076021223847286, 1, 1, 1]\n",
      "32 iteration -  Size: 61170.16309871953 [1, 1, 0.017168419101462558, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n",
      "Epoch [1/7], Step[54/54], Loss: 0.7368\n",
      "Epoch[1]: v_loss: 0.8448 v_acc: 89.01\n",
      "Epoch [2/7], Step[54/54], Loss: 0.7772\n",
      "Epoch[2]: v_loss: 0.84451 v_acc: 89.02\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8079\n",
      "Epoch[3]: v_loss: 0.84466 v_acc: 89.1\n",
      "Epoch [4/7], Step[54/54], Loss: 0.8551\n",
      "Epoch[4]: v_loss: 0.84401 v_acc: 89.08\n",
      "Epoch [5/7], Step[54/54], Loss: 0.7976\n",
      "Epoch[5]: v_loss: 0.84433 v_acc: 89.09\n",
      "Epoch [6/7], Step[54/54], Loss: 0.7626\n",
      "Epoch[6]: v_loss: 0.84348 v_acc: 89.22\n",
      "Epoch [7/7], Step[54/54], Loss: 0.7972\n",
      "Epoch[7]: v_loss: 0.84362 v_acc: 89.26\n",
      "Best model saved at epoch:  6\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 89.22 %\n",
      "33 iteration -  Size: 58413.546788847576 [1, 1, 0.015451577191316302, 1, 1, 1]\n",
      "34 iteration -  Size: 55932.59210996282 [1, 1, 0.013906419472184673, 1, 1, 1]\n",
      "35 iteration -  Size: 53699.73289896654 [1, 1, 0.012515777524966205, 1, 1, 1]\n",
      "36 iteration -  Size: 54644.532898966536 [0.45, 1, 0.012515777524966205, 1, 1, 1]\n",
      "37 iteration -  Size: 52634.95960906988 [0.45, 1, 0.011264199772469584, 1, 1, 1]\n",
      "38 iteration -  Size: 50659.27960906988 [0.405, 1, 0.011264199772469584, 1, 1, 1]\n",
      "Disabling: _fff.fff.depth\n",
      "Disabling: _fff.fff.node_weights\n",
      "Disabling: _fff.fff.node_biases\n",
      "Disabling: _fff.fff.w1s\n",
      "Disabling: _fff.fff.b1s\n",
      "Disabling: _fff.fff.w2s\n",
      "Disabling: _fff.fff.b2s\n",
      "Activating: torch.Size([7, 784])\n",
      "Activating: torch.Size([7, 1])\n",
      "Activating: torch.Size([8, 784, 32])\n",
      "Activating: torch.Size([8, 32])\n",
      "Activating: torch.Size([8, 32, 10])\n",
      "Activating: torch.Size([8, 10])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/7], Step[54/54], Loss: 0.9144\n",
      "Epoch[1]: v_loss: 0.92371 v_acc: 85.76\n",
      "Epoch [2/7], Step[54/54], Loss: 0.8656\n",
      "Epoch[2]: v_loss: 0.92091 v_acc: 85.94\n",
      "Epoch [3/7], Step[54/54], Loss: 0.8467\n",
      "Epoch[3]: v_loss: 0.91947 v_acc: 85.93\n",
      "Epoch [4/7], Step[54/54], Loss: 0.9735\n",
      "Epoch[4]: v_loss: 0.91822 v_acc: 85.91\n",
      "Epoch [5/7], Step[54/54], Loss: 0.8814\n",
      "Epoch[5]: v_loss: 0.91731 v_acc: 85.98\n",
      "Epoch [6/7], Step[54/54], Loss: 0.8783\n",
      "Epoch[6]: v_loss: 0.91669 v_acc: 85.99\n",
      "Epoch [7/7], Step[54/54], Loss: 0.8753\n",
      "Epoch[7]: v_loss: 0.91665 v_acc: 86.01\n",
      "Best model saved at epoch:  7\n",
      "Best acc model saved at epoch:  7\n",
      "Accuracy of the network on the 10000 test images: 86.01 %\n",
      "39 iteration -  Size: 48850.663648162896 [0.405, 1, 0.010137779795222625, 1, 1, 1]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Run:\t 7c66090a514e4080b639b6f261d5a134\n",
      "Depth:\t 2\n",
      "Input:\t 784\n",
      "Output:\t 10\n",
      "Leaf:\t 32\n",
      "Buffer:\t 90\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 47\u001b[0m\n\u001b[0;32m     45\u001b[0m model\u001b[38;5;241m.\u001b[39meval()\n\u001b[0;32m     46\u001b[0m t \u001b[38;5;241m=\u001b[39m time()\n\u001b[1;32m---> 47\u001b[0m train_loss, train_acc \u001b[38;5;241m=\u001b[39m \u001b[43mtest\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     48\u001b[0m test_loss, test_acc \u001b[38;5;241m=\u001b[39m test(model, test_loader)\n\u001b[0;32m     49\u001b[0m t \u001b[38;5;241m=\u001b[39m time() \u001b[38;5;241m-\u001b[39m t\n",
      "Cell \u001b[1;32mIn[8], line 31\u001b[0m, in \u001b[0;36mtest\u001b[1;34m(net, testloader)\u001b[0m\n\u001b[0;32m     28\u001b[0m \u001b[38;5;66;03m# Train the network for the given number of epochs\u001b[39;00m\n\u001b[0;32m     29\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m     30\u001b[0m     \u001b[38;5;66;03m# Iterate over data\u001b[39;00m\n\u001b[1;32m---> 31\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m data \u001b[38;5;129;01min\u001b[39;00m testloader:\n\u001b[0;32m     32\u001b[0m         images, labels \u001b[38;5;241m=\u001b[39m data[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mto(DEVICE), data[\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39mto(DEVICE)\n\u001b[0;32m     33\u001b[0m         outputs \u001b[38;5;241m=\u001b[39m net(images)\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:441\u001b[0m, in \u001b[0;36mDataLoader.__iter__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    439\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_iterator\n\u001b[0;32m    440\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 441\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_get_iterator\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:388\u001b[0m, in \u001b[0;36mDataLoader._get_iterator\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    386\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    387\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcheck_worker_number_rationality()\n\u001b[1;32m--> 388\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_MultiProcessingDataLoaderIter\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:1042\u001b[0m, in \u001b[0;36m_MultiProcessingDataLoaderIter.__init__\u001b[1;34m(self, loader)\u001b[0m\n\u001b[0;32m   1035\u001b[0m w\u001b[38;5;241m.\u001b[39mdaemon \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m   1036\u001b[0m \u001b[38;5;66;03m# NB: Process.start() actually take some time as it needs to\u001b[39;00m\n\u001b[0;32m   1037\u001b[0m \u001b[38;5;66;03m#     start a process and pass the arguments over via a pipe.\u001b[39;00m\n\u001b[0;32m   1038\u001b[0m \u001b[38;5;66;03m#     Therefore, we only add a worker to self._workers list after\u001b[39;00m\n\u001b[0;32m   1039\u001b[0m \u001b[38;5;66;03m#     it started, so that we do not call .join() if program dies\u001b[39;00m\n\u001b[0;32m   1040\u001b[0m \u001b[38;5;66;03m#     before it starts, and __del__ tries to join but will get:\u001b[39;00m\n\u001b[0;32m   1041\u001b[0m \u001b[38;5;66;03m#     AssertionError: can only join a started process.\u001b[39;00m\n\u001b[1;32m-> 1042\u001b[0m \u001b[43mw\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1043\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_index_queues\u001b[38;5;241m.\u001b[39mappend(index_queue)\n\u001b[0;32m   1044\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_workers\u001b[38;5;241m.\u001b[39mappend(w)\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\multiprocessing\\process.py:121\u001b[0m, in \u001b[0;36mBaseProcess.start\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    118\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _current_process\u001b[38;5;241m.\u001b[39m_config\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdaemon\u001b[39m\u001b[38;5;124m'\u001b[39m), \\\n\u001b[0;32m    119\u001b[0m        \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdaemonic processes are not allowed to have children\u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m    120\u001b[0m _cleanup()\n\u001b[1;32m--> 121\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_Popen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[0;32m    122\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sentinel \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_popen\u001b[38;5;241m.\u001b[39msentinel\n\u001b[0;32m    123\u001b[0m \u001b[38;5;66;03m# Avoid a refcycle if the target function holds an indirect\u001b[39;00m\n\u001b[0;32m    124\u001b[0m \u001b[38;5;66;03m# reference to the process object (see bpo-30775)\u001b[39;00m\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\multiprocessing\\context.py:224\u001b[0m, in \u001b[0;36mProcess._Popen\u001b[1;34m(process_obj)\u001b[0m\n\u001b[0;32m    222\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[0;32m    223\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_Popen\u001b[39m(process_obj):\n\u001b[1;32m--> 224\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_default_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_context\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mProcess\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_Popen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\multiprocessing\\context.py:336\u001b[0m, in \u001b[0;36mSpawnProcess._Popen\u001b[1;34m(process_obj)\u001b[0m\n\u001b[0;32m    333\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[0;32m    334\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_Popen\u001b[39m(process_obj):\n\u001b[0;32m    335\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mpopen_spawn_win32\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Popen\n\u001b[1;32m--> 336\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mPopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\multiprocessing\\popen_spawn_win32.py:93\u001b[0m, in \u001b[0;36mPopen.__init__\u001b[1;34m(self, process_obj)\u001b[0m\n\u001b[0;32m     91\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m     92\u001b[0m     reduction\u001b[38;5;241m.\u001b[39mdump(prep_data, to_child)\n\u001b[1;32m---> 93\u001b[0m     \u001b[43mreduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdump\u001b[49m\u001b[43m(\u001b[49m\u001b[43mprocess_obj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mto_child\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     94\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[0;32m     95\u001b[0m     set_spawning_popen(\u001b[38;5;28;01mNone\u001b[39;00m)\n",
      "File \u001b[1;32m~\\anaconda3\\envs\\unitn\\lib\\multiprocessing\\reduction.py:60\u001b[0m, in \u001b[0;36mdump\u001b[1;34m(obj, file, protocol)\u001b[0m\n\u001b[0;32m     58\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdump\u001b[39m(obj, file, protocol\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m     59\u001b[0m     \u001b[38;5;124;03m'''Replacement for pickle.dump() using ForkingPickler.'''\u001b[39;00m\n\u001b[1;32m---> 60\u001b[0m     \u001b[43mForkingPickler\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfile\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprotocol\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdump\u001b[49m\u001b[43m(\u001b[49m\u001b[43mobj\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "result = []\n",
    "for i in range (0,len(list_of_run)):\n",
    "    run_id = list_of_run[i]\n",
    "    mlflow.artifacts.download_artifacts(run_id=run_id, dst_path=\".\")\n",
    "    wrapped_model = pickle.load(open(\"./truncated_model.pkl\", \"rb\"))\n",
    "    depth = wrapped_model._fff.fff.depth.item()\n",
    "    input_width = wrapped_model._fff.fff.input_width\n",
    "    leaf_width = wrapped_model._fff.fff.leaf_width\n",
    "    output_width = wrapped_model._fff.fff.output_width\n",
    "    buffer_size = 2*(leaf_width + output_width + 3)\n",
    "    print(\"Run:\\t\", run_id)\n",
    "    print(\"Depth:\\t\", depth)\n",
    "    print(\"Input:\\t\", input_width)\n",
    "    print(\"Output:\\t\", output_width)\n",
    "    print(\"Leaf:\\t\", leaf_width)\n",
    "    print(\"Buffer:\\t\", buffer_size)\n",
    "    \n",
    "    # to reduce the sparsity and train only below a certain tresholds\n",
    "    list_of_sizes = [100, 90, 80, 70, 60, 50]\n",
    "    checked_sizes = [False for x in list_of_sizes]\n",
    "    current_size_index = 0\n",
    "    \n",
    "    start = 0.5\n",
    "    a = start\n",
    "    b = start\n",
    "    sizes=[]\n",
    "    before_trunc_sizes=[]\n",
    "    trunc_sizes=[]\n",
    "\n",
    "    model = wrapped_model.to(device)\n",
    "    \n",
    "    layers_list = []\n",
    "    for i, (name, p) in enumerate(model.named_parameters()):\n",
    "        if (len(list(p.shape)) > 1 and p.requires_grad):\n",
    "            layers_list.append(p)\n",
    "    un, comp, layers = print_size_fc(model, layers_list, [1,1,1,1,1,1])\n",
    "    \n",
    "    MODEL_NAME_COMPRESSED = \"mnist_\" + run_id + \"_compressed_full\"\n",
    "    \n",
    "    layers_sizes = layers.copy()\n",
    "    # Creating a list of 1s with the same length as original_list using multiplication\n",
    "    list_of_sparsity = [1] * len(layers_sizes)\n",
    "    \n",
    "    # Testing time and accuracy\n",
    "    model.eval()\n",
    "    t = time()\n",
    "    train_loss, train_acc = test(model, train_loader)\n",
    "    test_loss, test_acc = test(model, test_loader)\n",
    "    t = time() - t\n",
    "    un, comp, layers_sizes = print_size_fc(model, layers_list, list_of_sparsity)\n",
    "    sizes.append((comp/1000, test_acc, test_acc, t))\n",
    "    \n",
    "    for i in range(1, 100):\n",
    "        # get index of max\n",
    "        index_of_max = np.argmax(layers_sizes)\n",
    "        current_sparsity = list_of_sparsity[index_of_max]\n",
    "\n",
    "        # reduce\n",
    "        if (current_sparsity == 1):\n",
    "            list_of_sparsity[index_of_max] = start - (start * 0.1)\n",
    "        else:\n",
    "            list_of_sparsity[index_of_max] = current_sparsity - (current_sparsity * 0.1)\n",
    "\n",
    "        un, comp, layers_sizes = print_size_fc(model, layers_list, list_of_sparsity)\n",
    "        \n",
    "        if (current_size_index >= len(list_of_sizes)):\n",
    "            break\n",
    "        elif (comp / 1000 < list_of_sizes[current_size_index]):\n",
    "            # compress and save result\n",
    "            MODEL_NAME_COMPRESSED = \"mnist_\" + run_id + \"_compressed_\" + str(list_of_sizes[current_size_index])\n",
    "            current_size_index += 1\n",
    "            model.train()\n",
    "            accuracy = perform_compression(model, layers_list, list_of_sparsity, learning_rate, num_epochs,\n",
    "                                           train_loader, test_loader, device,\n",
    "                                           val_loader=val_loader, model_name=MODEL_NAME_COMPRESSED, given_criterion=criterion)\n",
    "            model.load_state_dict(torch.load(MODEL_NAME_COMPRESSED+\".h5\", map_location='cpu'))\n",
    "            model.eval()\n",
    "            t = time()\n",
    "            train_loss, train_acc = test(model, train_loader)\n",
    "            test_loss, test_acc = test(model, test_loader)\n",
    "            t = time() - t\n",
    "            sizes.append((comp / 1000, test_acc, accuracy, t))\n",
    "        # continue reducing sparsity\n",
    "        print(i, \"iteration - \", \"Size:\", comp, list_of_sparsity)\n",
    "    result.append({'run_id': run_id, 'depth': depth,'input_width':input_width, \n",
    "                   'output': output_width, 'leaf_width':leaf_width,\n",
    "                   'buffer_size': buffer_size, 'sizes': sizes})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dc73b66a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'run_id': '543fe9acb34441a3a82b09ca2ef6046c', 'depth': 4, 'input_width': 784, 'output': 10, 'leaf_width': 32, 'buffer_size': 90, 'sizes': [(1675.908, 0.9274, 0.9274, 52.47797870635986), (96.90447921813978, 0.8963, 89.63, 48.06929326057434), (89.85803129632579, 0.9121, 91.21, 48.54480266571045), (77.8086053500239, 0.9095, 90.95, 49.80986571311951), (68.04857033351935, 0.9013, 90.13, 49.315375089645386), (59.98136795040668, 0.9005, 90.05, 47.807331800460815), (49.72946803982941, 0.8852, 88.52, 43.566579818725586)]}\n",
      "{'run_id': '4ea2391b3d014e2fafff3accfb352d2c', 'depth': 4, 'input_width': 784, 'output': 10, 'leaf_width': 32, 'buffer_size': 90, 'sizes': [(1675.908, 0.8077, 0.8077, 34.018182039260864), (96.90447921813978, 0.8872, 88.72, 33.80422616004944), (89.85803129632579, 0.8957, 89.57, 34.04274368286133), (77.8086053500239, 0.8892, 88.92, 34.13646936416626), (68.04857033351935, 0.8897, 88.97, 34.019927740097046), (59.98136795040668, 0.8861, 88.61, 34.12063121795654), (49.72946803982941, 0.8931, 89.31, 34.19474697113037)]}\n",
      "{'run_id': '612d1573468c4ffaa89bf54d60ce4508', 'depth': 3, 'input_width': 784, 'output': 10, 'leaf_width': 32, 'buffer_size': 90, 'sizes': [(836.388, 0.9396, 0.9396, 43.909788846969604), (97.64177676520247, 0.9391, 93.91, 43.66941213607788), (85.474599179814, 0.9434, 94.34, 47.08596754074097), (75.61918533564933, 0.9445, 94.45, 44.3426558971405), (67.63630012187598, 0.9454, 94.54, 47.12140870094299), (58.413546788847576, 0.9392, 93.92, 44.264766454696655), (48.85066364816289, 0.9288, 92.88, 47.297189235687256)]}\n",
      "{'run_id': '24879e69bd174ef6a9973359ea1b9b6c', 'depth': 3, 'input_width': 784, 'output': 10, 'leaf_width': 32, 'buffer_size': 90, 'sizes': [(836.388, 0.7436, 0.7436, 35.06636023521423), (97.64177676520247, 0.8728, 87.28, 31.641921997070312), (85.474599179814, 0.8826, 88.26, 32.54939413070679), (75.61918533564933, 0.8868, 88.68, 37.319536447525024), (67.63630012187598, 0.89, 89.0, 35.424580335617065), (58.413546788847576, 0.8922, 89.22, 33.795485496520996), (48.85066364816289, 0.8601, 86.01, 31.7235746383667)]}\n"
     ]
    }
   ],
   "source": [
    "for r in result:\n",
    "    print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c02f791f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range (4,len(list_of_run)):\n",
    "    run_id = list_of_run[i]\n",
    "    mlflow.artifacts.download_artifacts(run_id=run_id, dst_path=\".\")\n",
    "    wrapped_model = pickle.load(open(\"./truncated_model.pkl\", \"rb\"))\n",
    "    depth = wrapped_model._fff.fff.depth.item()\n",
    "    input_width = wrapped_model._fff.fff.input_width\n",
    "    leaf_width = wrapped_model._fff.fff.leaf_width\n",
    "    output_width = wrapped_model._fff.fff.output_width\n",
    "    buffer_size = 2*(leaf_width + output_width + 3)\n",
    "    print(\"Run:\\t\", run_id)\n",
    "    print(\"Depth:\\t\", depth)\n",
    "    print(\"Input:\\t\", input_width)\n",
    "    print(\"Output:\\t\", output_width)\n",
    "    print(\"Leaf:\\t\", leaf_width)\n",
    "    print(\"Buffer:\\t\", buffer_size)\n",
    "    \n",
    "    # to reduce the sparsity and train only below a certain tresholds\n",
    "    list_of_sizes = [100, 90, 80, 70, 60, 50]\n",
    "    checked_sizes = [False for x in list_of_sizes]\n",
    "    current_size_index = 0\n",
    "    \n",
    "    start = 0.5\n",
    "    a = start\n",
    "    b = start\n",
    "    sizes=[]\n",
    "    before_trunc_sizes=[]\n",
    "    trunc_sizes=[]\n",
    "\n",
    "    model = wrapped_model.to(device)\n",
    "    \n",
    "    layers_list = []\n",
    "    for i, (name, p) in enumerate(model.named_parameters()):\n",
    "        if (len(list(p.shape)) > 1 and p.requires_grad):\n",
    "            layers_list.append(p)\n",
    "    un, comp, layers = print_size_fc(model, layers_list, [1,1,1,1,1,1])\n",
    "    \n",
    "    MODEL_NAME_COMPRESSED = \"mnist_\" + run_id + \"_compressed_full\"\n",
    "    \n",
    "    layers_sizes = layers.copy()\n",
    "    # Creating a list of 1s with the same length as original_list using multiplication\n",
    "    list_of_sparsity = [1] * len(layers_sizes)\n",
    "    \n",
    "    # Testing time and accuracy\n",
    "    model.eval()\n",
    "    t = time()\n",
    "    train_loss, train_acc = test(model, train_loader)\n",
    "    test_loss, test_acc = test(model, test_loader)\n",
    "    t = time() - t\n",
    "    un, comp, layers_sizes = print_size_fc(model, layers_list, list_of_sparsity)\n",
    "    sizes.append((comp/1000, test_acc, test_acc, t))\n",
    "    \n",
    "    for i in range(1, 100):\n",
    "        # get index of max\n",
    "        index_of_max = np.argmax(layers_sizes)\n",
    "        current_sparsity = list_of_sparsity[index_of_max]\n",
    "\n",
    "        # reduce\n",
    "        if (current_sparsity == 1):\n",
    "            list_of_sparsity[index_of_max] = start - (start * 0.1)\n",
    "        else:\n",
    "            list_of_sparsity[index_of_max] = current_sparsity - (current_sparsity * 0.1)\n",
    "\n",
    "        un, comp, layers_sizes = print_size_fc(model, layers_list, list_of_sparsity)\n",
    "        \n",
    "        if (current_size_index >= len(list_of_sizes)):\n",
    "            break\n",
    "        elif (comp / 1000 < list_of_sizes[current_size_index]):\n",
    "            # compress and save result\n",
    "            MODEL_NAME_COMPRESSED = \"mnist_\" + run_id + \"_compressed_\" + str(list_of_sizes[current_size_index])\n",
    "            current_size_index += 1\n",
    "            model.train()\n",
    "            accuracy = perform_compression(model, layers_list, list_of_sparsity, learning_rate, num_epochs,\n",
    "                                           train_loader, test_loader, device,\n",
    "                                           val_loader=val_loader, model_name=MODEL_NAME_COMPRESSED, given_criterion=criterion)\n",
    "            model.load_state_dict(torch.load(MODEL_NAME_COMPRESSED+\".h5\", map_location='cpu'))\n",
    "            model.eval()\n",
    "            t = time()\n",
    "            train_loss, train_acc = test(model, train_loader)\n",
    "            test_loss, test_acc = test(model, test_loader)\n",
    "            t = time() - t\n",
    "            sizes.append((comp / 1000, test_acc, accuracy, t))\n",
    "        # continue reducing sparsity\n",
    "        print(i, \"iteration - \", \"Size:\", comp, list_of_sparsity)\n",
    "    result.append({'run_id': run_id, 'depth': depth,'input_width':input_width, \n",
    "                   'output': output_width, 'leaf_width':leaf_width,\n",
    "                   'buffer_size': buffer_size, 'sizes': sizes})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f3397a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for r in result:\n",
    "    print(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fee5d86",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FFFWrapper(torch.nn.Module):\n",
    "    def __init__(self, fff):\n",
    "        super(FFFWrapper, self).__init__()\n",
    "        self._fff = fff\n",
    "        self._fastinference = [None for i in range(2 ** (self._fff.fff.depth.item()))]\n",
    "\n",
    "    def forward(self, x, return_nodes=False):\n",
    "        \"\"\"\n",
    "        Override the forward method in order to log the data distribution.\n",
    "        \"\"\"\n",
    "        x = x.view(len(x), -1)\n",
    "        original_shape = x.shape\n",
    "        batch_size = x.shape[0]\n",
    "        last_node = torch.zeros(len(x))\n",
    "\n",
    "        current_nodes = torch.zeros((batch_size,), dtype=torch.long, device=x.device)\n",
    "        for i in range(self._fff.fff.depth.item()):\n",
    "            plane_coeffs = self._fff.fff.node_weights.index_select(dim=0, index=current_nodes)\n",
    "            plane_offsets = self._fff.fff.node_biases.index_select(dim=0, index=current_nodes)\n",
    "            plane_coeff_score = torch.bmm(x.unsqueeze(1), plane_coeffs.unsqueeze(-1))\n",
    "            plane_score = plane_coeff_score.squeeze(-1) + plane_offsets\n",
    "            plane_choices = (plane_score.squeeze(-1) >= 0).long()\n",
    "\n",
    "            platform = torch.tensor(2 ** i - 1, dtype=torch.long, device=x.device)\n",
    "            next_platform = torch.tensor(2 ** (i+1) - 1, dtype=torch.long, device=x.device)\n",
    "            current_nodes = (current_nodes - platform) * 2 + plane_choices + next_platform\n",
    "\n",
    "        leaves = current_nodes - next_platform\n",
    "        new_logits = torch.empty((batch_size, self._fff.fff.output_width), dtype=torch.float, device=x.device)\n",
    "        last_node = leaves\n",
    "\n",
    "        for i in range(leaves.shape[0]):\n",
    "            leaf_index = leaves[i]\n",
    "            if self._fastinference[leaf_index] is not None:\n",
    "                new_logits[i] = self._fastinference[leaf_index]\n",
    "            else:\n",
    "                logits = torch.matmul( x[i].unsqueeze(0), self._fff.fff.w1s[leaf_index])\n",
    "                logits += self._fff.fff.b1s[leaf_index].unsqueeze(-2)\n",
    "                activations = self._fff.fff.activation(logits)\n",
    "                new_logits[i] = torch.matmul( activations, self._fff.fff.w2s[leaf_index]).squeeze(-2)\n",
    "\n",
    "        if return_nodes:\n",
    "            return new_logits.view(*original_shape[:-1], self._fff.fff.output_width), last_node\n",
    "        return new_logits.view(*original_shape[:-1], self._fff.fff.output_width)\n",
    "\n",
    "\n",
    "    def simplify_leaves(self, trainloader):\n",
    "        y, leaves = (get_dist(self, trainloader))\n",
    "        y = y.cpu().detach().numpy()\n",
    "        outputs = y.max() + 1\n",
    "        leaves = leaves.cpu().detach().numpy()\n",
    "\n",
    "        n_simplifications = 0\n",
    "        ratios = {}\n",
    "        for l in np.unique(leaves):\n",
    "            ratios[l] = torch.zeros(outputs)\n",
    "            indices = leaves == l\n",
    "\n",
    "            for i in range(outputs):\n",
    "                ratios[l][i] = (np.sum(y[indices] == i) / np.sum(indices))\n",
    "\n",
    "            argmax = np.argmax(ratios[l])\n",
    "            if ratios[l][argmax] >= 0:\n",
    "                output = torch.zeros(outputs)\n",
    "                output[argmax] = 1\n",
    "                self._fastinference[l] = output\n",
    "                n_simplifications += 1\n",
    "                print(f\"Leaf {l} has been replaced with {argmax}\")\n",
    "        print(self._fastinference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d7faec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import typer\n",
    "import mlflow\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def get_split_code(array, bias):\n",
    "    code = \"\"\"\n",
    "    acc = \"\"\" + \" + \".join(f\"x[{i}] * {v}\" for i, v in enumerate(array)) + \"\"\";\n",
    "    acc += \"\"\" + str(bias[0]) + \"\"\";\n",
    "    \"\"\"\n",
    "    return code\n",
    "\n",
    "\n",
    "def get_output_code(w1, b1, w2, b2):\n",
    "    code = \"\"\"\n",
    "    float hidden[\"\"\" + str(w1.shape[1]) + \"\"\"];\n",
    "    \"\"\"\n",
    "    for i in range(w1.shape[1]):\n",
    "        code += f\"hidden[{i}] = {b1[i]} + \" + \" + \".join(f\"x[{j}] * {v}\" for j, v in enumerate(w1[:, i])) + \";\\n\"\n",
    "        code += f\"hidden[{i}] = hidden[{i}] > 0 ? hidden[{i}] : 0;\\n\"\n",
    "    code += \"\"\"\n",
    "    float logits[\"\"\" + str(w2.shape[1]) + \"\"\"];\n",
    "    \"\"\"\n",
    "    for j in range(w2.shape[1]):\n",
    "        code += f\"logits[{j}] = {b2[j]} + \" + \" + \".join(f\"hidden[{i}] * {v}\" for i, v in enumerate(w2[:, j])) + \";\\n\"\n",
    "        code += f\"logits[{j}] = logits[{j}] > 0 ? logits[{j}] : 0;\\n\"\n",
    "\n",
    "    code += \"\"\"\n",
    "    max = 0.0;\n",
    "    argmax = 0;\n",
    "    for (int i = 0; i < \"\"\" + str(w2.shape[1]) + \"\"\"; i++) {\n",
    "        if (logits[i] > max) {\n",
    "            max = logits[i];\n",
    "            argmax = i;\n",
    "        }\n",
    "    }\n",
    "\n",
    "    return argmax;\n",
    "    \"\"\"\n",
    "\n",
    "    return code\n",
    "\n",
    "def get_splits(weights, biases):\n",
    "    code = \"\"\"int perform_inference(float* x) {\n",
    "    float acc;\n",
    "    float max;\n",
    "    int argmax;\n",
    "        <replaceme>\n",
    "}\"\"\"\n",
    "    for index, (array, bias) in enumerate(zip(weights, biases)):\n",
    "        code = code.replace(\"<replaceme>\", get_split_code(array, bias), 1)\n",
    "    return code\n",
    "\n",
    "\n",
    "class Node:\n",
    "    def __init__(self, array, bias, left, right):\n",
    "        self._array = array\n",
    "        self._bias = bias\n",
    "        self._left = left\n",
    "        self._right = right\n",
    "\n",
    "    def __str__(self):\n",
    "        code = get_split_code(self._array, self._bias)\n",
    "\n",
    "        code += \"\"\"\n",
    "        if (acc >= 0) {\n",
    "            \"\"\" + str(self._left) + \"\"\"\n",
    "        } else {\n",
    "            \"\"\" + str(self._right) + \"\"\"\n",
    "        }\n",
    "        \"\"\"\n",
    "        return code\n",
    "\n",
    "\n",
    "class Leaf(Node):\n",
    "    def __init__(self, w1, b1, w2, b2):\n",
    "        self._w1 = w1\n",
    "        self._b1 = b1\n",
    "        self._w2 = w2\n",
    "        self._b2 = b2\n",
    "\n",
    "    def __str__(self):\n",
    "        if self._w2 is None:\n",
    "            return f\"return {self._w1};\\n\"\n",
    "        return get_output_code(self._w1, self._b1, self._w2, self._b2)\n",
    "\n",
    "def print_parameters(params, key, lines, line, flit, skipTruncatedLeaves, sizes_print=False, index_print=True):\n",
    "    # the first check on the sparsity on the weight also consider the non zero values of truncated leaves\n",
    "    \n",
    "    weight = params[key]\n",
    "    dim = len(weight.shape)\n",
    "    \n",
    "    non_zero_values = np.count_nonzero(weight)\n",
    "    first_dim = weight.shape[0]\n",
    "    sparse = False\n",
    "    \n",
    "    if (dim == 2):\n",
    "        weight_row = weight.shape[0]\n",
    "        weight_col = weight.shape[1]\n",
    "        \n",
    "        sparse = non_zero_values < ((weight_row * (weight_col - 1) - 1) / 2)\n",
    "    elif (dim == 3):\n",
    "        weight_depth = weight.shape[0]\n",
    "        weight_row = weight.shape[1]\n",
    "        weight_col = weight.shape[2]\n",
    "        \n",
    "        size_requested = weight_depth * weight_row * weight_col\n",
    "        sparse = non_zero_values < ((size_requested - weight_depth) / 2)\n",
    "    if (not sparse):\n",
    "        # print the parameters as usal\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"#define \" + key + \"_SPARSE \" + str(0) + \"\\n\"\n",
    "        )\n",
    "        line += 1\n",
    "        line += 1\n",
    "        \n",
    "        param = weight\n",
    "        if (key in ['LEAF_HIDDEN_WEIGHTS', 'LEAF_HIDDEN_BIASES', 'LEAF_OUTPUT_WEIGHTS', 'LEAF_OUTPUT_BIASES']):\n",
    "            if (skipTruncatedLeaves):\n",
    "                param = param[params['FASTINFERENCE'] == -1]\n",
    "        param = param.flatten()\n",
    "        tmp = \"\"\n",
    "        if (flit and key not in ['FASTINFERENCE']):\n",
    "            tmp = \", \".join([\"F_LIT(\" + str(x) + \")\" for x in param])\n",
    "        else:\n",
    "            tmp = \", \".join([str(x) for x in param])\n",
    "        lines.insert(\n",
    "            line,\n",
    "            tmp\n",
    "        )\n",
    "    else:\n",
    "        # CSC or CSR format\n",
    "        leaves_values = np.empty([0], dtype=float)\n",
    "        leaves_offsets = np.empty([0], dtype=int)\n",
    "        leaves_sizes = None\n",
    "        printed_dim = key + \"_DIM_1\"\n",
    "        if (sizes_print):\n",
    "            leaves_sizes = np.empty([first_dim], dtype=int)\n",
    "        elif(index_print):\n",
    "            leaves_sizes = np.zeros([1], dtype=int)\n",
    "            printed_dim += \" + 1\"\n",
    "            \n",
    "        value_position = 0\n",
    "        actual_non_zero_values = 0\n",
    "        \n",
    "        for index, leaf_weight in enumerate(weight): # from 0 to first_dim\n",
    "            \n",
    "            if (key in ['LEAF_HIDDEN_WEIGHTS', 'LEAF_HIDDEN_BIASES', 'LEAF_OUTPUT_WEIGHTS', 'LEAF_OUTPUT_BIASES']):\n",
    "                # insert filters non zero values into the fitler sizes\n",
    "                if (params['FASTINFERENCE'][index] != -1 and skipTruncatedLeaves):\n",
    "                    leaves_sizes[index] = 0\n",
    "                    continue\n",
    "            \n",
    "            # insert filters non zero values into the fitler sizes\n",
    "            non_zero_values_here = np.count_nonzero(leaf_weight)\n",
    "            actual_non_zero_values += non_zero_values_here\n",
    "            # flatten the filter\n",
    "            flatten_leaf = leaf_weight.ravel()\n",
    "            offset = 1\n",
    "            \n",
    "            if (sizes_print):\n",
    "                # _sizes[DIM_1] = for each node/leaf we save the number of NNZ contained\n",
    "                leaves_sizes[index] = non_zero_values_here\n",
    "            elif (index_print):\n",
    "                # _sizes[DIM_1 + 1] = for each node/leaf we save the index in _data array of the starting value\n",
    "                # _sizes from {0 up to NNZ}\n",
    "                leaves_sizes = np.append(leaves_sizes, actual_non_zero_values)\n",
    "            \n",
    "            for value in flatten_leaf: # from 0 to (n_depth * n_height * n_width)\n",
    "                if (value == 0):\n",
    "                    # increase offset\n",
    "                    offset += 1\n",
    "                else:\n",
    "                    # save value, save index, reset offset, increase position\n",
    "                    leaves_values = np.append(leaves_values, value)\n",
    "                    leaves_offsets = np.append(leaves_offsets, offset)\n",
    "                    leaves_values[value_position] = value\n",
    "                    leaves_offsets[value_position] = offset\n",
    "                    offset = 1\n",
    "                    value_position += 1\n",
    "                \n",
    "            \n",
    "        tmp = \"\"\n",
    "        # substitute the definition\n",
    "        lines[line] = \"#define \" + key + \"_NNZ \" + str(actual_non_zero_values) + \"\\n\"\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"#define \" + key + \"_DIM \" + str(dim) + \"\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"#define \" + key + \"_SPARSE \" + str(1) + \"\\n\"\n",
    "        )\n",
    "        for d in range(0, dim):\n",
    "            line+=1\n",
    "            lines.insert(\n",
    "                line,\n",
    "                \"#define \" + key + \"_DIM_\" + str(d + 1) + \" \" + str(weight.shape[d]) + \"\\n\"\n",
    "            )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"__hifram fixed \" + key + \"_data[\" + key + \"_NNZ] = {\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        \n",
    "        if flit:\n",
    "            tmp = \", \".join([\"F_LIT(\" + str(x) + \")\" for x in leaves_values])\n",
    "        else:\n",
    "            tmp = \", \".join([str(x) for x in leaves_values])\n",
    "        lines.insert(\n",
    "            line,\n",
    "            tmp\n",
    "        )\n",
    "        line+=1\n",
    "        line+=1\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"\\n__hifram fixed \" + key + \"_offset[\" + key + \"_NNZ] = {\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        if (flit and False):\n",
    "            tmp = \", \".join([\"F_LIT(\" + str(x) + \")\" for x in leaves_offsets])\n",
    "        else:\n",
    "            tmp = \", \".join([str(x) for x in leaves_offsets])\n",
    "        lines.insert(\n",
    "            line,\n",
    "            tmp\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"};\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"\\n__hifram fixed \" + key + \"_sizes[\" + printed_dim + \"] = {\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        if (flit and False):\n",
    "            tmp = \", \".join([\"F_LIT(\" + str(x) + \")\" for x in leaves_sizes])\n",
    "        else:\n",
    "            tmp = \", \".join([str(x) for x in leaves_sizes])\n",
    "        lines.insert(\n",
    "            line,\n",
    "            tmp\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"\\n\"\n",
    "        )\n",
    "        line+=1\n",
    "        lines.insert(\n",
    "            line,\n",
    "            \"};\\n\\n\"\n",
    "        )\n",
    "    \n",
    "def make_program(wrapped_model, name, original_fastinference, flit=True, skipTruncatedLeaves=False):\n",
    "\n",
    "    node_weights = wrapped_model._fff.fff.node_weights.cpu().detach().numpy()\n",
    "    node_biases = wrapped_model._fff.fff.node_biases.cpu().detach().numpy()\n",
    "    w1s = wrapped_model._fff.fff.w1s\n",
    "    b1s = wrapped_model._fff.fff.b1s.cpu().detach().numpy()\n",
    "    w2s = wrapped_model._fff.fff.w2s\n",
    "    b2s = wrapped_model._fff.fff.b2s.cpu().detach().numpy()\n",
    "    fastinference = wrapped_model._fastinference\n",
    "\n",
    "    w1s = w1s.transpose(1, 2).cpu().detach().numpy()\n",
    "    w2s = w2s.transpose(1, 2).cpu().detach().numpy()\n",
    "\n",
    "    params = {}\n",
    "\n",
    "    params['NODE_WEIGHTS'] = node_weights#.flatten()\n",
    "    params['NODE_BIASES'] = node_biases#.flatten()\n",
    "    params['FASTINFERENCE'] = np.array([-1 if x is None else int(x.argmax()) for x in fastinference])\n",
    "    actual_leaves_weights = w1s #[params['FASTINFERENCE'] == -1]\n",
    "    actual_leaves_biases = b1s #[params['FASTINFERENCE'] == -1]\n",
    "    actual_leaves_out_weights = w2s #[params['FASTINFERENCE'] == -1]\n",
    "    actual_leaves_out_biases = b2s #[params['FASTINFERENCE'] == -1]\n",
    "    params['LEAF_HIDDEN_WEIGHTS'] = actual_leaves_weights#.flatten()\n",
    "    params['LEAF_HIDDEN_BIASES'] = actual_leaves_biases#.flatten()\n",
    "    params['LEAF_OUTPUT_WEIGHTS'] = actual_leaves_out_weights#.flatten()\n",
    "    params['LEAF_OUTPUT_BIASES'] = actual_leaves_out_biases#.flatten()\n",
    "    with open(\"weights_\" + name + \".h\", \"w\") as f:\n",
    "        with open(\"weights_template.h\") as in_f:\n",
    "            lines = in_f.readlines()\n",
    "\n",
    "            i = 0\n",
    "            while i < len(lines):\n",
    "                i += 1\n",
    "                if \"Add definitions here\" in lines[i]:\n",
    "                    break\n",
    "            lines[i] = (f\"\"\"// name of the model  {name}\n",
    "\n",
    "#define DEPTH {wrapped_model._fff.fff.depth.item()}\n",
    "#define N_LEAVES {2 ** wrapped_model._fff.fff.depth.item()}\n",
    "#define INPUT_SIZE {wrapped_model._fff.fff.input_width}\n",
    "#define LEAF_WIDTH {wrapped_model._fff.fff.leaf_width}\n",
    "#define OUTPUT_SIZE {wrapped_model._fff.fff.output_width}\n",
    "#define SIMPLIFIED_LEAVES {sum([f is not None for f in fastinference])}\n",
    "#define ORIGINAL_FASTINFERENCE {original_fastinference}\"\"\")\n",
    "            i+=1\n",
    "            lines.insert(i, \"\\n\")\n",
    "\n",
    "            # lines.insert(i+7, \"\"\"\n",
    "# float FASTINFERENCE[N_LEAVES] = {-1};\n",
    "# float NODE_WEIGHTS[N_LEAVES-1][INPUT_SIZE];\n",
    "# float NODE_BIASES[N_LEAVES-1];\n",
    "# float LEAF_HIDDEN_WEIGHTS[N_LEAVES-SIMPLIFIED_LEAVES][LEAF_WIDTH][INPUT_SIZE];\n",
    "# float LEAF_OUTPUT_WEIGHTS[N_LEAVES-SIMPLIFIED_LEAVES][OUTPUT_SIZE][LEAF_WIDTH];\n",
    "# float LEAF_HIDDEN_BIASES[N_LEAVES-SIMPLIFIED_LEAVES][LEAF_WIDTH];\n",
    "# float LEAF_OUTPUT_BIASES[N_LEAVES-SIMPLIFIED_LEAVES][OUTPUT_SIZE];\n",
    "            # \"\"\")\n",
    "\n",
    "            for key in params.keys():\n",
    "                i = 0\n",
    "                while i < len(lines):\n",
    "                    if f\"fixed {key}\" in lines[i]:\n",
    "                        break\n",
    "                    i += 1\n",
    "#                 i += 1\n",
    "                # line of the definition of the parameter\n",
    "                # calculate sparsity and choose representation\n",
    "                # if dense keep the definition, i+=1, and print\n",
    "                # if sparse\n",
    "                print_parameters(params, key, lines, i, flit, skipTruncatedLeaves)\n",
    "                # print Parameters\n",
    "                \n",
    "                # calculate sparsity and choose representation\n",
    "                # if dense -> print as we already are doing\n",
    "                # if not change the representation\n",
    "                \n",
    "\n",
    "            f.writelines(lines)\n",
    "    return wrapped_model\n",
    "\n",
    "\n",
    "def main(run_id, name, original_fastinference):\n",
    "#     import torch\n",
    "    net = make_program(run_id, name, original_fastinference)\n",
    "#     net._fff.eval()\n",
    "#     X = np.loadtxt('test.txt')\n",
    "#     X = torch.Tensor(X)\n",
    "#     with open('ref_outputs.txt', 'w') as f:\n",
    "#         y = net(X).argmax(1)\n",
    "#         y = [(str(x) + \"\\n\") for x in y.detach().cpu().numpy()]\n",
    "#         f.writelines(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a98e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range (0, len(list_of_run)):\n",
    "    run_id = list_of_run[i]\n",
    "    \n",
    "    mlflow.artifacts.download_artifacts(run_id=run_id, dst_path=\".\")\n",
    "    wrapped_model = pickle.load(open(\"./truncated_model.pkl\", \"rb\"))\n",
    "    \n",
    "    original_fastinference = str([-1 if x is None else int(x.argmax()) for x in wrapped_model._fastinference])\n",
    "    \n",
    "    # compress and save result\n",
    "    MODEL_NAME_COMPRESSED = \"mnist_\" + run_id + \"_compressed_\" + str(list_of_sizes[0])\n",
    "    wrapped_model.load_state_dict(torch.load(MODEL_NAME_COMPRESSED+\".h5\", map_location='cpu'))\n",
    "    \n",
    "    n_simplifications = wrapped_model.to(device).simplify_leaves(train_loader)\n",
    "    main(wrapped_model, MODEL_NAME_COMPRESSED, original_fastinference)"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}