05_federated_learning_beam_selection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/aldebaro/ai6g/blob/main/05_federated_learning_beam_selection.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc42cc4f",
   "metadata": {},
   "source": [
    "**Inteligência Artificial e Aprendizado de Máquina Aplicados a Redes 5G e 6G**.\n",
    "*Aldebaro Klautau* (UFPA). Minicurso 5 do SBrT - 25 de setembro de 2022."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the repository if running in Colab and install all the dependencies\n",
    "if 'google.colab' in str(get_ipython()):\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/drive/')\n",
    "    import sys\n",
    "    import os\n",
    "    try:\n",
    "      !git clone https://github.com/aldebaro/ai6g.git\n",
    "    except:\n",
    "      print(\"ai6g is already in the contents\")\n",
    "    %cd ai6g\n",
    "    !ln -s /content/drive/MyDrive/ai6g_files/files_05_federated/* ./files_05_federated\n",
    "    !pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Federated Learning\n",
    "\n",
    "## An example of Federated Learning applied to beam selection using LIDAR data - From ITU Challenge - https://github.com/ITU-AI-ML-in-5G-Challenge/PS-012-ML5G-PHY-Beam-Selection_Imperial_IPC1 -"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "from files_05_federated.dataloader import LidarDataset2D\n",
    "from files_05_federated.models import Lidar2D\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_federated as tff\n",
    "import scipy.io as sio\n",
    "import numpy as np\n",
    "import scipy.io as sio\n",
    "import numpy as np\n",
    "import scipy.stats\n",
    "import math\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "################################Arguments################################\n",
    "MONTECARLO = 1\n",
    "NUM_VEHICLES = 2\n",
    "LOCAL_EPOCHS = 5\n",
    "AGGREGATION_ROUNDS = 30\n",
    "\n",
    "BATCH_SIZE = 16\n",
    "SHUFFLE_BUFFER = 20\n",
    "PREFETCH_BUFFER=10\n",
    "\n",
    "data_path = \"./files_05_federated/files_05_data/\"\n",
    "#Remake train and test split data\n",
    "lidar_training_path = [data_path+\"lidar_input_train.npz\", data_path+\"lidar_input_validation.npz\"] ###Raymobtime s008\n",
    "beam_training_path = [data_path+\"beams_output_train.npz\", data_path+\"beams_output_validation.npz\"] ###Raymobtime s008\n",
    "\n",
    "lidar_test_path = [data_path+\"lidar_input_test.npz\"] ###Raymobtime s009\n",
    "beam_test_path = [data_path+\"beams_output_test.npz\"] ###Raymobtime s009"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset process functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_local_dataset(lidar_path, beam_path, num_vehicles, vehicle_ID):\n",
    "    training_data = LidarDataset2D(lidar_path, beam_path)\n",
    "    training_data.lidar_data = np.transpose(training_data.lidar_data, (0, 2, 3, 1))\n",
    "    x=training_data.lidar_data\n",
    "    xx = x[vehicle_ID*int(x.shape[0]/num_vehicles):(vehicle_ID+1)*int(x.shape[0]/num_vehicles),:,:,:] ###Split Lidar Data\n",
    "    y = training_data.beam_output\n",
    "    yy = y[vehicle_ID*int(y.shape[0]/num_vehicles):(vehicle_ID+1)*int(y.shape[0]/num_vehicles),:] ###Split Beam Labels\n",
    "    \n",
    "    dataset_train = tf.data.Dataset.from_tensor_slices((list(xx.astype(np.float32)),list(yy.astype(np.float32))))\n",
    "    #sio.savemat('label'+str(k)+'.mat',{'label'+str(k):yy})\n",
    "    return dataset_train\n",
    "\n",
    "def get_test_dataset(lidar_path, beam_path):\n",
    "    test_data = LidarDataset2D(lidar_path, beam_path)\n",
    "    test_data.lidar_data = np.transpose(test_data.lidar_data, (0, 2, 3, 1))\n",
    "    dataset_test = tf.data.Dataset.from_tensor_slices((list(test_data.lidar_data.astype(np.float32)),list(test_data.beam_output.astype(np.float32))))\n",
    "    return dataset_test\n",
    "\n",
    "def preprocess(dataset):\n",
    "  def batch_format_fn(element1,element2):\n",
    "    return collections.OrderedDict(x=element1, y=element2)\n",
    "  return dataset.repeat(LOCAL_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NN layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_keras_model():\n",
    "  return tf.keras.models.Sequential([\n",
    "    tf.keras.layers.Input(shape=(20, 200, 1)),\n",
    "    tf.keras.layers.Conv2D(5, 3, 1, padding='same'),# kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Conv2D(5, 3, 1, padding='same'),# kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Conv2D(5, 3, 2, padding='same'),# kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Conv2D(5, 3, 1, padding='same'),# kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Conv2D(5, 3, 2, padding='same'),#, kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Conv2D(1, 3, (1, 2), padding='same'),#, kernel_initializer=initializers.HeUniform),\n",
    "    tf.keras.layers.BatchNormalization(axis=3),\n",
    "    tf.keras.layers.PReLU(shared_axes=[1, 2]),\n",
    "    tf.keras.layers.Flatten(),\n",
    "    tf.keras.layers.Dense(16),\n",
    "    tf.keras.layers.ReLU(),\n",
    "    # layers.Dropout(0.7),\n",
    "    tf.keras.layers.Dense(256),\n",
    "    tf.keras.layers.Softmax()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "temp_dataset = get_local_dataset(lidar_training_path, beam_training_path,NUM_VEHICLES,0)\n",
    "preprocessed_example_dataset=preprocess(temp_dataset)\n",
    "example_element = next(iter((preprocessed_example_dataset)))\n",
    "\n",
    "def model_fn():\n",
    "  keras_model = create_keras_model()\n",
    "  top1 = tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='top_1_categorical_accuracy', dtype=None)\n",
    "  top10 = tf.keras.metrics.TopKCategoricalAccuracy(k=10, name='top_10_categorical_accuracy', dtype=None)\n",
    "  return tff.learning.from_keras_model(keras_model,\n",
    "      input_spec=preprocessed_example_dataset.element_spec,\n",
    "      loss=tf.keras.losses.CategoricalCrossentropy(),\n",
    "      metrics=[top1,top10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterative_process = tff.learning.build_federated_averaging_process(\n",
    "    model_fn,\n",
    "    client_optimizer_fn=lambda: tf.keras.optimizers.Adam(learning_rate=5e-3),\n",
    "    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=.25))\n",
    "\n",
    "evaluation = tff.learning.build_federated_evaluation(model_fn)\n",
    "\n",
    "test_data = LidarDataset2D(lidar_test_path, beam_test_path)\n",
    "test_data.lidar_data = np.transpose(test_data.lidar_data, (0, 2, 3, 1))\n",
    "\n",
    "accFL=0  \n",
    "for MONTECARLOi in range(MONTECARLO):\n",
    "    ###Generate Federated Train Dataset\n",
    "    federated_train_data=[]     \n",
    "    for i in range(NUM_VEHICLES):\n",
    "        print('Vehicle: ', i)\n",
    "        train_dataset = get_local_dataset(lidar_training_path, beam_training_path,NUM_VEHICLES,i)\n",
    "        federated_train_data.append(preprocess(train_dataset))\n",
    "        \n",
    "    ###Generate Test Dataset\n",
    "    test_dataset = get_test_dataset(lidar_test_path, beam_test_path)\n",
    "    federated_test_data=[preprocess(test_dataset)]\n",
    "    \n",
    "    top1=np.zeros(AGGREGATION_ROUNDS)\n",
    "    top10=np.zeros(AGGREGATION_ROUNDS)\n",
    "    \n",
    "    state = iterative_process.initialize() ###Initialize training\n",
    "    \n",
    "    #top1 = tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='top_1_categorical_accuracy', dtype=None)\n",
    "    #top10 = tf.keras.metrics.TopKCategoricalAccuracy(k=10, name='top_10_categorical_accuracy', dtype=None)\n",
    "    \n",
    "    ###Federated Training    \n",
    "    for round_num in range(AGGREGATION_ROUNDS):\n",
    "      state, metrics = iterative_process.next(state, federated_train_data)\n",
    "      test_metrics = evaluation(state.model, federated_test_data)['eval']\n",
    "      \n",
    "      print(str(metrics))\n",
    "      print(str(test_metrics))\n",
    "      print(test_metrics)\n",
    "      \n",
    "      top1[round_num]=test_metrics['top_1_categorical_accuracy']\n",
    "      top10[round_num]=test_metrics['top_10_categorical_accuracy']\n",
    "    \n",
    "      ###Generate Accuracy and Throughput Performance Curves\n",
    "      keras_model = create_keras_model()\n",
    "      #keras_model.compile(loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[top1,top10])\n",
    "      state.model.assign_weights_to(keras_model)\n",
    "      test_preds = keras_model.predict(test_data.lidar_data, batch_size=100)\n",
    "      test_preds_idx = np.argsort(test_preds, axis=1)\n",
    "      top_k = np.zeros(100)\n",
    "      throughput_ratio_at_k = np.zeros(100)\n",
    "      correct = 0\n",
    "      for i in range(100):\n",
    "        correct += np.sum(test_preds_idx[:, -1-i] == np.argmax(test_data.beam_output, axis=1))\n",
    "        top_k[i] = correct/test_data.beam_output.shape[0]\n",
    "        throughput_ratio_at_k[i] = np.sum(np.log2(np.max(np.take_along_axis(test_data.beam_output_true, test_preds_idx, axis=1)[:, -1-i:], axis=1) + 1.0))/\\\n",
    "                                   np.sum(np.log2(np.max(test_data.beam_output_true, axis=1) + 1.0))\n",
    "        \n",
    "      sio.savemat('federated_accuracy'+str(round_num)+'.mat',{'accuracy':top_k})\n",
    "      sio.savemat('federated_throughput'+str(round_num)+'.mat',{'throughput':throughput_ratio_at_k})\n",
    "    \n",
    "    sio.savemat('top1.mat',{'top1':top1})\n",
    "    sio.savemat('top10.mat',{'top10':top10})\n",
    "    \n",
    "    np.savez(\"federated.npz\", classification=top_k, throughput_ratio=throughput_ratio_at_k)\n",
    "    accFL=accFL+metrics['train']['top_10_categorical_accuracy']/MONTECARLO\n",
    "    \n",
    "    print(MONTECARLOi)\n",
    "    \n",
    "print(accFL)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e85753",
   "metadata": {},
   "source": [
    "## Authors\n",
    "\n",
    "- [Ailton Oliveira - LASSE/UFPA](https://github.com/AiltonOliveir)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.14 ('ai6g-xssCGMI2')",
   "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.14"
  },
  "vscode": {
   "interpreter": {
    "hash": "dac2514730ec6dd4fc1d04cb6927967d7376870f25ae7edf50b318831380a3e7"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}