PPG Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c0bb7f94",
   "metadata": {},
   "source": [
    "# Learned Kernels\n",
    "This notebook contains all of the code needed to set up and run SMoLK model on the various PPG quality datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18d6f91c",
   "metadata": {},
   "source": [
    "## Setup\n",
    "Here, we import all of the necessary libraries we need to run this code, as well as set the random seed for reproducibility. Note that we set device to \"cuda\" to run on a CUDA-enabled GPU. If you do not have a CUDA enabled GPU, you can set this to \"cpu\" instead, though it will be much slower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28f62bf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os #path/directory stuff\n",
    "import pickle\n",
    "\n",
    "#Deep learning\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "#Math\n",
    "from sklearn.metrics import f1_score, r2_score\n",
    "from scipy.signal import savgol_filter\n",
    "import numpy as np\n",
    "import copy\n",
    "import random\n",
    "\n",
    "#Set seed for reproducibility\n",
    "seed = 0\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "#Progress bar\n",
    "from tqdm import tqdm\n",
    "\n",
    "device = \"cuda\" # device to use"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b2ba0f56",
   "metadata": {},
   "source": [
    "# Load the Data\n",
    "Here, we load the training and test sets. We use only the DaLiA training set for training, and the DaLiA test set for testing. We load the other datasets to test out-of-distribution performance. These datasets are all included in the GitHub repository."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95ed9ab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = \"PPG Data\" # Base directory for the PPG data\n",
    "subdirs = [ # sub dirs that contain each PPG dataset\n",
    "\"new_PPG_DaLiA_test/processed_dataset\",\n",
    "\"new_PPG_DaLiA_train/processed_dataset\",\n",
    "\"TROIKA_channel_1/processed_dataset\",\n",
    "\"WESAD_all/processed_dataset\"]\n",
    "\n",
    "# We use the DaLiA train set exclusively for training\n",
    "X_train = np.load(os.path.join(base, subdirs[1], \"scaled_ppgs.npy\"))\n",
    "Y_train = np.load(os.path.join(base, subdirs[1], \"seg_labels.npy\"))\n",
    "\n",
    "# The rest of these datasets are test\n",
    "DaLiA_X = np.load(os.path.join(base, subdirs[0], \"scaled_ppgs.npy\"))\n",
    "DaLiA_Y = np.load(os.path.join(base, subdirs[0], \"seg_labels.npy\"))\n",
    "\n",
    "TROIKA_X = np.load(os.path.join(base, subdirs[2], \"scaled_ppgs.npy\"))\n",
    "TROIKA_Y = np.load(os.path.join(base, subdirs[2], \"seg_labels.npy\"))\n",
    "\n",
    "WESAD_X = np.load(os.path.join(base, subdirs[3], \"scaled_ppgs.npy\"))\n",
    "WESAD_Y = np.load(os.path.join(base, subdirs[3], \"seg_labels.npy\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "21b428bc",
   "metadata": {},
   "source": [
    "# Define the model\n",
    "Below is the definition of the Learned Kernels model. It simply consists of three `Conv1d` operations aggregated and summed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1800aba6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LearnedFilters(nn.Module):\n",
    "    def __init__(self, num_kernels=24):\n",
    "        super(LearnedFilters, self).__init__()\n",
    "        self.conv1 = nn.Conv1d(1, num_kernels, 192, stride=1, padding=\"same\", bias=True)\n",
    "        self.conv2 = nn.Conv1d(1, num_kernels, 96, stride=1, padding=\"same\", bias=True)\n",
    "        self.conv3 = nn.Conv1d(1, num_kernels, 64, stride=1, padding=\"same\", bias=True)\n",
    "        \n",
    "        self.w1 = torch.nn.Parameter(torch.zeros(num_kernels), requires_grad=True) #these are learned weights for the kernels        \n",
    "        self.w2 = torch.nn.Parameter(torch.zeros(num_kernels), requires_grad=True)        \n",
    "        self.w3 = torch.nn.Parameter(torch.zeros(num_kernels), requires_grad=True)\n",
    "    \n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.shape[0]\n",
    "        \n",
    "        c1 = F.relu(self.conv1(x)) * self.w1[None,:,None]\n",
    "        c2 = F.relu(self.conv2(x)) * self.w2[None,:,None]\n",
    "        c3 = F.relu(self.conv3(x)) * self.w3[None,:,None]\n",
    "        \n",
    "        aggregate = torch.cat([c1,c2,c3], dim=1)\n",
    "        aggregate = aggregate.sum(dim=1).view(batch_size, -1)\n",
    "        aggregate = torch.sigmoid(aggregate)\n",
    "        \n",
    "        return aggregate"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5a38a7c2",
   "metadata": {},
   "source": [
    "# Train the models\n",
    "This should take around a day to train all 90 model folds on a decent consumer GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee97573f",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setup\n",
    "filter_nums = [4, 8, 16, 24, 32, 64, 128, 256, 512] #number of filters to train models for\n",
    "folds = 10 #number of folds to use for cross validation\n",
    "epochs = 512 #number of epochs to train for\n",
    "lr = 0.01 #learning rate\n",
    "wd = 1e-4 #weight decay\n",
    "decay_range = [1.0, 0.2] #range of decay values to use for learning rate decay\n",
    "\n",
    "save_dir = \"models\" #directory to save models to\n",
    "if not os.path.exists(save_dir):\n",
    "    os.makedirs(save_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4264919a",
   "metadata": {},
   "source": [
    "## Main training loop\n",
    "Here, we train `folds` number of models for each `filter_nums` number of filters. Keep in mind that the end model will have `filter_num*3` filters, since the `filter_num` refers to the number of filters in each kernel group, for which there are three"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54eaba4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop through different numbers of filters\n",
    "for filter_num in filter_nums:\n",
    "    # Loop through different versions of the model\n",
    "    for fold in range(0, folds):\n",
    "        # Initialize a new instance of the LearnedFilters class with the current number of filters\n",
    "        net = LearnedFilters(filter_num).to(device)\n",
    "\n",
    "        # Compute the total number of model parameters and print it\n",
    "        params = sum([np.prod(p.size()) for p in filter(lambda p: p.requires_grad, net.parameters())])\n",
    "        print(f\"Training kernel with {filter_num} filters (fold {fold+1})...\")\n",
    "        print(\"Num params: %i\" % params)\n",
    "\n",
    "        # Initialize the optimizer\n",
    "        optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)\n",
    "\n",
    "        # Initialize a linear learning rate scheduler\n",
    "        scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=decay_range[0], end_factor=decay_range[1], total_iters=epochs)\n",
    "\n",
    "        # Initialize a progress bar for visualization\n",
    "        pbar = tqdm(range(0, epochs))\n",
    "\n",
    "        # Normalize the input data by subtracting the mean and dividing by the standard deviation\n",
    "        x = copy.deepcopy(X_train) #we deepcopy the data so we don't modify the original\n",
    "        \n",
    "        #normalize each signal\n",
    "        for i in range(0, len(x)):\n",
    "            x[i] = (x[i] - np.mean(x[i]))/np.std(x[i])\n",
    "\n",
    "        # Convert the input and output data to PyTorch tensors and move them to the device\n",
    "        x = torch.tensor(x, dtype=torch.float32, device=device).reshape(len(x), 1, 1920)\n",
    "        y = torch.tensor(Y_train, dtype=torch.float32, device=device)\n",
    "\n",
    "        loss_hist = [] #history of losses (used to update the progress bar primarily)\n",
    "\n",
    "        # Train the model\n",
    "        for step in pbar:\n",
    "            # Zero out the gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Initialize the total loss to 0\n",
    "            total_loss = 0\n",
    "\n",
    "            # Split the input data into smaller batches if the number of filters is greater than 32 to avoid running out of memory\n",
    "            split = 32 if filter_num > 32 else 1\n",
    "            for i in range(0, split):\n",
    "                split_len = x.shape[0] // split\n",
    "                out = net(x[i*split_len:(i+1)*split_len]) # Forward pass\n",
    "                loss = F.binary_cross_entropy(out.view(-1), y[i*split_len:(i+1)*split_len].view(-1)) / split # Compute the loss\n",
    "                total_loss += loss.item() # Add the loss to the total loss\n",
    "                loss.backward() # Backward pass\n",
    "\n",
    "            optimizer.step() # Take an optimizer step\n",
    "            scheduler.step() # Adjust the learning rate\n",
    "\n",
    "            loss_hist.append(total_loss) # Add the total loss to the loss history\n",
    "            pbar.set_description(\"Loss: %.5f\" % np.mean(loss_hist[-256:])) # Update the progress bar with the current loss\n",
    "\n",
    "        # Save the model\n",
    "        torch.save(net, os.path.join(save_dir, f\"learned_filters_{filter_num}_{fold}.pt\"))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "43a0d317",
   "metadata": {},
   "source": [
    "# Test the models\n",
    "In this section, we simply run each model fold again the various test sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7b91424",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(net, X, Y):\n",
    "    \"\"\"\n",
    "    Test the given network on the provided data.\n",
    "\n",
    "    Args:\n",
    "        net (nn.Module): The trained model.\n",
    "        X (list): The list of data samples.\n",
    "        Y (list): The list of corresponding labels.\n",
    "\n",
    "    Returns:\n",
    "        float: The DICE score of the model's predictions.\n",
    "    \"\"\"\n",
    "    dtype = net.state_dict()['w1'].dtype # Get the data type of the model\n",
    "    device = net.state_dict()['w1'].device # Get the device of the model\n",
    "    data = X\n",
    "    labels = Y\n",
    "\n",
    "    preds = [] # To store the model predictions\n",
    "    true = [] # To store the true labels\n",
    "\n",
    "    for i in range(len(data)):\n",
    "        x = data[i]\n",
    "        x = (x - np.mean(x))/np.std(x) # Normalize the data\n",
    "\n",
    "        y_pred = net(torch.tensor(x, dtype=dtype, device=device).view(1, -1)) # Predictions\n",
    "\n",
    "        # Apply a Savitzky-Golay filter to the predictions and convert to binary form\n",
    "        preds += list(np.float32(savgol_filter(y_pred.detach().cpu().float().numpy()[0], 151, 3) > 0.5)) \n",
    "        true += list(labels[i]) # Concatenate labels as list\n",
    "\n",
    "    return f1_score(preds, true) # Return the F1 score, which in this case, is equivalent to DICE score\n",
    "\n",
    "def test_suite(net, verbose=True):\n",
    "    \"\"\"\n",
    "    Test the given network on all the datasets and return and/or print the DICE scores\n",
    "\n",
    "    Args:\n",
    "        net (nn.Module): The trained model.\n",
    "        verbose (bool, optional): Whether to print the DICE scores. Defaults to True.\n",
    "\n",
    "    Returns:\n",
    "        list: The DICE scores for each dataset.\n",
    "    \"\"\"\n",
    "    results = [0, 0, 0] # Will store the F1 scores for each dataset\n",
    "\n",
    "    # Test on the DaLiA dataset\n",
    "    results[0] = test_model(net, DaLiA_X, DaLiA_Y)\n",
    "    if verbose:\n",
    "        print(\"DaLiA DICE score: %.4f\" % results[0])\n",
    "\n",
    "    # Test on the TROIKA dataset\n",
    "    results[1] = test_model(net, TROIKA_X, TROIKA_Y)\n",
    "    if verbose:\n",
    "        print(\"TROIKA DICE score: %.4f\" % results[1])\n",
    "\n",
    "    # Test on the WESAD dataset\n",
    "    results[2] = test_model(net, WESAD_X, WESAD_Y)\n",
    "    \n",
    "    if verbose:\n",
    "        print(\"WESAD DICE score: %.4f\" % results[2])\n",
    "    \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a06ac3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_results = {} #will store the results of each model\n",
    "\n",
    "# Loop through different numbers of filters\n",
    "for filter_num in filter_nums:\n",
    "    print(f\"Testing kernel with {filter_num} filters...\")\n",
    "    # Loop through different versions of the model\n",
    "    pbar = tqdm(range(0, folds))\n",
    "    results = []\n",
    "    for fold in pbar: #loop through all the folds\n",
    "        # Load the model\n",
    "        net = torch.load(os.path.join(save_dir, f\"learned_filters_{filter_num}_{fold}.pt\"), map_location=device)\n",
    "\n",
    "        # Test the model\n",
    "        results.append(test_suite(net, verbose=False))\n",
    "\n",
    "        # Update the progress bar\n",
    "        pbar.set_description(f\"DaLiA: %.4f, TROIKA: %.4f, WESAD: %.4f\" % tuple(np.mean(results, axis=0)))\n",
    "    \n",
    "    test_results[filter_num] = np.transpose(results) #transpose the results so that the rows are the datasets and the columns are the folds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e06a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "#save results\n",
    "with open(os.path.join(save_dir, \"test_results.pkl\"), \"wb\") as f:\n",
    "    pickle.dump(test_results, f)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3f6ec078",
   "metadata": {},
   "source": [
    "# Kernel Pruning\n",
    "This code removes similar kernels from each model and tests the results. We defer to the methods section of our manuscript for a detailed discussion of the pruning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b25bd58",
   "metadata": {},
   "outputs": [],
   "source": [
    "def similarity(v1, v2):\n",
    "    \"\"\"\n",
    "    Calculate the cosine similarity between two vectors.\n",
    "\n",
    "    Args:\n",
    "        v1, v2 (torch.Tensor): The input vectors.\n",
    "\n",
    "    Returns:\n",
    "        float: The cosine similarity between v1 and v2.\n",
    "    \"\"\"\n",
    "    norm_v1 = v1 / v1.norm()\n",
    "    norm_v2 = v2 / v2.norm()\n",
    "    \n",
    "    return (norm_v1*norm_v2).sum().item()\n",
    "\n",
    "def compute_param_num(num_conv1, num_conv2, num_conv3):\n",
    "    \"\"\"\n",
    "    Compute the number of parameters in a network given the number of kernels in each layer.\n",
    "\n",
    "    Args:\n",
    "        num_conv1, num_conv2, num_conv3 (int): The number of kernels in each convolutional layer.\n",
    "\n",
    "    Returns:\n",
    "        int: The total number of parameters in the network.\n",
    "    \"\"\"\n",
    "    params = num_conv1*192 +num_conv2*96 + num_conv3*64 #kernel params\n",
    "    params += num_conv1 + num_conv2 + num_conv3 #biases (1 per kernel)\n",
    "    params += num_conv1 + num_conv2 + num_conv3 #weights (1 per kernel)\n",
    "    \n",
    "    return params\n",
    "\n",
    "def get_most_similar_kernels(similarity_flat, coords):\n",
    "    \"\"\"\n",
    "    Get the indices of the most similar kernels based on their similarity scores.\n",
    "\n",
    "    Args:\n",
    "        similarity_flat (np.array): The flattened array of similarity scores.\n",
    "        coords (np.array): The flattened array of kernel index pairs.\n",
    "\n",
    "    Returns:\n",
    "        np.array: The indices of the most similar kernels.\n",
    "    \"\"\"\n",
    "    return coords[np.argsort(similarity_flat)]\n",
    "\n",
    "\n",
    "def compute_similarity(state_dict, conv_i, num_kernels):\n",
    "    \"\"\"\n",
    "    Compute the similarity between convolutional kernels for a given layer.\n",
    "\n",
    "    Args:\n",
    "        state_dict (dict): The state dict of the network.\n",
    "        conv_i (int): Index of the convolutional layer.\n",
    "        num_kernels (int): Number of kernels in each layer.\n",
    "\n",
    "    Returns:\n",
    "        tuple: Two numpy arrays containing the flattened similarity scores and their corresponding coordinates.\n",
    "    \"\"\"\n",
    "    coords = []\n",
    "    similarity_flat = []\n",
    "    \n",
    "    # Iterate over all pairs of kernels\n",
    "    for i in range(num_kernels):\n",
    "        for j in range(i, num_kernels):\n",
    "            if i != j:\n",
    "                sim = similarity(state_dict[f'conv{conv_i}.weight'][i], state_dict[f'conv{conv_i}.weight'][j])\n",
    "                similarity_flat.append(sim)\n",
    "                coords.append((j, i))\n",
    "\n",
    "    return np.asarray(similarity_flat), np.asarray(coords)\n",
    "\n",
    "\n",
    "def prune(state_dict, conv_i, num_kernels, prune_ratio):\n",
    "    \"\"\"\n",
    "    Prune the least important kernels from a kernel group based on cosine similarity and kernel importance.\n",
    "\n",
    "    Args:\n",
    "        state_dict (dict): The state dict of the network.\n",
    "        conv_i (int): Index of the kernel group.\n",
    "        num_kernels (int): Number of kernels in each group.\n",
    "        prune_ratio (float): The proportion of kernels to prune.\n",
    "\n",
    "    Returns:\n",
    "        dict: The updated state dict after pruning.\n",
    "    \"\"\"\n",
    "    # Compute similarity of kernels\n",
    "    sim_flat, coords = compute_similarity(state_dict, conv_i, num_kernels)\n",
    "    \n",
    "    # Get the most similar kernels\n",
    "    most_similar_kernels = get_most_similar_kernels(sim_flat, coords)\n",
    "\n",
    "    # Prune if the ratio is greater than zero, otherwise do nothing\n",
    "    if prune_ratio > 0:\n",
    "        # Iterate over the most similar kernels\n",
    "        for item in most_similar_kernels[-int(num_kernels*prune_ratio):]:\n",
    "            # Calculate weights for two kernels under consideration\n",
    "            item0_weight = state_dict[f'w{conv_i}'][item[0]]*state_dict[f'conv{conv_i}.weight'][item[0]].abs().mean()\n",
    "            item1_weight = state_dict[f'w{conv_i}'][item[1]]*state_dict[f'conv{conv_i}.weight'][item[1]].abs().mean()\n",
    "\n",
    "            # Decide which kernel to remove and which to keep\n",
    "            remove, keep, keep_weight, remove_weight = (item[1], item[0], item0_weight, item1_weight) if item0_weight > item1_weight else (item[0], item[1], item1_weight, item0_weight)\n",
    "            \n",
    "            # Update state_dict\n",
    "            state_dict[f'w{conv_i}'][keep] = (keep_weight + remove_weight) / state_dict[f'conv{conv_i}.weight'][keep].abs().mean()\n",
    "            state_dict[f'conv{conv_i}.bias'][keep] += state_dict[f'conv{conv_i}.bias'][remove]\n",
    "            \n",
    "            # this step is actually enough to \"prune the kernel\", for computation/measurement sake. In reality, we'd want to remove the kernel from the network for a speedup\n",
    "            state_dict[f'w{conv_i}'][remove] = 0.0\n",
    "            state_dict[f'conv{conv_i}.bias'][remove] = 0.0\n",
    "\n",
    "            # this is an extra step used for counting the kernel that we remove in the end\n",
    "            # no matter what this value is set to, it will not have any effect on the network since the weight is set to zero\n",
    "            # however, this value *does* need to be non-zero, since PyTorch handles completely zeroed convolutions a bit weirdly\n",
    "            # and you'll get weird results convolving a purely 0 kernel\n",
    "            state_dict[f'conv{conv_i}.weight'][remove] = 1e-5\n",
    "            \n",
    "    return state_dict\n",
    "\n",
    "def prune_network(net, num_kernels, prune_ratio):\n",
    "    \"\"\"\n",
    "    Prune the least important kernels from the model.\n",
    "\n",
    "    Args:\n",
    "        net (nn.Module): The model.\n",
    "        num_kernels (int): Number of kernels in each group.\n",
    "        prune_ratio (list): The proportion of kernels to prune in each group.\n",
    "\n",
    "    Returns:\n",
    "        nn.Module: The pruned network.\n",
    "    \"\"\"\n",
    "    state_dict = net.state_dict()\n",
    "\n",
    "    # Iterate over all layers and prune\n",
    "    for conv_i in range(1, 4):\n",
    "        state_dict = prune(state_dict, conv_i, num_kernels, prune_ratio[conv_i-1])\n",
    "    net.load_state_dict(state_dict)\n",
    "    return net\n",
    "\n",
    "\n",
    "def count_nonzero_weights(state_dict, num_kernels):\n",
    "    \"\"\"\n",
    "    Count the number of non-zero weights in each kernel group of the model.\n",
    "\n",
    "    Args:\n",
    "        state_dict (dict): The state dict of the network.\n",
    "        num_kernels (int): Number of kernels in each layer.\n",
    "\n",
    "    Returns:\n",
    "        list: The number of non-zero weights in each layer.\n",
    "    \"\"\"\n",
    "    zero_weights = [0, 0, 0]\n",
    "\n",
    "    # Iterate over all layers and kernels\n",
    "    for i in range(1, 4):\n",
    "        for j in range(0, num_kernels):\n",
    "            # Count the zero weights\n",
    "            if (state_dict[f'conv{i}.weight'][j] == 1e-5).all(): #this is the value we set the weights to in the prune function, it is arbitrary\n",
    "                zero_weights[i-1] += 1\n",
    "    nonzero_weights = [num_kernels - zero_weights[i] for i in range(3)]\n",
    "\n",
    "    return nonzero_weights\n",
    "\n",
    "\n",
    "# Initiate lists to store results\n",
    "pre_prunes = [] # DICE scores before pruning\n",
    "post_prunes = [] # DICE scores after pruning\n",
    "reductions = [] # Reduction in parameters\n",
    "num_kernels = 128 # Number of kernels in each kernel group\n",
    "\n",
    "# Iterate over models and prune\n",
    "for j in range(folds):\n",
    "    # Load the network\n",
    "    net = torch.load(f\"models/learned_filters_{num_kernels}_{j}.pt\", map_location=device)\n",
    "\n",
    "    print(\"-------Before pruning-------\")\n",
    "    # Test the network before pruning\n",
    "    pre_prune = test_suite(net, verbose=True)\n",
    "    pre_prunes.append(pre_prune)\n",
    "\n",
    "    # Define the pruning ratio for each layer, essentially, this is the proportion of kernel *pairs* for which one pair will be removed\n",
    "    # In other words, if the pruning ratio is 1.0, then *at most* half of all the kernels for that layer will be removed\n",
    "    # However, it is not *guaranteed* that half will be removed, since the similarity ordering can cause some kernels to be the most similar to multiple other kernels\n",
    "    # thus, this kernel could be removed first, leading to the other pairs to have \"already been pruned\".\n",
    "    # This is also why we have to manually compute the number of parameters removed\n",
    "    prune_ratio = [0.35, 0.0, 0.0]\n",
    "    \n",
    "    # Prune the network\n",
    "    net = prune_network(net, num_kernels, prune_ratio)\n",
    "\n",
    "    # Count the number of non-zero weights\n",
    "    nonzero_weights = count_nonzero_weights(net.state_dict(), num_kernels)\n",
    "    new_kernel_num = nonzero_weights\n",
    "\n",
    "    # Compute the new parameter count and the reduction\n",
    "    new_param_count = compute_param_num(new_kernel_num[0], new_kernel_num[1], new_kernel_num[2])\n",
    "    reduction_percentage = (1 - new_param_count/compute_param_num(num_kernels, num_kernels, num_kernels))*100\n",
    "    reductions.append(reduction_percentage)\n",
    "\n",
    "    print(f\"\\nRemoved {reduction_percentage:.2f}% of params\")\n",
    "    print(\"-------After pruning-------\")\n",
    "\n",
    "    # Test the network after pruning\n",
    "    post_prune = test_suite(net, verbose=False)\n",
    "    post_prunes.append(post_prune)\n",
    "\n",
    "    # Print the results\n",
    "    print(f\"DaLiA DICE score: {post_prune[0]:.4f} ({post_prune[0]/pre_prune[0]*100:.2f}% of original)\")\n",
    "    print(f\"TROIKA DICE score: {post_prune[1]:.4f} ({post_prune[1]/pre_prune[1]*100:.2f}% of original)\")\n",
    "    print(f\"WESAD DICE score: {post_prune[2]:.4f} ({post_prune[2]/pre_prune[2]*100:.2f}% of original)\")\n",
    "    print(\"=====================================\")"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}