train_test_plot_withSSLpretraining.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7db72646",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1625,
     "status": "ok",
     "timestamp": 1633304691157,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "7db72646",
    "outputId": "5422674a-185f-4f29-f30f-076994750869"
   },
   "outputs": [],
   "source": [
    "import torch_geometric\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "from collections import defaultdict\n",
    "import os\n",
    "import itertools\n",
    "import sys\n",
    "folder_path = './'\n",
    "sys.path.insert(1, folder_path)\n",
    "from train_pipeline import train\n",
    "from SSL_pipeline import pretrain_SSL\n",
    "import seaborn\n",
    "seaborn.set_style('darkgrid')\n",
    "plt.rcParams.update({'font.size': 14, 'lines.markersize': 8, 'lines.linewidth': 2})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78a9fab4",
   "metadata": {
    "executionInfo": {
     "elapsed": 11,
     "status": "ok",
     "timestamp": 1633304691158,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "78a9fab4"
   },
   "outputs": [],
   "source": [
    "random_seeds = [13531, 24642, 35753]\n",
    "k_range = [4, 6, 8, 10] # number of transmitter-receiver pairs for training\n",
    "num_train_samples_range = [256] # number of training samples\n",
    "num_test_samples = 256 # number of test samples\n",
    "modes = ['Supervised', 'Unsupervised'] # modes of training\n",
    "figsize = (8, 6)\n",
    "os.makedirs(folder_path + 'saved_training_data/results', exist_ok=True)\n",
    "os.makedirs(folder_path + 'saved_training_data/models', exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "_D5q6PW4q2y2",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1835248,
     "status": "ok",
     "timestamp": 1633306526396,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "_D5q6PW4q2y2",
    "outputId": "d19079ec-0caa-4c8c-f78b-9e9b833b47c9"
   },
   "outputs": [],
   "source": [
    "# perform SSL\n",
    "for k, num_train_samples, seed in itertools.product(k_range, num_train_samples_range, random_seeds):\n",
    "    num_samples = {'train': np.max(num_train_samples_range), 'test': num_test_samples}\n",
    "    mode = 'SSL'\n",
    "    experiment_name = 'k_train_{}_num_train_samples_{}_mode_{}_seed_{}'.format(k, str(num_samples), mode, seed)\n",
    "    print('Now running experiment \\\"{}\\\" ...'.format(experiment_name))\n",
    "    all_epoch_results, best_model = pretrain_SSL(k, num_samples, seed)\n",
    "    torch.save(all_epoch_results, folder_path + 'saved_training_data/results/{}.json'.format(experiment_name))\n",
    "    torch.save(best_model, folder_path + 'saved_training_data/models/{}.pt'.format(experiment_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "AlotSRJLd0oi",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 2017244,
     "status": "ok",
     "timestamp": 1633308543624,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "AlotSRJLd0oi",
    "outputId": "d2fcfe10-ebd6-47df-bbf7-3ca2300e9790"
   },
   "outputs": [],
   "source": [
    "# fine-tune using supervised/unsupervised training\n",
    "for k, num_train_samples, mode, seed in itertools.product(k_range, num_train_samples_range, modes, random_seeds):\n",
    "    num_samples = {'train': num_train_samples, 'test': num_test_samples}\n",
    "    # load pre-trained SSL model\n",
    "    SSL_experiment_name = 'k_train_{}_num_train_samples_{}_mode_{}_seed_{}'.format(k, str(num_samples), 'SSL', seed)\n",
    "    best_SSL_model = torch.load(folder_path + 'saved_training_data/models/{}.pt'.format(SSL_experiment_name))\n",
    "    # start training\n",
    "    experiment_name = 'k_train_{}_num_train_samples_{}_mode_{}_seed_{}_afterSSL'.format(k, str(num_samples), mode, seed)\n",
    "    print('Now running experiment \\\"{}\\\" ...'.format(experiment_name))\n",
    "    all_epoch_results, best_model = train(k, num_samples, mode, seed, best_SSL_model)\n",
    "    torch.save(all_epoch_results, folder_path + 'saved_training_data/results/{}.json'.format(experiment_name))\n",
    "    torch.save(best_model, folder_path + 'saved_training_data/models/{}.pt'.format(experiment_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bBrvxrFGS-S",
   "metadata": {
    "executionInfo": {
     "elapsed": 4999,
     "status": "ok",
     "timestamp": 1633309441687,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "5bBrvxrFGS-S"
   },
   "outputs": [],
   "source": [
    "# load all results with SSL\n",
    "all_epoch_results_SSL = dict()\n",
    "for k, num_train_samples, mode, seed in itertools.product(k_range, num_train_samples_range, modes, random_seeds):\n",
    "    num_samples = {'train': num_train_samples, 'test': num_test_samples}\n",
    "    experiment_name = 'k_train_{}_num_train_samples_{}_mode_{}_seed_{}_afterSSL'.format(k, str(num_samples), mode, seed)\n",
    "    all_epoch_results_SSL[k, str(num_samples), mode, seed] = torch.load(folder_path + 'saved_training_data/results/{}.json'.format(experiment_name))\n",
    "\n",
    "# load all results without SSL\n",
    "all_epoch_results = dict()\n",
    "for k, num_train_samples, mode, seed in itertools.product(k_range, num_train_samples_range, modes, random_seeds):\n",
    "    num_samples = {'train': num_train_samples, 'test': num_test_samples}\n",
    "    experiment_name = 'k_train_{}_num_train_samples_{}_mode_{}_seed_{}'.format(k, str(num_samples), mode, seed)\n",
    "    all_epoch_results[k, str(num_samples), mode, seed] = torch.load(folder_path + 'saved_training_data/results/{}.json'.format(experiment_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "S2dt_NGdUnFX",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 431
    },
    "executionInfo": {
     "elapsed": 797,
     "status": "ok",
     "timestamp": 1633309860465,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "S2dt_NGdUnFX",
    "outputId": "917ba8e9-8513-42da-bfc3-fc5543f6e687"
   },
   "outputs": [],
   "source": [
    "# convergence behavior\n",
    "plt.figure(figsize=figsize)\n",
    "threshold = 0.8\n",
    "markers = ['o', 'x']\n",
    "linestyles = [':', '-']\n",
    "num_samples = {'train': np.max(num_train_samples_range), 'test': num_test_samples}\n",
    "for i, mode in enumerate(modes):\n",
    "    # without SSL\n",
    "    epochs_all_seeds = []\n",
    "    for seed in random_seeds:\n",
    "        test_epochs = []\n",
    "        for k in k_range:\n",
    "            res = all_epoch_results[k, str(num_samples), mode, seed]\n",
    "            first_epoch_above_threshold = 1 + np.where(np.array(res['test', k, 'normalized_sum_rate']) >= threshold)[0][0]\n",
    "            test_epochs.append(first_epoch_above_threshold)\n",
    "        epochs_all_seeds.append(test_epochs)\n",
    "    plt.plot(k_range, np.mean(epochs_all_seeds, axis=0), linestyles[0], label=mode, marker=markers[i])\n",
    "    \n",
    "    # without SSL\n",
    "    epochs_all_seeds = []\n",
    "    for seed in random_seeds:\n",
    "        test_epochs = []\n",
    "        for k in k_range:\n",
    "            res = all_epoch_results_SSL[k, str(num_samples), mode, seed]\n",
    "            first_epoch_above_threshold = 1 + np.where(np.array(res['test', k, 'normalized_sum_rate']) >= threshold)[0][0]\n",
    "            test_epochs.append(first_epoch_above_threshold)\n",
    "        epochs_all_seeds.append(test_epochs)\n",
    "    plt.plot(k_range, np.mean(epochs_all_seeds, axis=0), linestyles[1], label=mode + ' (with SSL pre-training)', marker=markers[i], color=plt.gca().lines[-1].get_color())\n",
    "plt.xlabel(r'$K$')\n",
    "plt.ylabel(r'First epoch with normalized sum-rate $\\geq$ 0.8')\n",
    "plt.xticks(k_range)\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "yTq3uT32SvfB",
   "metadata": {
    "executionInfo": {
     "elapsed": 7,
     "status": "ok",
     "timestamp": 1633308549513,
     "user": {
      "displayName": "Navid NaderiAlizadeh",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ghjujt3S2rcWpfcLut6MmW6YZc9vrkQL0oal1FrBT4=s64",
      "userId": "14430719405503000266"
     },
     "user_tz": 420
    },
    "id": "yTq3uT32SvfB"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "run_pipeline_and_plot_SSL.ipynb",
   "provenance": []
  },
  "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}