find_subnets_torch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QWwPhoB4VkHs"
   },
   "source": [
    "# Find RL Agent Subnetworks using Supermasks Techniques\n",
    "\n",
    "---\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O2PtbrFz01L5"
   },
   "source": [
    "## SetUp & Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 139
    },
    "id": "cVH7LTvX00Y5",
    "outputId": "d3706608-c04f-47f7-ce02-4637d455e786"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cloning into 'Preference_Extraction'...\n",
      "remote: Enumerating objects: 201, done.\u001b[K\n",
      "remote: Counting objects: 100% (201/201), done.\u001b[K\n",
      "remote: Compressing objects: 100% (142/142), done.\u001b[K\n",
      "remote: Total 1201 (delta 135), reused 104 (delta 59), pack-reused 1000\u001b[K\n",
      "Receiving objects: 100% (1201/1201), 38.18 MiB | 10.96 MiB/s, done.\n",
      "Resolving deltas: 100% (449/449), done.\n"
     ]
    }
   ],
   "source": [
    "!git clone https://github.com/arunraja-hub/Preference_Extraction.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 364
    },
    "id": "kFZkjgJt4W2H",
    "outputId": "29e8bec3-7d89-47f8-f7f2-f27c879d9d4e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wed Sep 30 07:52:16 2020       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 455.23.05    Driver Version: 418.67       CUDA Version: 10.1     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
      "| N/A   48C    P8    30W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
      "|                               |                      |                 ERR! |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|  No running processes found                                                 |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "# Check what GPU is active\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "Yd-5ibD8Zsfl"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  %tensorflow_version 2.x\n",
    "except Exception:\n",
    "  pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "5xmOhgVIHOwh"
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import argparse\n",
    "import os\n",
    "import math\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn import metrics\n",
    "from sklearn.metrics import roc_curve, auc, roc_auc_score\n",
    "from sklearn.preprocessing import label_binarize\n",
    "import joblib\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision.transforms.functional as TF\n",
    "from torch.optim.lr_scheduler import CosineAnnealingLR\n",
    "from torchvision import datasets, transforms\n",
    "import torch.autograd as autograd\n",
    "from torchsummary import summary\n",
    "\n",
    "from sklearn.utils import shuffle\n",
    "import tensorflow as tf\n",
    "import concurrent.futures\n",
    "import itertools\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "import re\n",
    "import io\n",
    "import itertools\n",
    "import sys\n",
    "\n",
    "sys.path.append('Preference_Extraction')\n",
    "from extractors.data_getter import get_data_from_folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "id": "5hU3eIxdF6IH",
    "outputId": "6a2b291e-fcb4-4436-98b5-c872b0c6a673"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PARAMS = {}\n",
    "np.random.seed(0)\n",
    "torch.manual_seed(0)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "use_cuda = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uIHGI0a71NVI"
   },
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "dAN7OYeaDMk5"
   },
   "outputs": [],
   "source": [
    "# Data source options, choices: 'mnist', 'grid' and 'doom'\n",
    "PARAMS['env'] = 'grid' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 69
    },
    "id": "-XVXv9cXDeL2",
    "outputId": "91538c8a-e53f-499f-8cc9-7afa2fb3d92a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying gs://pref-extr-data/agentv29/experienceData.pkl...\n",
      "- [1 files][  1.3 GiB/  1.3 GiB]   46.8 MiB/s                                   \n",
      "Operation completed over 1 objects/1.3 GiB.                                      \n"
     ]
    }
   ],
   "source": [
    "# Load doom data from GCP storage, run only if env == 'doom'\n",
    "# Requires GCP authentication\n",
    "# Change last line to load different experience data\n",
    "from google.colab import auth\n",
    "auth.authenticate_user()\n",
    "!gsutil cp gs://pref-extr-data/agentv29/experienceData.pkl experienceData.pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "id": "bSOw_1tcDWJt",
    "outputId": "f75f07e5-898a-413d-bd8a-668fb6be64d3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xs (23750, 5, 14, 16) ys (23750,)\n",
      "ys 1 9569\n"
     ]
    }
   ],
   "source": [
    "# Data Loader\n",
    "xs = []\n",
    "ys = []\n",
    "\n",
    "if PARAMS['env'] == 'grid':\n",
    "    all_raw_data = get_data_from_folder(\"data/simple_env_1/\")\n",
    "    for data in all_raw_data:\n",
    "        for i in range(data.observation.shape[0]):\n",
    "            x = np.copy(data.observation[i])\n",
    "            y = (data.policy_info['satisfaction'].as_list()[i] > -6).astype(int)\n",
    "\n",
    "            xs.append(x)\n",
    "            ys.append(y)\n",
    "\n",
    "    ys = np.array(ys).astype(int)\n",
    "    xs = np.rollaxis(np.array(xs), 3, 1) # Torch wants channel-first\n",
    "\n",
    "elif PARAMS['env'] == 'doom':\n",
    "    all_raw_data = joblib.load('/content/experienceData.pkl')\n",
    "    for data in all_raw_data:\n",
    "        for i in range(data.observation.shape[0]):\n",
    "            x = np.copy(data.observation[i])\n",
    "            label_object = data.policy_info['satisfaction'][i]\n",
    "            if len(label_object) == 0: # When label is empty, i.e. human is dead, skip frame\n",
    "                continue\n",
    "            else:\n",
    "                y = label_object['object_angle'] < 90 or label_object['object_angle'] > 270\n",
    "\n",
    "            xs.append(x)\n",
    "            ys.append(y)\n",
    "    \n",
    "    ys = np.array(ys).astype(int)\n",
    "    xs = np.rollaxis(np.array(xs), 3, 1) # Torch wants channel-first\n",
    "\n",
    "elif PARAMS['env'] == 'mnist':\n",
    "    mnist_choice = 3 # Change this value to do binary classification against other mnist digit\n",
    "    (mnist_xs, mnist_ys), _ = tf.keras.datasets.mnist.load_data()\n",
    "    for ix, label in enumerate(mnist_ys):\n",
    "        x = np.copy(mnist_xs[ix])\n",
    "        y = (label == mnist_choice).astype(int)\n",
    "\n",
    "        xs.append(np.array([x]))\n",
    "        ys.append(y)\n",
    "\n",
    "    ys = np.array(ys).astype(int)\n",
    "    xs = np.array(xs) # Mnist has only one channel\n",
    "    del mnist_xs, mnist_ys\n",
    "\n",
    "else:\n",
    "    print('Invalid enviroment choice!!')\n",
    "\n",
    "print(\"xs\", xs.shape, \"ys\", ys.shape)\n",
    "print(\"ys 1\", np.sum(ys))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "id": "JmNMl0RHdBSe",
    "outputId": "57b8efdd-24ca-4a2f-dba2-bd2273f455ec"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xs (3656, 6, 60, 100) ys (3656,)\n",
      "ys 1 1828\n"
     ]
    }
   ],
   "source": [
    "# Rebalancing data to minority class\n",
    "points = xs\n",
    "labels = ys\n",
    "\n",
    "# indexes of 1s and 0s\n",
    "indexes1 = [i for i in range(len(points)) if labels[i] == 1]\n",
    "indexes0 = [i for i in range(len(points)) if labels[i] == 0]\n",
    "\n",
    "# separate 0s and 1s\n",
    "x0, x1, y0, y1 = points[indexes0], points[indexes1], labels[indexes0], labels[indexes1]\n",
    "\n",
    "minority_points, minority_labels = x1, y1  # points and labels for the minority class\n",
    "majority_points, majority_labels = x0, y0  # points and labels for the majority class\n",
    "\n",
    "# get a random permutation of indexes of the majority that includes a number of indexes equal to the minority\n",
    "sample_ind = np.random.permutation(len(majority_labels))[:len(minority_labels)]\n",
    "\n",
    "# subsample the majority\n",
    "majority_points, majority_labels = majority_points[sample_ind], majority_labels[sample_ind]\n",
    "\n",
    "# concat the minority and the sub-sampled majority\n",
    "xs = np.concatenate((majority_points, minority_points))\n",
    "ys = np.concatenate((majority_labels, minority_labels))\n",
    "\n",
    "del points, labels, x0, x1, y0, y1\n",
    "del minority_points, majority_points, minority_labels, majority_labels,\n",
    "\n",
    "print(\"xs\", xs.shape, \"ys\", ys.shape)\n",
    "print(\"ys 1\", np.sum(ys))\n",
    "xs, ys = shuffle(xs, ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y1R6bB1u1B-l"
   },
   "source": [
    "## Nets and Subnets Architectures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "metadata": {
    "id": "fzdNqZBANfjh"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Original code from What's hidden in a randomly weighted neural network? paper\n",
    "    Implemented at https://github.com/allenai/hidden-networks\n",
    "\"\"\"\n",
    "\n",
    "class GetSubnet(autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, scores, k):\n",
    "        # Get the supermask by sorting the scores and using the top k%\n",
    "        out = scores.clone()\n",
    "        _, idx = scores.flatten().sort()\n",
    "        j = int((1 - k) * scores.numel())\n",
    "\n",
    "        # flat_out and out access the same memory.\n",
    "        flat_out = out.flatten()\n",
    "        flat_out[idx[:j]] = 0\n",
    "        flat_out[idx[j:]] = 1\n",
    "\n",
    "        return out\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, g):\n",
    "        # send the gradient g straight-through on the backward pass.\n",
    "        return g, None\n",
    "\n",
    "class SupermaskConv(nn.Conv2d):\n",
    "    def __init__(self, *args, k, scores_init='kaiming_uniform', **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.k = k\n",
    "        self.scores_init = scores_init\n",
    "\n",
    "        # initialize the scores\n",
    "        self.scores = nn.Parameter(torch.Tensor(self.weight.size()))\n",
    "        if self.scores_init == 'kaiming_normal':\n",
    "          nn.init.kaiming_normal_(self.scores)\n",
    "        elif self.scores_init == 'kaiming_uniform':\n",
    "          nn.init.kaiming_uniform_(self.scores, a=math.sqrt(5))\n",
    "        elif self.scores_init == 'xavier_normal':\n",
    "          nn.init.xavier_normal_(self.scores)\n",
    "        elif self.scores_init == 'best_activation':\n",
    "          nn.init.ones_(self.scores)\n",
    "        else:\n",
    "          nn.init.uniform_(self.scores)\n",
    "        \n",
    "        # NOTE: turn the gradient on the weights off\n",
    "        self.weight.requires_grad = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        subnet = GetSubnet.apply(self.scores.abs(), self.k)\n",
    "        w = self.weight * subnet\n",
    "        x = F.conv2d(\n",
    "            x, w, self.bias, self.stride, self.padding, self.dilation, self.groups\n",
    "        )\n",
    "        return x\n",
    "\n",
    "class SupermaskLinear(nn.Linear):\n",
    "    def __init__(self, *args, k, scores_init='kaiming_uniform', **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.k = k\n",
    "        self.scores_init = scores_init\n",
    "\n",
    "        # initialize the scores and weights\n",
    "        self.scores = nn.Parameter(torch.Tensor(self.weight.size()))\n",
    "        if self.scores_init == 'kaiming_normal':\n",
    "          nn.init.kaiming_normal_(self.scores)\n",
    "        elif self.scores_init == 'kaiming_uniform':\n",
    "          nn.init.kaiming_uniform_(self.scores, a=math.sqrt(5))\n",
    "        elif self.scores_init == 'xavier_normal':\n",
    "          nn.init.xavier_normal_(self.scores)\n",
    "        elif self.scores_init == 'best_activation':\n",
    "          nn.init.ones_(self.scores)\n",
    "        else:\n",
    "          nn.init.uniform_(self.scores)\n",
    "\n",
    "        # NOTE: turn the gradient on the weights off\n",
    "        self.weight.requires_grad = False\n",
    "\n",
    "    def forward(self, x):\n",
    "        subnet = GetSubnet.apply(self.scores.abs(), self.k)\n",
    "        w = self.weight * subnet\n",
    "        return F.linear(x, w, self.bias)\n",
    "        return x\n",
    "\n",
    "# NOTE: not used here but we use NON-AFFINE Normalization!\n",
    "# So there is no learned parameters for your nomralization layer.\n",
    "class NonAffineBatchNorm(nn.BatchNorm2d):\n",
    "    def __init__(self, dim):\n",
    "        super(NonAffineBatchNorm, self).__init__(dim, affine=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {
    "id": "MLgKTR1sqhB-"
   },
   "outputs": [],
   "source": [
    "class DQNMaskNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Class for Supermask Networks for DQN Agents\n",
    "    When picking an agent with a different architecture\n",
    "    You need to change the specs of this class as well (e.g. add a layer)\n",
    "    If q_head_index is None, this uses a linear model on the normalized q outputs\n",
    "    otherwise, it gets the Q head with the specified index.\n",
    "    \"\"\" \n",
    "    def __init__(self, fine_tune, k, q_head_index, use_last_linear,\n",
    "                 means_stds=[[0, 0, 0], [1, 1, 1]], init_from_act_index=None):\n",
    "        super(DQNMaskNet, self).__init__()\n",
    "        \n",
    "        if not PARAMS['env'] == 'mnist':\n",
    "            channels_in = 5\n",
    "            flattened_shape = 960\n",
    "        else:\n",
    "            channels_in = 1\n",
    "            flattened_shape = 4608\n",
    "\n",
    "        if fine_tune:\n",
    "            conv_layer = nn.Conv2d\n",
    "            dense_layer = nn.Linear\n",
    "            additional_args = {}\n",
    "            init_from_act_index = None\n",
    "        else:\n",
    "            conv_layer = SupermaskConv\n",
    "            dense_layer = SupermaskLinear\n",
    "            additional_args = {'k': k}\n",
    "            if init_from_act_index is not None:\n",
    "                additional_args['scores_init'] = 'best_activation'\n",
    "        \n",
    "        self.conv1 = conv_layer(in_channels=channels_in, out_channels=16, \n",
    "                                kernel_size=3, stride=1, bias=True, **additional_args)\n",
    "        self.conv2 = conv_layer(in_channels=16, out_channels=32, \n",
    "                                kernel_size=3, stride=2, bias=True, **additional_args)\n",
    "        self.fc1 = dense_layer(in_features=flattened_shape, out_features=64, \n",
    "                               bias=True, **additional_args)\n",
    "        self.fc2 = dense_layer(in_features=64, out_features=3, bias=True, **additional_args)\n",
    "        \n",
    "        if init_from_act_index is not None:\n",
    "            init_scores = np.zeros((3, 64))\n",
    "            init_scores[:, init_from_act_index] = 1.0\n",
    "            self.fc2.scores.data = torch.from_numpy(init_scores).float()\n",
    "\n",
    "        self.fc3 = dense_layer(in_features=3, out_features=1, bias=True, **additional_args)\n",
    "        self.linear = nn.Linear(1, 1, bias=True)\n",
    "\n",
    "        self.qix = q_head_index\n",
    "        self.mu_s = means_stds\n",
    "        self.use_last_linear = use_last_linear\n",
    "\n",
    "    def fwd_conv1(self, x):\n",
    "        x = self.conv1(x)\n",
    "        return F.relu(x)\n",
    "\n",
    "    def fwd_conv2(self, x):\n",
    "        x = self.fwd_conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        return F.relu(x)\n",
    "\n",
    "    def fwd_flat(self, x):\n",
    "        x = self.fwd_conv2(x)\n",
    "        # Pre-flattening transpose is necessary for TF-Torch conversion\n",
    "        return torch.flatten(torch.transpose(x, 1, 3), 1)\n",
    "\n",
    "    def fwd_fc1(self, x):\n",
    "        x = self.fwd_flat(x)\n",
    "        x = self.fc1(x)\n",
    "        return F.relu(x)\n",
    "    \n",
    "    def fwd_fc2(self, x):\n",
    "        x = self.fwd_fc1(x)\n",
    "        return self.fc2(x)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fwd_fc2(x)\n",
    "\n",
    "        x -= torch.tensor(self.mu_s[0], device=device)\n",
    "        x /= torch.tensor(self.mu_s[1], device=device)\n",
    "\n",
    "        if self.qix == None:\n",
    "          x = self.fc3(x)\n",
    "        else:\n",
    "          x = x[: ,self.qix:self.qix+1]\n",
    "\n",
    "        if self.use_last_linear:\n",
    "          x = self.linear(x)\n",
    "\n",
    "        x = torch.sigmoid(x)\n",
    "        return x.flatten()\n",
    "\n",
    "    def layer_to_norm(self, x):\n",
    "        return self.fwd_fc2(x)\n",
    "\n",
    "    def load_weights(self, original_weights):\n",
    "        if not PARAMS['env'] == 'mnist':\n",
    "            self.conv1.weight.data = torch.from_numpy(np.transpose(original_weights[0]))\n",
    "            self.fc1.weight.data = torch.from_numpy(np.transpose(original_weights[4]))\n",
    "        else:\n",
    "            self.conv1.weight.data = torch.from_numpy(np.transpose(original_weights[0][:,:,:1,:]))\n",
    "            mnist_flt_weights = np.random.rand(64, 4608)\n",
    "            mnist_flt_weights[:, :original_weights[4].shape[0]] = np.transpose(original_weights[4])\n",
    "            mnist_flt_weights = mnist_flt_weights.astype(np.float32)\n",
    "            self.fc1.weight.data = torch.from_numpy(mnist_flt_weights)\n",
    "\n",
    "        self.conv1.bias.data = torch.from_numpy(original_weights[1])\n",
    "        self.conv2.weight.data = torch.from_numpy(np.transpose(original_weights[2]))\n",
    "        self.conv2.bias.data = torch.from_numpy(original_weights[3])\n",
    "        self.fc1.bias.data = torch.from_numpy(original_weights[5])\n",
    "        self.fc2.weight.data = torch.from_numpy(np.transpose(original_weights[6]))\n",
    "        self.fc2.bias.data = torch.from_numpy(original_weights[7])\n",
    "        self.fc3.weight.data = torch.from_numpy(np.ones(shape=[1,3], dtype=np.float32))\n",
    "        self.fc3.bias.data = torch.from_numpy(np.zeros(shape=[1], dtype=np.float32))\n",
    "        self.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {
    "id": "22d-UaZtMPMT"
   },
   "outputs": [],
   "source": [
    "class PPOMaskNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Class for Supermask Networks for PPO Agents\n",
    "    When picking an agent with a different architecture\n",
    "    You need to change the specs of this class as well (e.g. add a layer)\n",
    "    \"\"\" \n",
    "    def __init__(self, fine_tune, k, use_last_linear, means_stds=[[0, 0, 0], [1, 1, 1]],\n",
    "                 dropout=None):\n",
    "        super(PPOMaskNet, self).__init__()\n",
    "        \n",
    "        if fine_tune:\n",
    "            conv_layer = nn.Conv2d\n",
    "            dense_layer = nn.Linear\n",
    "            additional_args = {}\n",
    "            init_from_act_index = None\n",
    "        else:\n",
    "            conv_layer = SupermaskConv\n",
    "            dense_layer = SupermaskLinear\n",
    "            additional_args = {'k': k}\n",
    "        \n",
    "        self.fc1 = dense_layer(in_features=36000, out_features=200, bias=True, **additional_args)\n",
    "        self.fc2 = dense_layer(in_features=200, out_features=100, bias=True, **additional_args)\n",
    "        #self.fc3 = dense_layer(in_features=100, out_features=4, bias=True, **additional_args)\n",
    "        self.drop = None\n",
    "        if dropout is not None:\n",
    "            self.drop = nn.Dropout(p=dropout)\n",
    "        self.linear = nn.Linear(100, 1, bias=True)\n",
    "        self.use_last_linear = use_last_linear\n",
    "\n",
    "        self.mu_s = means_stds\n",
    "\n",
    "    def fwd_flat(self, x):\n",
    "        # Pre-flattening transpose is necessary for TF-Torch conversion\n",
    "        return torch.flatten(torch.transpose(x, 1, 3), 1)\n",
    "\n",
    "    def fwd_fc1(self, x):\n",
    "        x = self.fwd_flat(x)\n",
    "        x = self.fc1(x)\n",
    "        return F.relu(x)\n",
    "    \n",
    "    def fwd_fc2(self, x):\n",
    "        x = self.fwd_fc1(x)\n",
    "        x = self.fc2(x)\n",
    "        return F.relu(x)\n",
    "\n",
    "    def fwd_fc3(self, x):\n",
    "        x = self.fwd_fc2(x)\n",
    "        return self.fc3(x)\n",
    "\n",
    "    def fwd_fc4(self, x):\n",
    "        x = self.fwd_fc3(x)\n",
    "        return self.fc4(x)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fwd_fc2(x)\n",
    "\n",
    "        x -= torch.tensor(self.mu_s[0], device=device)\n",
    "        x /= torch.tensor(self.mu_s[1], device=device)\n",
    "\n",
    "        if self.use_last_linear:\n",
    "            if self.drop is not None:\n",
    "                x = self.drop(x)\n",
    "            x = self.linear(x)\n",
    "\n",
    "        x = torch.sigmoid(x)\n",
    "        return x.flatten()\n",
    "\n",
    "    def layer_to_norm(self, x):\n",
    "        # Change the function below when you want to do activation\n",
    "        # normalizaiton on a different layer of the agent\n",
    "        return self.fwd_fc2(x)\n",
    "\n",
    "    def load_weights(self, original_weights):\n",
    "        self.fc1.weight.data = torch.from_numpy(np.transpose(original_weights[0]))\n",
    "        self.fc1.bias.data = torch.from_numpy(original_weights[1])\n",
    "        self.fc2.weight.data = torch.from_numpy(np.transpose(original_weights[2]))\n",
    "        self.fc2.bias.data = torch.from_numpy(original_weights[3])\n",
    "        #self.fc3.weight.data = torch.from_numpy(np.transpose(original_weights[4]))\n",
    "        #self.fc3.bias.data = torch.from_numpy(original_weights[5])\n",
    "        self.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "metadata": {
    "id": "epVaY6QPLyFe"
   },
   "outputs": [],
   "source": [
    "class RandomMaskNet(nn.Module):\n",
    "    \"\"\"\n",
    "    Big Randomly Weighted Neural Net\n",
    "    Used as a baseline for supermask technique\n",
    "    Copies original Net class from paper\n",
    "    See https://github.com/allenai/hidden-networks/blob/master/simple_mnist_example.py\n",
    "    \"\"\" \n",
    "    def __init__(self, k, dropout):\n",
    "        super(RandomMaskNet, self).__init__()\n",
    "        if PARAMS['env'] == 'doom':\n",
    "            channels_in = 6\n",
    "            flattened_shape = 86016\n",
    "        elif PARAMS['env'] == 'grid':\n",
    "            channels_in = 5\n",
    "            flattened_shape = 960\n",
    "        elif PARAMS['env'] == 'mnist':\n",
    "            channels_in = 1\n",
    "            flattened_shape = 4608\n",
    "\n",
    "        self.conv1 = SupermaskConv(in_channels=channels_in, out_channels=32, \n",
    "                                   kernel_size=3, stride=1, bias=False, k=k)\n",
    "        self.conv2 = SupermaskConv(in_channels=32, out_channels=64, \n",
    "                                   kernel_size=3, stride=1, bias=False, k=k)\n",
    "        self.dropout1 = nn.Dropout2d(dropout)\n",
    "        self.dropout2 = nn.Dropout2d(dropout)\n",
    "        self.fc1 = SupermaskLinear(in_features=flattened_shape, out_features=128, bias=False, k=k)\n",
    "        self.fc2 = SupermaskLinear(in_features=128, out_features=1, bias=False, k=k)\n",
    "\n",
    "        nn.init.kaiming_normal_(self.conv1.weight, mode=\"fan_in\", nonlinearity=\"relu\")\n",
    "        nn.init.kaiming_normal_(self.conv2.weight, mode=\"fan_in\", nonlinearity=\"relu\")\n",
    "        nn.init.kaiming_normal_(self.fc1.weight, mode=\"fan_in\", nonlinearity=\"relu\")\n",
    "        nn.init.kaiming_normal_(self.fc2.weight, mode=\"fan_in\", nonlinearity=\"relu\")\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.max_pool2d(x, 2)\n",
    "        x = self.dropout1(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        x = torch.sigmoid(x)\n",
    "        return x.flatten()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wpSHf83VSgoJ"
   },
   "source": [
    "## Loading Weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 69
    },
    "id": "9Bmgjet_cJ7m",
    "outputId": "a70ce8ce-2972-4edb-a9f8-bafc749acb4f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying gs://pref-extr-data/agentv29/actorNet.keras...\n",
      "\\ [1 files][ 27.6 MiB/ 27.6 MiB]                                                \n",
      "Operation completed over 1 objects/27.6 MiB.                                     \n"
     ]
    }
   ],
   "source": [
    "!gsutil cp gs://pref-extr-data/agentv29/actorNet.keras actorNet.keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "DD-16xFHSi5r"
   },
   "outputs": [],
   "source": [
    "# Shortlisting different agents\n",
    "dqn_grid_path = \"Preference_Extraction/saved_model2\"\n",
    "ppo_doom_path = \"/content/actorNet.keras\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 186,
   "metadata": {
    "id": "DduBU77wPgSI"
   },
   "outputs": [],
   "source": [
    "# Choose agent to use (ppo, dqn or random)\n",
    "PARAMS['agent_type'] = 'ppo'\n",
    "\n",
    "if PARAMS['agent_type'] == 'ppo':\n",
    "    model_path = ppo_doom_path\n",
    "elif PARAMS['agent_type'] == 'dqn':\n",
    "    model_path = dqn_grid_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 187,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 312
    },
    "id": "ml_EV0Y_KrHR",
    "outputId": "26300241-fe21-4c72-f1b3-99c54ae442c4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.\n",
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "flatten (Flatten)            (1, 36000)                0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (1, 200)                  7200200   \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (1, 100)                  20100     \n",
      "_________________________________________________________________\n",
      "logits (Dense)               (1, 4)                    404       \n",
      "=================================================================\n",
      "Total params: 7,220,704\n",
      "Trainable params: 7,220,704\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "restored_model = tf.keras.models.load_model(model_path)\n",
    "restored_model.summary()\n",
    "original_weights = restored_model.get_weights()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "id": "AwWx3ioNqhhC"
   },
   "outputs": [],
   "source": [
    "if PARAMS['agent_type'] == 'dqn':\n",
    "    test_model = DQNMaskNet(k=1, fine_tune=False, q_head_index=None,\n",
    "                            use_last_linear=True).to(device)\n",
    "elif PARAMS['agent_type'] == 'ppo':\n",
    "    test_model = PPOMaskNet(k=1, fine_tune=False, use_last_linear=True).to(device)\n",
    "\n",
    "test_model.load_weights(original_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {
    "id": "_i1akicOSL3n"
   },
   "outputs": [],
   "source": [
    "# Testing that weigths loaded property\n",
    "# Change this list to test different torch models\n",
    "torch_model_layers = [test_model.fwd_conv1,\n",
    "                    test_model.fwd_conv2,\n",
    "                    test_model.fwd_fc1, \n",
    "                    test_model.fwd_fc2]\n",
    "                    #test_model.fwd_fc3]\n",
    "\n",
    "# Does not work of mnist because original agents don't have the shape!\n",
    "if not PARAMS['env'] == 'mnist':\n",
    "\n",
    "    def check_same(torch_layer, tf_layer):\n",
    "        torch_out = np.transpose(torch_layer(single_observation_torch).detach().cpu().numpy())\n",
    "        torch_out = torch_out.reshape(torch_out.shape[:-1])\n",
    "        tf_out = tf_layer(np.array([np.transpose(single_observation)]))[0].numpy()\n",
    "        np.testing.assert_allclose(torch_out, tf_out, rtol=.1, atol=5)  \n",
    "\n",
    "    for i in range(100):\n",
    "        single_observation = xs[i]\n",
    "        single_observation_torch = torch.Tensor(np.array([xs[i]]))\n",
    "        single_observation_torch = single_observation_torch.to(device)\n",
    "        \n",
    "        index_shift = 0\n",
    "        for ix, original_lyr in enumerate(restored_model.layers):\n",
    "            if ix < len(torch_model_layers):\n",
    "                if original_lyr.name == 'flatten':\n",
    "                    index_shift = 1\n",
    "                else:\n",
    "                    tf_sub_model = tf.keras.models.Model(inputs=restored_model.input, outputs=original_lyr.output)\n",
    "                    check_same(torch_model_layers[ix-index_shift], tf_sub_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cjzhnnt9UqBB"
   },
   "source": [
    "### Data getters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {
    "id": "fv8Z2oPJb270"
   },
   "outputs": [],
   "source": [
    "PARAMS['num_train'] = 500\n",
    "PARAMS['num_val'] = 500\n",
    "PARAMS['batch_size'] = 128\n",
    "\n",
    "def get_data_sample(xs=None, ys=None):\n",
    "    xs, ys = shuffle(xs, ys)\n",
    "    \n",
    "    train_split = PARAMS['num_train']\n",
    "    test_split = PARAMS['num_train']+PARAMS['num_val']\n",
    "\n",
    "    tr_data_loader = torch.utils.data.DataLoader(\n",
    "        torch.utils.data.TensorDataset(torch.Tensor(xs[:train_split]), torch.Tensor(ys[:train_split])),\n",
    "        batch_size=PARAMS['batch_size'])\n",
    "    \n",
    "    val_data_loader = torch.utils.data.DataLoader(\n",
    "        torch.utils.data.TensorDataset(torch.Tensor(xs[train_split:test_split]), torch.Tensor(ys[train_split:test_split])),\n",
    "        batch_size=PARAMS['batch_size'])\n",
    "    \n",
    "    return tr_data_loader, val_data_loader, xs[:train_split]\n",
    "\n",
    "def get_heads_mu_and_sigma(model, obs):\n",
    "    \n",
    "    model.eval()\n",
    "    obs_tensor = torch.Tensor(obs)\n",
    "    obs_tensor = obs_tensor.to(device)\n",
    "    heads_values = model.layer_to_norm(obs_tensor).detach().cpu().numpy()\n",
    "\n",
    "    mu = heads_values.mean(axis=0)\n",
    "    s = heads_values.std(axis=0)\n",
    "    s[s == 0] = 1\n",
    "    \n",
    "    return np.array([mu, s])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Th-WxNDGTKth"
   },
   "source": [
    "### Single run train/test methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "metadata": {
    "id": "xEj7Qn1-H-Es"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "    Train/Test function for Randomly Weighted Hidden Neural Networks Techniques\n",
    "    Adapted from https://github.com/NesterukSergey/hidden-networks/blob/master/demos/mnist.ipynb\n",
    "\"\"\"\n",
    "\n",
    "def compute_metrics(predictions, true_labels):\n",
    "    predictions = np.array(predictions)\n",
    "    true_labels = np.array(true_labels)\n",
    "    accuracy = np.sum(np.equal((predictions > 0.5).astype(int), true_labels)) / len(true_labels)\n",
    "    fpr, tpr, thresholds = metrics.roc_curve(true_labels, predictions)\n",
    "    auc = metrics.auc(fpr, tpr)\n",
    "    return accuracy, auc\n",
    "\n",
    "def train(model, device, train_loader, optimizer, criterion):\n",
    "    \n",
    "    train_loss = 0\n",
    "    true_labels = []\n",
    "    predictions = [] # labels\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    for data, target in itertools.islice(train_loader, PARAMS['num_train']):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = criterion(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss\n",
    "        predictions.extend(output.detach().cpu().numpy())\n",
    "        true_labels.extend(target.detach().cpu().numpy())\n",
    "    \n",
    "    train_loss /= len(train_loader.dataset)\n",
    "    accuracy, auc = compute_metrics(predictions, true_labels)\n",
    "\n",
    "    return train_loss.item(), accuracy, auc\n",
    "\n",
    "def test(model, device, test_loader, criterion):\n",
    "    true_labels = []\n",
    "    predictions = [] # labels\n",
    "\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in itertools.islice(test_loader, PARAMS['num_val']):\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += criterion(output, target)\n",
    "            predictions.extend(output.detach().cpu().numpy())\n",
    "            true_labels.extend(target.detach().cpu().numpy())\n",
    "    \n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    accuracy, auc = compute_metrics(predictions, true_labels)\n",
    "\n",
    "    return test_loss.item(), accuracy, auc\n",
    "\n",
    "def run_model(model, learning_rate, weight_decay, num_epochs):\n",
    "\n",
    "  tr_data_loader, val_data_loader, x_train = get_data_sample(xs, ys)\n",
    "\n",
    "  # Normalise last layer using training data\n",
    "  if hasattr(model, 'layer_to_norm'):\n",
    "      model.mu_s = get_heads_mu_and_sigma(model, x_train)\n",
    "  \n",
    "  optimizer = optim.Adam(\n",
    "      [p for p in model.parameters() if p.requires_grad],\n",
    "      lr=learning_rate,\n",
    "      weight_decay=weight_decay\n",
    "  )\n",
    "\n",
    "  criterion = nn.BCELoss().to(device)\n",
    "  scheduler = CosineAnnealingLR(optimizer, T_max=len(tr_data_loader))\n",
    "\n",
    "  train_losses = []\n",
    "  test_losses = []\n",
    "  train_accs = []\n",
    "  train_aucs = []\n",
    "  test_accs = []\n",
    "  test_aucs = []  \n",
    "\n",
    "  best_test_loss = np.inf\n",
    "  test_loss_up_since = 0\n",
    "  early_stop = 100\n",
    "  verbose = False\n",
    "  for epoch in range(num_epochs):\n",
    "      train_loss, train_accuracy, train_auc = train(model, device, tr_data_loader, optimizer, criterion)\n",
    "      test_loss, test_accuracy, test_auc = test(model, device, val_data_loader, criterion)\n",
    "      scheduler.step()\n",
    "      if test_loss < best_test_loss:\n",
    "          best_test_loss = test_loss\n",
    "          test_loss_up_since = 0\n",
    "      test_loss_up_since += 1\n",
    "      if test_loss_up_since > early_stop:\n",
    "          print('Epoch - ', epoch, 'Early stopping')\n",
    "          break\n",
    "      if verbose:\n",
    "          print('Epoch - ', epoch)\n",
    "          print('Train metrics: loss', train_loss, 'accuracy', train_accuracy, 'auc', train_auc)\n",
    "          print('Val metrics: loss', test_loss, 'accuracy', test_accuracy, 'auc', test_auc)\n",
    "\n",
    "      train_losses.append(train_loss)\n",
    "      test_losses.append(test_loss)\n",
    "      train_accs.append(train_accuracy)\n",
    "      train_aucs.append(train_auc)\n",
    "      test_accs.append(test_accuracy)\n",
    "      test_aucs.append(test_auc)\n",
    "\n",
    "  return {'trainLoss': train_losses, 'testLoss': test_losses, \n",
    "          'trainAccuracy': train_accs, 'testAccuracy': test_accs,\n",
    "          'trainAUC': train_aucs, 'testAUC': test_aucs}\n",
    "\n",
    "def plot_metric(results_dict, metric):\n",
    "    plt.title(metric)\n",
    "    plt.xlabel('Epochs')\n",
    "    plt.plot(list(range(1, len(results_dict[f'train{metric}']) + 1)), results_dict[f'train{metric}'], label=f'Train {metric}')\n",
    "    plt.plot(list(range(1, len(results_dict[f'train{metric}']) + 1)), results_dict[f'test{metric}'], label=f'Test {metric}')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "def multi_runs(model, learning_rate, weight_decay, plots=False):\n",
    "\n",
    "    averaged_results = {}    \n",
    "    for run_ix in range(PARAMS['num_run']):\n",
    " \n",
    "        results = run_model(model, learning_rate=learning_rate, weight_decay=weight_decay, num_epochs=PARAMS['num_epochs'])       \n",
    "        print(f'Train pass no. {run_ix+1}')\n",
    "\n",
    "        if (run_ix == 0):\n",
    "            print(model) \n",
    "            if plots:\n",
    "                print('Debug charts for first training run')\n",
    "                plot_metric(results, 'Loss')\n",
    "                plot_metric(results, 'Accuracy')\n",
    "                plot_metric(results, 'AUC')\n",
    "\n",
    "        for res in results:\n",
    "            if len(results[res]) > 0:\n",
    "                if res not in averaged_results:\n",
    "                    averaged_results[res] = [results[res][-1]]\n",
    "                else:\n",
    "                    averaged_results[res].append(results[res][-1])         \n",
    "    \n",
    "    return {x: sum(averaged_results[x]) / PARAMS['num_run'] for x in averaged_results}\n",
    "\n",
    "def multi_agent_train(hparams):\n",
    "    \n",
    "    load_weigths = True\n",
    "    if PARAMS['agent_type'] == 'ppo':\n",
    "        model = PPOMaskNet(\n",
    "            k=hparams['k'], \n",
    "            fine_tune=hparams['fine_tune'],\n",
    "            means_stds=None,\n",
    "            dropout=hparams['dropout'],\n",
    "            use_last_linear=hparams['use_last_linear']).to(device)\n",
    "    \n",
    "    elif PARAMS['agent_type'] == 'dqn':\n",
    "        model = DQNMaskNet(\n",
    "            k=hparams['k'], \n",
    "            fine_tune=hparams['fine_tune'],\n",
    "            q_head_index=hparams['q_head_index'],\n",
    "            means_stds=None,\n",
    "            use_last_linear=hparams['use_last_linear']).to(device)\n",
    "\n",
    "    else:\n",
    "        model = RandomMaskNet(\n",
    "            k=hparams['k'],\n",
    "            dropout=hparams['dropout'],).to(device)\n",
    "        load_weigths = False\n",
    "    \n",
    "    if load_weigths:\n",
    "       model.load_weights(original_weights)\n",
    "\n",
    "    return multi_runs(model, learning_rate=hparams['learning_rate'], \n",
    "                      weight_decay=hparams['weight_decay'], plots=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 184,
   "metadata": {
    "id": "LvnzysBFGrrj"
   },
   "outputs": [],
   "source": [
    "# Set up hyperparamets for model training\n",
    "PARAMS['num_run'] = 1\n",
    "PARAMS['num_epochs'] = 100\n",
    "\n",
    "all_hparam_possibilities = [\n",
    "   {\n",
    "    \"dropout\": [0.2],\n",
    "    \"k\": [0.1],  \n",
    "    \"fine_tune\": [False],\n",
    "    \"use_last_linear\": [True],\n",
    "    \"learning_rate\": [1e-4], \n",
    "    \"weight_decay\": [0],\n",
    "    \"q_head_index\": [None]\n",
    "   }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 188,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "0PQ3G6oeGm1i",
    "outputId": "5aaaef0e-f788-44fa-ae15-d4d8f60aa900"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PARAMS {'env': 'doom', 'num_train': 500, 'num_val': 500, 'batch_size': 128, 'num_run': 1, 'num_epochs': 100, 'agent_type': 'ppo', 'best_hparams': {'dropout': 0.2, 'k': 0.1, 'fine_tune': False, 'use_last_linear': True, 'learning_rate': 0.0001, 'weight_decay': 0, 'q_head_index': None}}\n",
      "len(hparam_combinations) 1 hparam_combinations [{'dropout': 0.2, 'k': 0.1, 'fine_tune': False, 'use_last_linear': True, 'learning_rate': 0.0001, 'weight_decay': 0, 'q_head_index': None}]\n",
      "hparams {'dropout': 0.2, 'k': 0.1, 'fine_tune': False, 'use_last_linear': True, 'learning_rate': 0.0001, 'weight_decay': 0, 'q_head_index': None}\n",
      "Epoch -  0\n",
      "Train metrics: loss 0.012701721861958504 accuracy 0.508 auc 0.4779825908858166\n",
      "Val metrics: loss 0.005908384453505278 accuracy 0.508 auc 0.48295854955474404\n",
      "Epoch -  1\n",
      "Train metrics: loss 0.006377368234097958 accuracy 0.494 auc 0.48345494111623144\n",
      "Val metrics: loss 0.00567703926935792 accuracy 0.484 auc 0.5202447306041387\n",
      "Epoch -  2\n",
      "Train metrics: loss 0.006441389676183462 accuracy 0.434 auc 0.4406362007168459\n",
      "Val metrics: loss 0.005627544596791267 accuracy 0.488 auc 0.49119097956307267\n",
      "Epoch -  3\n",
      "Train metrics: loss 0.006069375202059746 accuracy 0.518 auc 0.5098566308243728\n",
      "Val metrics: loss 0.005608385894447565 accuracy 0.484 auc 0.4966846050355564\n",
      "Epoch -  4\n",
      "Train metrics: loss 0.006140026263892651 accuracy 0.478 auc 0.47726254480286734\n",
      "Val metrics: loss 0.005608385894447565 accuracy 0.484 auc 0.4966846050355564\n",
      "Epoch -  5\n",
      "Train metrics: loss 0.006320706102997065 accuracy 0.496 auc 0.44998079877112135\n",
      "Val metrics: loss 0.005597704090178013 accuracy 0.514 auc 0.5198202959830867\n",
      "Epoch -  6\n",
      "Train metrics: loss 0.006229417864233255 accuracy 0.514 auc 0.5171050947260625\n",
      "Val metrics: loss 0.005550191272050142 accuracy 0.662 auc 0.6578896790313281\n",
      "Epoch -  7\n",
      "Train metrics: loss 0.005952253472059965 accuracy 0.52 auc 0.5290738607270866\n",
      "Val metrics: loss 0.005557605996727943 accuracy 0.534 auc 0.5853994490358128\n",
      "Epoch -  8\n",
      "Train metrics: loss 0.006137315649539232 accuracy 0.502 auc 0.49427163338453656\n",
      "Val metrics: loss 0.005900253541767597 accuracy 0.488 auc 0.5575309116535332\n",
      "Epoch -  9\n",
      "Train metrics: loss 0.005748345050960779 accuracy 0.504 auc 0.5361143113159242\n",
      "Val metrics: loss 0.005627711769193411 accuracy 0.516 auc 0.5899481068614261\n",
      "Epoch -  10\n",
      "Train metrics: loss 0.005776767618954182 accuracy 0.52 auc 0.5307059651817717\n",
      "Val metrics: loss 0.005556228570640087 accuracy 0.524 auc 0.6270741239028765\n",
      "Epoch -  11\n",
      "Train metrics: loss 0.005814075004309416 accuracy 0.502 auc 0.49331157194060415\n",
      "Val metrics: loss 0.005517211277037859 accuracy 0.54 auc 0.6308539944903581\n",
      "Epoch -  12\n",
      "Train metrics: loss 0.005608655977994204 accuracy 0.522 auc 0.5395545314900154\n",
      "Val metrics: loss 0.005517211277037859 accuracy 0.54 auc 0.6308539944903581\n",
      "Epoch -  13\n",
      "Train metrics: loss 0.005717067513614893 accuracy 0.506 auc 0.5066884280593957\n",
      "Val metrics: loss 0.00558756897225976 accuracy 0.502 auc 0.6229819334999038\n",
      "Epoch -  14\n",
      "Train metrics: loss 0.005732696037739515 accuracy 0.532 auc 0.5645961341525858\n",
      "Val metrics: loss 0.005755905527621508 accuracy 0.48 auc 0.5690467038247166\n",
      "Epoch -  15\n",
      "Train metrics: loss 0.0055700368247926235 accuracy 0.552 auc 0.5811091909882232\n",
      "Val metrics: loss 0.005731650628149509 accuracy 0.516 auc 0.6067493112947658\n",
      "Epoch -  16\n",
      "Train metrics: loss 0.005845072213560343 accuracy 0.52 auc 0.4968157962109575\n",
      "Val metrics: loss 0.005534779280424118 accuracy 0.53 auc 0.6044909987827535\n",
      "Epoch -  17\n",
      "Train metrics: loss 0.006067382637411356 accuracy 0.528 auc 0.5699084741423451\n",
      "Val metrics: loss 0.006192658096551895 accuracy 0.494 auc 0.5710968031264014\n",
      "Epoch -  18\n",
      "Train metrics: loss 0.005832807160913944 accuracy 0.554 auc 0.5745007680491552\n",
      "Val metrics: loss 0.005498514976352453 accuracy 0.544 auc 0.6135722980331859\n",
      "Epoch -  19\n",
      "Train metrics: loss 0.005618811119347811 accuracy 0.502 auc 0.5339381720430109\n",
      "Val metrics: loss 0.005480929743498564 accuracy 0.546 auc 0.6187295790889871\n",
      "Epoch -  20\n",
      "Train metrics: loss 0.005607046186923981 accuracy 0.492 auc 0.5313940092165899\n",
      "Val metrics: loss 0.005480929743498564 accuracy 0.546 auc 0.6187295790889871\n",
      "Epoch -  21\n",
      "Train metrics: loss 0.005556359421461821 accuracy 0.546 auc 0.5609799027137737\n",
      "Val metrics: loss 0.005477254744619131 accuracy 0.538 auc 0.618297136267538\n",
      "Epoch -  22\n",
      "Train metrics: loss 0.005620752461254597 accuracy 0.554 auc 0.5409786226318485\n",
      "Val metrics: loss 0.0054742321372032166 accuracy 0.552 auc 0.6215324492280094\n",
      "Epoch -  23\n",
      "Train metrics: loss 0.005606944672763348 accuracy 0.53 auc 0.5445308499743985\n",
      "Val metrics: loss 0.0055177584290504456 accuracy 0.544 auc 0.6161509385610866\n",
      "Epoch -  24\n",
      "Train metrics: loss 0.005591595079749823 accuracy 0.56 auc 0.5558755760368663\n",
      "Val metrics: loss 0.005522614344954491 accuracy 0.544 auc 0.60690947530271\n",
      "Epoch -  25\n",
      "Train metrics: loss 0.00554589694365859 accuracy 0.542 auc 0.5635240655401946\n",
      "Val metrics: loss 0.00549625139683485 accuracy 0.552 auc 0.6160548401563202\n",
      "Epoch -  26\n",
      "Train metrics: loss 0.0055749686434865 accuracy 0.53 auc 0.547603046594982\n",
      "Val metrics: loss 0.0054771751165390015 accuracy 0.552 auc 0.6251681722083413\n",
      "Epoch -  27\n",
      "Train metrics: loss 0.005521927494555712 accuracy 0.578 auc 0.5852694572452637\n",
      "Val metrics: loss 0.005466926377266645 accuracy 0.544 auc 0.6272503043116151\n",
      "Epoch -  28\n",
      "Train metrics: loss 0.005485245026648045 accuracy 0.546 auc 0.5763408858166922\n",
      "Val metrics: loss 0.005466926377266645 accuracy 0.544 auc 0.6272503043116151\n",
      "Epoch -  29\n",
      "Train metrics: loss 0.005474342964589596 accuracy 0.55 auc 0.582789298515105\n",
      "Val metrics: loss 0.005456145387142897 accuracy 0.548 auc 0.6300371580498431\n",
      "Epoch -  30\n",
      "Train metrics: loss 0.005550154019147158 accuracy 0.536 auc 0.5514272913466461\n",
      "Val metrics: loss 0.005441828165203333 accuracy 0.552 auc 0.6326478313793324\n",
      "Epoch -  31\n",
      "Train metrics: loss 0.0054379659704864025 accuracy 0.588 auc 0.5985183051715309\n",
      "Val metrics: loss 0.005425231531262398 accuracy 0.574 auc 0.636075341149337\n",
      "Epoch -  32\n",
      "Train metrics: loss 0.005433543585240841 accuracy 0.538 auc 0.5818932411674347\n",
      "Val metrics: loss 0.0054307784885168076 accuracy 0.544 auc 0.6416170158242039\n",
      "Epoch -  33\n",
      "Train metrics: loss 0.005514631979167461 accuracy 0.542 auc 0.5751728110599079\n",
      "Val metrics: loss 0.005454302299767733 accuracy 0.544 auc 0.6382535716573772\n",
      "Epoch -  34\n",
      "Train metrics: loss 0.005410149227827787 accuracy 0.584 auc 0.6051747311827957\n",
      "Val metrics: loss 0.005413522943854332 accuracy 0.552 auc 0.6369882759946185\n",
      "Epoch -  35\n",
      "Train metrics: loss 0.005340244621038437 accuracy 0.586 auc 0.624647977470558\n",
      "Val metrics: loss 0.005407344549894333 accuracy 0.556 auc 0.6395669165225191\n",
      "Epoch -  36\n",
      "Train metrics: loss 0.005513249430805445 accuracy 0.55 auc 0.5775089605734767\n",
      "Val metrics: loss 0.005407344549894333 accuracy 0.556 auc 0.6395669165225191\n",
      "Epoch -  37\n",
      "Train metrics: loss 0.005447436589747667 accuracy 0.554 auc 0.5815252176139274\n",
      "Val metrics: loss 0.0054007540456950665 accuracy 0.562 auc 0.6427862130821962\n",
      "Epoch -  38\n",
      "Train metrics: loss 0.005373360123485327 accuracy 0.586 auc 0.6204557091653866\n",
      "Val metrics: loss 0.005394937936216593 accuracy 0.564 auc 0.645492984816452\n",
      "Epoch -  39\n",
      "Train metrics: loss 0.00541610736399889 accuracy 0.592 auc 0.6051907322068613\n",
      "Val metrics: loss 0.005404097493737936 accuracy 0.56 auc 0.6439714267409827\n",
      "Epoch -  40\n",
      "Train metrics: loss 0.005387199576944113 accuracy 0.566 auc 0.6207757296466974\n",
      "Val metrics: loss 0.005442459136247635 accuracy 0.56 auc 0.6411205073995772\n",
      "Epoch -  41\n",
      "Train metrics: loss 0.0054199183359742165 accuracy 0.554 auc 0.6007744495647722\n",
      "Val metrics: loss 0.0053429617546498775 accuracy 0.586 auc 0.6620859760394644\n",
      "Epoch -  42\n",
      "Train metrics: loss 0.005353693850338459 accuracy 0.582 auc 0.6201036866359446\n",
      "Val metrics: loss 0.005336219910532236 accuracy 0.592 auc 0.6648728297776925\n",
      "Epoch -  43\n",
      "Train metrics: loss 0.005251187365502119 accuracy 0.6 auc 0.6535138248847927\n",
      "Val metrics: loss 0.005338545423001051 accuracy 0.564 auc 0.6640079441347941\n",
      "Epoch -  44\n",
      "Train metrics: loss 0.005231444723904133 accuracy 0.612 auc 0.66084229390681\n",
      "Val metrics: loss 0.005338545423001051 accuracy 0.564 auc 0.6640079441347941\n",
      "Epoch -  45\n",
      "Train metrics: loss 0.005356925074011087 accuracy 0.568 auc 0.624375960061444\n",
      "Val metrics: loss 0.005350570194423199 accuracy 0.572 auc 0.661605484015632\n",
      "Epoch -  46\n",
      "Train metrics: loss 0.005390211474150419 accuracy 0.594 auc 0.6193676395289298\n",
      "Val metrics: loss 0.00545898312702775 accuracy 0.566 auc 0.6521878403485168\n",
      "Epoch -  47\n",
      "Train metrics: loss 0.005439892411231995 accuracy 0.604 auc 0.60767089093702\n",
      "Val metrics: loss 0.005310393404215574 accuracy 0.6 auc 0.6741463258376578\n",
      "Epoch -  48\n",
      "Train metrics: loss 0.0053424169309437275 accuracy 0.584 auc 0.6280081925243216\n",
      "Val metrics: loss 0.005282344296574593 accuracy 0.61 auc 0.68391633032225\n",
      "Epoch -  49\n",
      "Train metrics: loss 0.005247656721621752 accuracy 0.594 auc 0.6499615975422427\n",
      "Val metrics: loss 0.005289341323077679 accuracy 0.586 auc 0.6863668396437952\n",
      "Epoch -  50\n",
      "Train metrics: loss 0.00525208143517375 accuracy 0.624 auc 0.6560099846390169\n",
      "Val metrics: loss 0.005287186708301306 accuracy 0.586 auc 0.6856300852072522\n",
      "Epoch -  51\n",
      "Train metrics: loss 0.005315887741744518 accuracy 0.614 auc 0.6400729646697388\n",
      "Val metrics: loss 0.005286556202918291 accuracy 0.596 auc 0.6846691011595875\n",
      "Epoch -  52\n",
      "Train metrics: loss 0.005354379769414663 accuracy 0.58 auc 0.6205997183819765\n",
      "Val metrics: loss 0.005286556202918291 accuracy 0.596 auc 0.6846691011595875\n",
      "Epoch -  53\n",
      "Train metrics: loss 0.005284537561237812 accuracy 0.598 auc 0.6490975422427037\n",
      "Val metrics: loss 0.005291712004691362 accuracy 0.588 auc 0.6841245435325773\n",
      "Epoch -  54\n",
      "Train metrics: loss 0.005279864650219679 accuracy 0.616 auc 0.6516417050691244\n",
      "Val metrics: loss 0.005281850229948759 accuracy 0.6 auc 0.687904414120059\n",
      "Epoch -  55\n",
      "Train metrics: loss 0.0052210246212780476 accuracy 0.614 auc 0.6607302867383513\n",
      "Val metrics: loss 0.005261273588985205 accuracy 0.616 auc 0.6941187776282913\n",
      "Epoch -  56\n",
      "Train metrics: loss 0.005214662291109562 accuracy 0.622 auc 0.6655945980542756\n",
      "Val metrics: loss 0.005254644434899092 accuracy 0.612 auc 0.6923890063424947\n",
      "Epoch -  57\n",
      "Train metrics: loss 0.005228282418102026 accuracy 0.604 auc 0.6544738863287249\n",
      "Val metrics: loss 0.005239460617303848 accuracy 0.62 auc 0.7008136331603563\n",
      "Epoch -  58\n",
      "Train metrics: loss 0.005319183226674795 accuracy 0.598 auc 0.6366327444956478\n",
      "Val metrics: loss 0.005231957882642746 accuracy 0.616 auc 0.6997725671087194\n",
      "Epoch -  59\n",
      "Train metrics: loss 0.005107528530061245 accuracy 0.65 auc 0.6924283154121864\n",
      "Val metrics: loss 0.0052283755503594875 accuracy 0.614 auc 0.7026555192517139\n",
      "Epoch -  60\n",
      "Train metrics: loss 0.005118419416248798 accuracy 0.642 auc 0.697100614439324\n",
      "Val metrics: loss 0.0052283755503594875 accuracy 0.614 auc 0.7026555192517139\n",
      "Epoch -  61\n",
      "Train metrics: loss 0.005251210182905197 accuracy 0.608 auc 0.6560899897593446\n",
      "Val metrics: loss 0.005291318986564875 accuracy 0.614 auc 0.6868473316676277\n",
      "Epoch -  62\n",
      "Train metrics: loss 0.0052399528212845325 accuracy 0.614 auc 0.6567460317460317\n",
      "Val metrics: loss 0.005303733982145786 accuracy 0.56 auc 0.7017265680056377\n",
      "Epoch -  63\n",
      "Train metrics: loss 0.0051757520996034145 accuracy 0.618 auc 0.670586917562724\n",
      "Val metrics: loss 0.00521039217710495 accuracy 0.628 auc 0.695496188096611\n",
      "Epoch -  64\n",
      "Train metrics: loss 0.0051680663600564 accuracy 0.632 auc 0.6787154377880185\n",
      "Val metrics: loss 0.005211273208260536 accuracy 0.612 auc 0.7091581779742456\n",
      "Epoch -  65\n",
      "Train metrics: loss 0.005118841305375099 accuracy 0.644 auc 0.6937403993855608\n",
      "Val metrics: loss 0.0051583400927484035 accuracy 0.636 auc 0.723236594272535\n",
      "Epoch -  66\n",
      "Train metrics: loss 0.005037043709307909 accuracy 0.652 auc 0.7179339477726574\n",
      "Val metrics: loss 0.005144011694937944 accuracy 0.636 auc 0.7172464603754244\n",
      "Epoch -  67\n",
      "Train metrics: loss 0.0050437976606190205 accuracy 0.672 auc 0.7171658986175116\n",
      "Val metrics: loss 0.005146785173565149 accuracy 0.63 auc 0.7151162790697674\n",
      "Epoch -  68\n",
      "Train metrics: loss 0.005041946657001972 accuracy 0.644 auc 0.7123815924219151\n",
      "Val metrics: loss 0.005146785173565149 accuracy 0.63 auc 0.7151162790697674\n",
      "Epoch -  69\n",
      "Train metrics: loss 0.005116896238178015 accuracy 0.622 auc 0.6881080389144906\n",
      "Val metrics: loss 0.0051565817557275295 accuracy 0.62 auc 0.7157088858991607\n",
      "Epoch -  70\n",
      "Train metrics: loss 0.005030268803238869 accuracy 0.636 auc 0.7211981566820276\n",
      "Val metrics: loss 0.005172376520931721 accuracy 0.61 auc 0.7136908193990648\n",
      "Epoch -  71\n",
      "Train metrics: loss 0.0049277557991445065 accuracy 0.68 auc 0.7430075524833588\n",
      "Val metrics: loss 0.0051397220231592655 accuracy 0.682 auc 0.7119770645140624\n",
      "Epoch -  72\n",
      "Train metrics: loss 0.005059725139290094 accuracy 0.654 auc 0.7041570660522274\n",
      "Val metrics: loss 0.005153368227183819 accuracy 0.614 auc 0.7196329040937921\n",
      "Epoch -  73\n",
      "Train metrics: loss 0.005060712806880474 accuracy 0.65 auc 0.7056291602662569\n",
      "Val metrics: loss 0.005090649705380201 accuracy 0.642 auc 0.7267762188481005\n",
      "Epoch -  74\n",
      "Train metrics: loss 0.005147574003785849 accuracy 0.636 auc 0.6828757040450589\n",
      "Val metrics: loss 0.005068551283329725 accuracy 0.656 auc 0.7311166634633866\n",
      "Epoch -  75\n",
      "Train metrics: loss 0.005133618600666523 accuracy 0.634 auc 0.6812596006144394\n",
      "Val metrics: loss 0.005065482575446367 accuracy 0.646 auc 0.7317092702927798\n",
      "Epoch -  76\n",
      "Train metrics: loss 0.0049133943393826485 accuracy 0.686 auc 0.7408154121863799\n",
      "Val metrics: loss 0.005065482575446367 accuracy 0.646 auc 0.7317092702927798\n",
      "Epoch -  77\n",
      "Train metrics: loss 0.005008849315345287 accuracy 0.666 auc 0.7226542498719919\n",
      "Val metrics: loss 0.005076784174889326 accuracy 0.636 auc 0.7323178935229675\n",
      "Epoch -  78\n",
      "Train metrics: loss 0.00497392239049077 accuracy 0.682 auc 0.7399673579109063\n",
      "Val metrics: loss 0.005063830874860287 accuracy 0.656 auc 0.7413511435710167\n",
      "Epoch -  79\n",
      "Train metrics: loss 0.004921108949929476 accuracy 0.69 auc 0.7477918586789555\n",
      "Val metrics: loss 0.0050255232490599155 accuracy 0.714 auc 0.7498718687936448\n",
      "Epoch -  80\n",
      "Train metrics: loss 0.004979608580470085 accuracy 0.66 auc 0.7274065540194572\n",
      "Val metrics: loss 0.00501586776226759 accuracy 0.662 auc 0.7493273111666345\n",
      "Epoch -  81\n",
      "Train metrics: loss 0.0049528093077242374 accuracy 0.676 auc 0.7336149513568868\n",
      "Val metrics: loss 0.004978986922651529 accuracy 0.67 auc 0.7600262668973028\n",
      "Epoch -  82\n",
      "Train metrics: loss 0.004927119705826044 accuracy 0.688 auc 0.7429115463389656\n",
      "Val metrics: loss 0.0049956138245761395 accuracy 0.65 auc 0.7569511179447754\n",
      "Epoch -  83\n",
      "Train metrics: loss 0.004865963943302631 accuracy 0.686 auc 0.7569764464925754\n",
      "Val metrics: loss 0.004981888458132744 accuracy 0.65 auc 0.7590492664488436\n",
      "Epoch -  84\n",
      "Train metrics: loss 0.005018358584493399 accuracy 0.636 auc 0.7129416282642088\n",
      "Val metrics: loss 0.004981888458132744 accuracy 0.65 auc 0.7590492664488436\n",
      "Epoch -  85\n",
      "Train metrics: loss 0.00489033805206418 accuracy 0.692 auc 0.7514080901177675\n",
      "Val metrics: loss 0.004963920451700687 accuracy 0.68 auc 0.7637580882824012\n",
      "Epoch -  86\n",
      "Train metrics: loss 0.0049224174581468105 accuracy 0.666 auc 0.7383672555043523\n",
      "Val metrics: loss 0.0049601406790316105 accuracy 0.676 auc 0.7648632199372157\n",
      "Epoch -  87\n",
      "Train metrics: loss 0.004849368240684271 accuracy 0.694 auc 0.7642249103942653\n",
      "Val metrics: loss 0.0049646347761154175 accuracy 0.65 auc 0.7588570696393107\n",
      "Epoch -  88\n",
      "Train metrics: loss 0.004958536010235548 accuracy 0.67 auc 0.7278385816692268\n",
      "Val metrics: loss 0.00492424052208662 accuracy 0.654 auc 0.7617480299827022\n",
      "Epoch -  89\n",
      "Train metrics: loss 0.004862597677856684 accuracy 0.668 auc 0.7407674091141834\n",
      "Val metrics: loss 0.004871800076216459 accuracy 0.672 auc 0.7703728618104939\n",
      "Epoch -  90\n",
      "Train metrics: loss 0.004779163282364607 accuracy 0.69 auc 0.7733294930875576\n",
      "Val metrics: loss 0.004886352922767401 accuracy 0.66 auc 0.768514959318342\n",
      "Epoch -  91\n",
      "Train metrics: loss 0.004674865864217281 accuracy 0.738 auc 0.7916666666666667\n",
      "Val metrics: loss 0.004875454120337963 accuracy 0.666 auc 0.7697962713818951\n",
      "Epoch -  92\n",
      "Train metrics: loss 0.004656155128031969 accuracy 0.718 auc 0.8014912954429084\n",
      "Val metrics: loss 0.004875454120337963 accuracy 0.666 auc 0.7697962713818951\n",
      "Epoch -  93\n",
      "Train metrics: loss 0.004680217243731022 accuracy 0.738 auc 0.7945948540706607\n",
      "Val metrics: loss 0.0048643010668456554 accuracy 0.674 auc 0.7714299442629253\n",
      "Epoch -  94\n",
      "Train metrics: loss 0.004712502472102642 accuracy 0.728 auc 0.7892185099846389\n",
      "Val metrics: loss 0.0048463731072843075 accuracy 0.684 auc 0.7743769620090972\n",
      "Epoch -  95\n",
      "Train metrics: loss 0.004722994286566973 accuracy 0.688 auc 0.7773937532002048\n",
      "Val metrics: loss 0.0048697046004235744 accuracy 0.656 auc 0.7678422704849766\n",
      "Epoch -  96\n",
      "Train metrics: loss 0.004692602902650833 accuracy 0.708 auc 0.7825940860215054\n",
      "Val metrics: loss 0.004842197988182306 accuracy 0.688 auc 0.7725030431161508\n",
      "Epoch -  97\n",
      "Train metrics: loss 0.004754356108605862 accuracy 0.708 auc 0.7717453917050692\n",
      "Val metrics: loss 0.004851192235946655 accuracy 0.676 auc 0.7679223524889487\n",
      "Epoch -  98\n",
      "Train metrics: loss 0.004769170191138983 accuracy 0.686 auc 0.7639688940092166\n",
      "Val metrics: loss 0.004825575742870569 accuracy 0.672 auc 0.771189698251009\n",
      "Epoch -  99\n",
      "Train metrics: loss 0.004622390028089285 accuracy 0.734 auc 0.8057315668202765\n",
      "Val metrics: loss 0.0048184506595134735 accuracy 0.672 auc 0.772198731501057\n",
      "Train pass no. 1\n",
      "PPOMaskNet(\n",
      "  (fc1): SupermaskLinear(in_features=36000, out_features=200, bias=True)\n",
      "  (fc2): SupermaskLinear(in_features=200, out_features=100, bias=True)\n",
      "  (drop): Dropout(p=0.2, inplace=False)\n",
      "  (linear): Linear(in_features=100, out_features=1, bias=True)\n",
      ")\n",
      "{'trainLoss': 0.004622390028089285, 'testLoss': 0.0048184506595134735, 'trainAccuracy': 0.734, 'testAccuracy': 0.672, 'trainAUC': 0.8057315668202765, 'testAUC': 0.772198731501057}\n",
      "Retraining on the best_hparams to make sure we didn't just get good results by random chance.\n",
      "best_hparams {'dropout': 0.2, 'k': 0.1, 'fine_tune': False, 'use_last_linear': True, 'learning_rate': 0.0001, 'weight_decay': 0, 'q_head_index': None}\n",
      "Epoch -  0\n",
      "Train metrics: loss 0.19443157315254211 accuracy 0.51 auc 0.5124550590652286\n",
      "Val metrics: loss 0.28002244234085083 accuracy 0.496 auc 0.48086277521761395\n",
      "Epoch -  1\n",
      "Train metrics: loss 0.12192317098379135 accuracy 0.512 auc 0.4687981510015408\n",
      "Val metrics: loss 0.007519424892961979 accuracy 0.504 auc 0.5385624679979518\n",
      "Epoch -  2\n",
      "Train metrics: loss 0.010774712078273296 accuracy 0.482 auc 0.49587506420133537\n",
      "Val metrics: loss 0.0071472954005002975 accuracy 0.486 auc 0.4992959549411162\n",
      "Epoch -  3\n",
      "Train metrics: loss 0.008315208368003368 accuracy 0.558 auc 0.5266114535182331\n",
      "Val metrics: loss 0.0073217772878706455 accuracy 0.482 auc 0.4947516641065028\n",
      "Epoch -  4\n",
      "Train metrics: loss 0.007855113595724106 accuracy 0.538 auc 0.5127279147406266\n",
      "Val metrics: loss 0.0073217772878706455 accuracy 0.482 auc 0.4947516641065028\n",
      "Epoch -  5\n",
      "Train metrics: loss 0.00718134269118309 accuracy 0.518 auc 0.5097907036466358\n",
      "Val metrics: loss 0.0062310658395290375 accuracy 0.524 auc 0.5167210701484894\n",
      "Epoch -  6\n",
      "Train metrics: loss 0.006885257549583912 accuracy 0.516 auc 0.522663071391885\n",
      "Val metrics: loss 0.005760325118899345 accuracy 0.568 auc 0.5632200460829493\n",
      "Epoch -  7\n",
      "Train metrics: loss 0.006003798916935921 accuracy 0.504 auc 0.5276226245505906\n",
      "Val metrics: loss 0.005860531236976385 accuracy 0.442 auc 0.49164746543778803\n",
      "Epoch -  8\n",
      "Train metrics: loss 0.005998401436954737 accuracy 0.51 auc 0.48939072932717\n",
      "Val metrics: loss 0.005651704501360655 accuracy 0.448 auc 0.5118407578084998\n",
      "Epoch -  9\n",
      "Train metrics: loss 0.005733804311603308 accuracy 0.536 auc 0.5473324345146379\n",
      "Val metrics: loss 0.00538498954847455 accuracy 0.606 auc 0.6479294674859192\n",
      "Epoch -  10\n",
      "Train metrics: loss 0.005638162139803171 accuracy 0.554 auc 0.5784219311761685\n",
      "Val metrics: loss 0.005845590028911829 accuracy 0.496 auc 0.5928699436763953\n",
      "Epoch -  11\n",
      "Train metrics: loss 0.005639628507196903 accuracy 0.592 auc 0.6065581664098614\n",
      "Val metrics: loss 0.005439154338091612 accuracy 0.522 auc 0.6559139784946237\n",
      "Epoch -  12\n",
      "Train metrics: loss 0.005562870763242245 accuracy 0.56 auc 0.5890472521828455\n",
      "Val metrics: loss 0.005439154338091612 accuracy 0.522 auc 0.6559139784946237\n",
      "Epoch -  13\n",
      "Train metrics: loss 0.005788011010736227 accuracy 0.56 auc 0.567957113507961\n",
      "Val metrics: loss 0.0055066789500415325 accuracy 0.548 auc 0.6602822580645161\n",
      "Epoch -  14\n",
      "Train metrics: loss 0.005757823120802641 accuracy 0.546 auc 0.558166409861325\n",
      "Val metrics: loss 0.005498613230884075 accuracy 0.506 auc 0.650825652841782\n",
      "Epoch -  15\n",
      "Train metrics: loss 0.005590866319835186 accuracy 0.554 auc 0.5768008474576272\n",
      "Val metrics: loss 0.00538299698382616 accuracy 0.658 auc 0.6408890168970813\n",
      "Epoch -  16\n",
      "Train metrics: loss 0.005746574606746435 accuracy 0.548 auc 0.5513771186440679\n",
      "Val metrics: loss 0.005321058444678783 accuracy 0.606 auc 0.6756752432155658\n",
      "Epoch -  17\n",
      "Train metrics: loss 0.0054639349691569805 accuracy 0.582 auc 0.6291730868002054\n",
      "Val metrics: loss 0.00521581806242466 accuracy 0.634 auc 0.7048931131592422\n",
      "Epoch -  18\n",
      "Train metrics: loss 0.0053295292891561985 accuracy 0.578 auc 0.6271025937339497\n",
      "Val metrics: loss 0.005297360010445118 accuracy 0.53 auc 0.7088293650793651\n",
      "Epoch -  19\n",
      "Train metrics: loss 0.005433878395706415 accuracy 0.542 auc 0.6027060862865947\n",
      "Val metrics: loss 0.005209554918110371 accuracy 0.678 auc 0.7124775985663082\n",
      "Epoch -  20\n",
      "Train metrics: loss 0.005250909831374884 accuracy 0.634 auc 0.6576142783769902\n",
      "Val metrics: loss 0.005209554918110371 accuracy 0.678 auc 0.7124775985663082\n",
      "Epoch -  21\n",
      "Train metrics: loss 0.0053331139497458935 accuracy 0.592 auc 0.6326239085772983\n",
      "Val metrics: loss 0.005163433030247688 accuracy 0.642 auc 0.709005376344086\n",
      "Epoch -  22\n",
      "Train metrics: loss 0.00535224936902523 accuracy 0.598 auc 0.635641371340524\n",
      "Val metrics: loss 0.005182628054171801 accuracy 0.662 auc 0.7193900409626217\n",
      "Epoch -  23\n",
      "Train metrics: loss 0.005146901123225689 accuracy 0.638 auc 0.6800365947611711\n",
      "Val metrics: loss 0.005179313477128744 accuracy 0.656 auc 0.7258384536610345\n",
      "Epoch -  24\n",
      "Train metrics: loss 0.005089054349809885 accuracy 0.64 auc 0.6902446070878275\n",
      "Val metrics: loss 0.006050538271665573 accuracy 0.52 auc 0.7172779057859704\n",
      "Epoch -  25\n",
      "Train metrics: loss 0.0058654616586863995 accuracy 0.568 auc 0.6079224447868516\n",
      "Val metrics: loss 0.005145584233105183 accuracy 0.628 auc 0.7189420122887864\n",
      "Epoch -  26\n",
      "Train metrics: loss 0.005309199448674917 accuracy 0.598 auc 0.6867616846430404\n",
      "Val metrics: loss 0.005031278356909752 accuracy 0.686 auc 0.7517601126472094\n",
      "Epoch -  27\n",
      "Train metrics: loss 0.0050416914746165276 accuracy 0.63 auc 0.711848356445814\n",
      "Val metrics: loss 0.005109910387545824 accuracy 0.604 auc 0.7548323092677931\n",
      "Epoch -  28\n",
      "Train metrics: loss 0.005003814119845629 accuracy 0.642 auc 0.7226502311248075\n",
      "Val metrics: loss 0.005109910387545824 accuracy 0.604 auc 0.7548323092677931\n",
      "Epoch -  29\n",
      "Train metrics: loss 0.0049840835854411125 accuracy 0.656 auc 0.7109174370826913\n",
      "Val metrics: loss 0.004977562930434942 accuracy 0.676 auc 0.7480958781362008\n",
      "Epoch -  30\n",
      "Train metrics: loss 0.005144439171999693 accuracy 0.624 auc 0.6998587570621468\n",
      "Val metrics: loss 0.00499841570854187 accuracy 0.738 auc 0.7558563748079877\n",
      "Epoch -  31\n",
      "Train metrics: loss 0.0050495886243879795 accuracy 0.634 auc 0.7048022598870057\n",
      "Val metrics: loss 0.004955276381224394 accuracy 0.738 auc 0.7647689452124935\n",
      "Epoch -  32\n",
      "Train metrics: loss 0.004970194306224585 accuracy 0.648 auc 0.718156137647663\n",
      "Val metrics: loss 0.00494353286921978 accuracy 0.698 auc 0.7734735023041474\n",
      "Epoch -  33\n",
      "Train metrics: loss 0.004946683533489704 accuracy 0.658 auc 0.7208365434001027\n",
      "Val metrics: loss 0.0048667509108781815 accuracy 0.728 auc 0.7711533538146442\n",
      "Epoch -  34\n",
      "Train metrics: loss 0.004911867901682854 accuracy 0.676 auc 0.7262455059065229\n",
      "Val metrics: loss 0.004838376771658659 accuracy 0.748 auc 0.7806899641577062\n",
      "Epoch -  35\n",
      "Train metrics: loss 0.00476863095536828 accuracy 0.674 auc 0.7587795326142783\n",
      "Val metrics: loss 0.004842700436711311 accuracy 0.748 auc 0.783298131080389\n",
      "Epoch -  36\n",
      "Train metrics: loss 0.0048420061357319355 accuracy 0.676 auc 0.7484752182845403\n",
      "Val metrics: loss 0.004842700436711311 accuracy 0.748 auc 0.783298131080389\n",
      "Epoch -  37\n",
      "Train metrics: loss 0.004962423350661993 accuracy 0.66 auc 0.7134052388289677\n",
      "Val metrics: loss 0.004842475987970829 accuracy 0.738 auc 0.7847542242703534\n",
      "Epoch -  38\n",
      "Train metrics: loss 0.004777100868523121 accuracy 0.684 auc 0.7595178479712379\n",
      "Val metrics: loss 0.0048004998825490475 accuracy 0.738 auc 0.7846102150537634\n",
      "Epoch -  39\n",
      "Train metrics: loss 0.004841402173042297 accuracy 0.684 auc 0.7396475346687211\n",
      "Val metrics: loss 0.004767927806824446 accuracy 0.746 auc 0.7911066308243728\n",
      "Epoch -  40\n",
      "Train metrics: loss 0.004815671592950821 accuracy 0.7 auc 0.744526836158192\n",
      "Val metrics: loss 0.004739432130008936 accuracy 0.762 auc 0.7973470302099335\n",
      "Epoch -  41\n",
      "Train metrics: loss 0.004763131495565176 accuracy 0.664 auc 0.7538841807909604\n",
      "Val metrics: loss 0.004674581810832024 accuracy 0.75 auc 0.8061795954941116\n",
      "Epoch -  42\n",
      "Train metrics: loss 0.004680856131017208 accuracy 0.69 auc 0.7704320749871598\n",
      "Val metrics: loss 0.004703786224126816 accuracy 0.72 auc 0.8154921915002559\n",
      "Epoch -  43\n",
      "Train metrics: loss 0.004883706569671631 accuracy 0.672 auc 0.7377696456086287\n",
      "Val metrics: loss 0.00468863733112812 accuracy 0.778 auc 0.8196844598054276\n",
      "Epoch -  44\n",
      "Train metrics: loss 0.004716844297945499 accuracy 0.688 auc 0.7624711093990755\n",
      "Val metrics: loss 0.00468863733112812 accuracy 0.778 auc 0.8196844598054276\n",
      "Epoch -  45\n",
      "Train metrics: loss 0.004810457117855549 accuracy 0.648 auc 0.7373041859270673\n",
      "Val metrics: loss 0.004686044994741678 accuracy 0.754 auc 0.8191884280593957\n",
      "Epoch -  46\n",
      "Train metrics: loss 0.004600294400006533 accuracy 0.732 auc 0.7852465331278891\n",
      "Val metrics: loss 0.0045903706923127174 accuracy 0.754 auc 0.8090277777777778\n",
      "Epoch -  47\n",
      "Train metrics: loss 0.0046315547078847885 accuracy 0.696 auc 0.7680726759116591\n",
      "Val metrics: loss 0.0046090176329016685 accuracy 0.716 auc 0.8007552483358935\n",
      "Epoch -  48\n",
      "Train metrics: loss 0.004707897547632456 accuracy 0.678 auc 0.761203133025167\n",
      "Val metrics: loss 0.004589603282511234 accuracy 0.712 auc 0.8044194828469022\n",
      "Epoch -  49\n",
      "Train metrics: loss 0.004745906218886375 accuracy 0.688 auc 0.7479455572675913\n",
      "Val metrics: loss 0.004529326688498259 accuracy 0.738 auc 0.814484126984127\n",
      "Epoch -  50\n",
      "Train metrics: loss 0.0050677224062383175 accuracy 0.618 auc 0.69522021058038\n",
      "Val metrics: loss 0.005337247159332037 accuracy 0.514 auc 0.8256688428059396\n",
      "Epoch -  51\n",
      "Train metrics: loss 0.004936945624649525 accuracy 0.626 auc 0.7718926553672315\n",
      "Val metrics: loss 0.00449099438264966 accuracy 0.752 auc 0.816260240655402\n",
      "Epoch -  52\n",
      "Train metrics: loss 0.004588742274791002 accuracy 0.688 auc 0.7696777092963535\n",
      "Val metrics: loss 0.00449099438264966 accuracy 0.752 auc 0.816260240655402\n",
      "Epoch -  53\n",
      "Train metrics: loss 0.004481765907257795 accuracy 0.696 auc 0.7860972008217771\n",
      "Val metrics: loss 0.004503398202359676 accuracy 0.714 auc 0.8113639272913467\n",
      "Epoch -  54\n",
      "Train metrics: loss 0.004569381475448608 accuracy 0.712 auc 0.7784090909090909\n",
      "Val metrics: loss 0.004678053315728903 accuracy 0.678 auc 0.8191244239631337\n",
      "Epoch -  55\n",
      "Train metrics: loss 0.004475465510040522 accuracy 0.726 auc 0.8008956086286594\n",
      "Val metrics: loss 0.004495556931942701 accuracy 0.762 auc 0.8133640552995391\n",
      "Epoch -  56\n",
      "Train metrics: loss 0.004573650658130646 accuracy 0.682 auc 0.7712506420133539\n",
      "Val metrics: loss 0.004452373832464218 accuracy 0.758 auc 0.8119399641577061\n",
      "Epoch -  57\n",
      "Train metrics: loss 0.004582903813570738 accuracy 0.708 auc 0.7753755778120185\n",
      "Val metrics: loss 0.004439385607838631 accuracy 0.766 auc 0.8359735023041476\n",
      "Epoch -  58\n",
      "Train metrics: loss 0.004828028846532106 accuracy 0.66 auc 0.7334681561376476\n",
      "Val metrics: loss 0.004891420714557171 accuracy 0.54 auc 0.8372375832053252\n",
      "Epoch -  59\n",
      "Train metrics: loss 0.004761882126331329 accuracy 0.674 auc 0.7551200564971752\n",
      "Val metrics: loss 0.0057858871296048164 accuracy 0.594 auc 0.8139080901177675\n",
      "Epoch -  60\n",
      "Train metrics: loss 0.006038418505340815 accuracy 0.568 auc 0.7906554956343091\n",
      "Val metrics: loss 0.0057858871296048164 accuracy 0.594 auc 0.8139080901177675\n",
      "Epoch -  61\n",
      "Train metrics: loss 0.005035179667174816 accuracy 0.642 auc 0.714175654853621\n",
      "Val metrics: loss 0.005705039948225021 accuracy 0.498 auc 0.8257648489503329\n",
      "Epoch -  62\n",
      "Train metrics: loss 0.005010063294321299 accuracy 0.612 auc 0.7591326399589111\n",
      "Val metrics: loss 0.004565862938761711 accuracy 0.81 auc 0.8295410906298002\n",
      "Epoch -  63\n",
      "Train metrics: loss 0.004811356775462627 accuracy 0.66 auc 0.7719247560349255\n",
      "Val metrics: loss 0.004455507267266512 accuracy 0.764 auc 0.8311411930363543\n",
      "Epoch -  64\n",
      "Train metrics: loss 0.0046444060280919075 accuracy 0.678 auc 0.7773176682074987\n",
      "Val metrics: loss 0.004379598889499903 accuracy 0.736 auc 0.829237071172555\n",
      "Epoch -  65\n",
      "Train metrics: loss 0.0045764269307255745 accuracy 0.682 auc 0.7827908320493067\n",
      "Val metrics: loss 0.0044890353456139565 accuracy 0.744 auc 0.8287730414746544\n",
      "Epoch -  66\n",
      "Train metrics: loss 0.004441087134182453 accuracy 0.692 auc 0.801024011299435\n",
      "Val metrics: loss 0.004433492664247751 accuracy 0.764 auc 0.8218445980542755\n",
      "Epoch -  67\n",
      "Train metrics: loss 0.004497724585235119 accuracy 0.696 auc 0.7816833590138674\n",
      "Val metrics: loss 0.004426955245435238 accuracy 0.758 auc 0.8179403481822837\n",
      "Epoch -  68\n",
      "Train metrics: loss 0.004382624756544828 accuracy 0.73 auc 0.8033192090395479\n",
      "Val metrics: loss 0.004426955245435238 accuracy 0.758 auc 0.8179403481822837\n",
      "Epoch -  69\n",
      "Train metrics: loss 0.0044364952482283115 accuracy 0.712 auc 0.7934643040575243\n",
      "Val metrics: loss 0.0045030377805233 accuracy 0.726 auc 0.8249967997951869\n",
      "Epoch -  70\n",
      "Train metrics: loss 0.0043641054071486 accuracy 0.722 auc 0.8063045711350796\n",
      "Val metrics: loss 0.0048047457821667194 accuracy 0.694 auc 0.8007712493599589\n",
      "Epoch -  71\n",
      "Train metrics: loss 0.0046800291165709496 accuracy 0.684 auc 0.754879301489471\n",
      "Val metrics: loss 0.004442843608558178 accuracy 0.754 auc 0.830805171530978\n",
      "Epoch -  72\n",
      "Train metrics: loss 0.004586536902934313 accuracy 0.7 auc 0.7734495377503852\n",
      "Val metrics: loss 0.00451754592359066 accuracy 0.69 auc 0.8339733742959548\n",
      "Epoch -  73\n",
      "Train metrics: loss 0.004318252205848694 accuracy 0.72 auc 0.8289515921931176\n",
      "Val metrics: loss 0.004465959034860134 accuracy 0.712 auc 0.8217165898617512\n",
      "Epoch -  74\n",
      "Train metrics: loss 0.004380085039883852 accuracy 0.742 auc 0.8036562660503339\n",
      "Val metrics: loss 0.00446104584261775 accuracy 0.716 auc 0.836021505376344\n",
      "Epoch -  75\n",
      "Train metrics: loss 0.004268001765012741 accuracy 0.738 auc 0.8357569337442218\n",
      "Val metrics: loss 0.004324436187744141 accuracy 0.76 auc 0.8335413466461854\n",
      "Epoch -  76\n",
      "Train metrics: loss 0.00419256929308176 accuracy 0.752 auc 0.8313751926040062\n",
      "Val metrics: loss 0.004324436187744141 accuracy 0.76 auc 0.8335413466461854\n",
      "Epoch -  77\n",
      "Train metrics: loss 0.004218412563204765 accuracy 0.748 auc 0.8270095017976373\n",
      "Val metrics: loss 0.004279494751244783 accuracy 0.77 auc 0.8376696108550948\n",
      "Epoch -  78\n",
      "Train metrics: loss 0.0042624203488230705 accuracy 0.718 auc 0.8186472778633795\n",
      "Val metrics: loss 0.0042873346246778965 accuracy 0.782 auc 0.844134024577573\n",
      "Epoch -  79\n",
      "Train metrics: loss 0.004498734604567289 accuracy 0.716 auc 0.7803190806368772\n",
      "Val metrics: loss 0.004207671619951725 accuracy 0.74 auc 0.8563748079877112\n",
      "Epoch -  80\n",
      "Train metrics: loss 0.00442863954231143 accuracy 0.706 auc 0.7852465331278891\n",
      "Val metrics: loss 0.004481232259422541 accuracy 0.602 auc 0.8566788274449565\n",
      "Epoch -  81\n",
      "Train metrics: loss 0.004039878025650978 accuracy 0.752 auc 0.8376187724704675\n",
      "Val metrics: loss 0.004094269126653671 accuracy 0.798 auc 0.8587909626216079\n",
      "Epoch -  82\n",
      "Train metrics: loss 0.0041510434821248055 accuracy 0.726 auc 0.8219857473035439\n",
      "Val metrics: loss 0.004063726402819157 accuracy 0.776 auc 0.8533026113671276\n",
      "Epoch -  83\n",
      "Train metrics: loss 0.004277592059224844 accuracy 0.734 auc 0.8052934001027221\n",
      "Val metrics: loss 0.004071374423801899 accuracy 0.796 auc 0.8577348950332822\n",
      "Epoch -  84\n",
      "Train metrics: loss 0.003939111717045307 accuracy 0.754 auc 0.8488540061633282\n",
      "Val metrics: loss 0.004071374423801899 accuracy 0.796 auc 0.8577348950332822\n",
      "Epoch -  85\n",
      "Train metrics: loss 0.004230302758514881 accuracy 0.728 auc 0.812788906009245\n",
      "Val metrics: loss 0.0042488654144108295 accuracy 0.776 auc 0.8443900409626216\n",
      "Epoch -  86\n",
      "Train metrics: loss 0.004316245671361685 accuracy 0.714 auc 0.8011363636363635\n",
      "Val metrics: loss 0.0041615767404437065 accuracy 0.754 auc 0.843926011264721\n",
      "Epoch -  87\n",
      "Train metrics: loss 0.004282649140805006 accuracy 0.718 auc 0.8199634052388288\n",
      "Val metrics: loss 0.0042330678552389145 accuracy 0.742 auc 0.8465661802355351\n",
      "Epoch -  88\n",
      "Train metrics: loss 0.004342720378190279 accuracy 0.72 auc 0.8040414740626606\n",
      "Val metrics: loss 0.004123299848288298 accuracy 0.778 auc 0.8533026113671275\n",
      "Epoch -  89\n",
      "Train metrics: loss 0.004176123533397913 accuracy 0.754 auc 0.8327234206471497\n",
      "Val metrics: loss 0.004121938720345497 accuracy 0.798 auc 0.8603270609318996\n",
      "Epoch -  90\n",
      "Train metrics: loss 0.004115651361644268 accuracy 0.748 auc 0.8415511042629686\n",
      "Val metrics: loss 0.004176871385425329 accuracy 0.794 auc 0.8627752176139274\n",
      "Epoch -  91\n",
      "Train metrics: loss 0.004024388734251261 accuracy 0.774 auc 0.8571680790960452\n",
      "Val metrics: loss 0.0041273292154073715 accuracy 0.798 auc 0.8627912186379929\n",
      "Epoch -  92\n",
      "Train metrics: loss 0.004084355663508177 accuracy 0.75 auc 0.8491589625064202\n",
      "Val metrics: loss 0.0041273292154073715 accuracy 0.798 auc 0.8627912186379929\n",
      "Epoch -  93\n",
      "Train metrics: loss 0.004088669549673796 accuracy 0.774 auc 0.8505232408834104\n",
      "Val metrics: loss 0.004088259767740965 accuracy 0.78 auc 0.8596070148489503\n",
      "Epoch -  94\n",
      "Train metrics: loss 0.004200250841677189 accuracy 0.736 auc 0.8339111453518233\n",
      "Val metrics: loss 0.004052270203828812 accuracy 0.802 auc 0.8655113927291347\n",
      "Epoch -  95\n",
      "Train metrics: loss 0.0045564440079033375 accuracy 0.674 auc 0.7771732152028762\n",
      "Val metrics: loss 0.004856917541474104 accuracy 0.69 auc 0.8513824884792627\n",
      "Epoch -  96\n",
      "Train metrics: loss 0.004867659881711006 accuracy 0.646 auc 0.7357473035439137\n",
      "Val metrics: loss 0.004055737983435392 accuracy 0.772 auc 0.8595270097286226\n",
      "Epoch -  97\n",
      "Train metrics: loss 0.004408145323395729 accuracy 0.718 auc 0.7922926296866974\n",
      "Val metrics: loss 0.0041628205217421055 accuracy 0.772 auc 0.8521825396825397\n",
      "Epoch -  98\n",
      "Train metrics: loss 0.004080129787325859 accuracy 0.752 auc 0.8433326913199793\n",
      "Val metrics: loss 0.0042171780951321125 accuracy 0.742 auc 0.8453661034306197\n",
      "Epoch -  99\n",
      "Train metrics: loss 0.0044344221241772175 accuracy 0.704 auc 0.7868355161787366\n",
      "Val metrics: loss 0.004107359331101179 accuracy 0.782 auc 0.8568868407578084\n",
      "Train pass no. 1\n",
      "PPOMaskNet(\n",
      "  (fc1): SupermaskLinear(in_features=36000, out_features=200, bias=True)\n",
      "  (fc2): SupermaskLinear(in_features=200, out_features=100, bias=True)\n",
      "  (drop): Dropout(p=0.2, inplace=False)\n",
      "  (linear): Linear(in_features=100, out_features=1, bias=True)\n",
      ")\n",
      "Result of retrain on the best hyperparameters {'trainLoss': 0.0044344221241772175, 'testLoss': 0.004107359331101179, 'trainAccuracy': 0.704, 'testAccuracy': 0.782, 'trainAUC': 0.7868355161787366, 'testAUC': 0.8568868407578084}\n"
     ]
    }
   ],
   "source": [
    "hparam_combinations = []\n",
    "for hparam_possibilities in all_hparam_possibilities:\n",
    "  hparam_keys, hparam_values = zip(*hparam_possibilities.items())\n",
    "  hparam_combinations.extend([dict(zip(hparam_keys, v)) for v in itertools.product(*hparam_values)])\n",
    "random.shuffle(hparam_combinations)\n",
    "print('PARAMS', PARAMS)\n",
    "print(\"len(hparam_combinations)\", len(hparam_combinations), \"hparam_combinations\", hparam_combinations)\n",
    "\n",
    "best_test_auc = -float('inf')\n",
    "for hparams in hparam_combinations:\n",
    "    print(\"hparams\", hparams)\n",
    "    results = multi_agent_train(hparams)\n",
    "    print(results)\n",
    "    test_auc = results['testAUC']\n",
    "    if test_auc > best_test_auc:\n",
    "        best_test_auc = test_auc\n",
    "        best_hparams = hparams\n",
    "\n",
    "print(\"Retraining on the best_hparams to make sure we didn't just get good results by random chance.\")\n",
    "print(\"best_hparams\", best_hparams)\n",
    "print(\"Result of retrain on the best hyperparameters\", multi_agent_train(best_hparams))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Clre1lUXPXYg"
   },
   "source": [
    "## Initialise Subnets Search with activation that obtained optimal AUC in previous experiment\n",
    "\n",
    "We do this both as a sanity check as well as a potential improvement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 70
    },
    "id": "RO5NwFAmPh2l",
    "outputId": "934c0786-040b-4462-dee0-a533b1ed5702"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUC from only picking a single activation\n",
      "34 train 0.8168389427248516 val 0.8203183087179845\n",
      "13 train 0.6133060108723783 val 0.6179787990821131\n"
     ]
    }
   ],
   "source": [
    "from sklearn import metrics \n",
    "\n",
    "acts = []\n",
    "prefs = []\n",
    "\n",
    "for data in all_raw_data:\n",
    "    for i in range(data.observation.shape[0]):\n",
    "        acts.append(np.copy(data.policy_info[\"activations\"][i]))\n",
    "        prefs.append((data.policy_info['satisfaction'].as_list()[i] > -6).astype(int))\n",
    "\n",
    "acts = np.array(acts)\n",
    "prefs = np.array(prefs)\n",
    "\n",
    "def display_auc_info(xs, ys):\n",
    "    \n",
    "    def calc_auc(xs, ys, i):\n",
    "        fpr, tpr, thresholds = metrics.roc_curve(ys, xs[:,i], pos_label=1)\n",
    "        return metrics.auc(fpr, tpr)\n",
    "\n",
    "\n",
    "    multi_runs_aucs = []\n",
    "    for run_ix in range(50):\n",
    "        xs, ys = shuffle(xs, ys)\n",
    "        flat_xs = np.reshape(xs, (xs.shape[0], -1))\n",
    "        aucs = []    \n",
    "        \n",
    "        for i in range(flat_xs.shape[1]):\n",
    "            auc = calc_auc(flat_xs[:params['num_train']], ys[:params['num_train']], i)\n",
    "            aucs.append(auc)  \n",
    "\n",
    "        aucs = np.array(aucs)\n",
    "        multi_runs_aucs.append(aucs)\n",
    "\n",
    "    aucs = np.array(multi_runs_aucs)\n",
    "    aucs = aucs.mean(axis=0)\n",
    "\n",
    "    print(\"AUC from only picking a single activation\")\n",
    "    print(np.argmin(aucs), \"train\", 1-np.min(aucs), \"val\", 1-calc_auc(flat_xs[params['num_train']:], ys[params['num_train']:], np.argmin(aucs)))\n",
    "    print(np.argmax(aucs), \"train\", np.max(aucs), \"val\", calc_auc(flat_xs[params['num_train']:], ys[params['num_train']:], np.argmax(aucs)))\n",
    "  \n",
    "display_auc_info(acts, prefs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 123
    },
    "id": "M9Yw8LJRYEnL",
    "outputId": "9dbfdd4e-a3a5-4d67-9017-e75a952d36eb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train pass no. 1\n",
      "Train pass no. 2\n",
      "Train pass no. 3\n",
      "Train pass no. 4\n",
      "Train pass no. 5\n",
      "0.8254807323986824\n"
     ]
    }
   ],
   "source": [
    "## TODO: fix broken code after refactoring notebook\n",
    "\n",
    "best_act_index = 34\n",
    "K = 66774 / 67152  # Num of weigths with all dense activations except one set to 0 / Number of total weights\n",
    "\n",
    "params['num_epochs'] = 1\n",
    "results = multi_runs(fine_tune=False, K=1, q_head_index=None, q_means_stds=q_mu_s, \n",
    "                     use_last_linear=False, init_from_act_index=best_act_index,\n",
    "                     learning_rate=0.000, weight_decay=0.000, plots=False)\n",
    "\n",
    "####results is a tuple where the first element is the dict\n",
    "print(1-max(results[0]['testAUC']))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "find-subnets-torch_ppo",
   "provenance": []
  },
  "environment": {
   "name": "tf2-2-2-gpu.2-2.m49",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-2-2-gpu.2-2:m49"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}