variational_autoencoder_spark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "258ef4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "KafkaServerIP = dbutils.widgets.get(\"ip\")\n",
    "KafkaServerPort = dbutils.widgets.get(\"port\")\n",
    "srcTopicName = dbutils.widgets.get(\"topicname\")\n",
    "\n",
    "frameSize = [int(i) for i in dbutils.widgets.get(\"frameSize\").split(\":\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "39757723",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.streaming import StreamingContext\n",
    "from pyspark.streaming.kafka import KafkaUtils\n",
    "from pyspark.sql import Row, SparkSession\n",
    "\n",
    "from keras.models import model_from_json\n",
    "from keras.models import Sequential\n",
    "from keras.layers import Conv2D, Dropout, Merge, Dense\n",
    "from keras.layers.embeddings import Embedding\n",
    "from keras.callbacks import TensorBoard\n",
    "from keras.utils.np_utils import to_categorical\n",
    "from keras import optimizers\n",
    "from keras.layers import advanced_activations\n",
    "from keras import initializers\n",
    "from keras.layers.pooling import GlobalMaxPooling2D\n",
    "\n",
    "from tensorflow.keras.models import Model, load_model\n",
    "from tensorflow.keras.layers import Lambda, Input, Dense, Conv2D, Flatten, Reshape, Conv2DTranspose\n",
    "from tensorflow.keras.losses import mse, binary_crossentropy, kl_divergence\n",
    "from tensorflow.keras import optimizers\n",
    "from tensorflow.keras import backend as K\n",
    "\n",
    "from tensorflow.python.framework.ops import disable_eager_execution\n",
    "disable_eager_execution()\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "70955063",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This spark session isn't needed in DATABRICKS,\n",
    "uncomment the next line if you do not wish to\n",
    "execute the notebook in Databricks\n",
    "\"\"\"\n",
    "#spark = (SparkSession.builder.master(\"local[8]\").config(\"spark.driver.cores\", 8).appName(\"KerasStream\").getOrCreate() )\n",
    "\n",
    "\n",
    "sc = spark.sparkContext\n",
    "ssc = StreamingContext(sc, 3)\n",
    "\n",
    "kafkaStream = KafkaUtils.createDirectStream(ssc, [f'{srcTopicName}'], {\"metadata.broker.list\": f'{KafkaServerIP}:{KafkaServerPort}'})\n",
    "\n",
    "def getSparkSessionInstance(sparkConf):\n",
    "    if (\"sparkSessionSingletonInstance\" not in globals()):\n",
    "        globals()[\"sparkSessionSingletonInstance\"] = SparkSession \\\n",
    "            .builder \\\n",
    "            .config(conf=sparkConf) \\\n",
    "            .getOrCreate()\n",
    "    return globals()[\"sparkSessionSingletonInstance\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1cf2ac17",
   "metadata": {},
   "outputs": [],
   "source": [
    "#computes latent space parameters\n",
    "def latent(args):\n",
    "    z_mean, z_log_var = args\n",
    "    batch = K.shape(z_mean)[0]\n",
    "    dim = K.int_shape(z_mean)[1]\n",
    "    epsilon = K.random_normal(shape=(batch, dim))\n",
    "    return z_mean + K.exp(0.5 * z_log_var) * epsilon\n",
    "\n",
    "\n",
    "# computes the reconstruction loss\n",
    "def total_loss(true, pred):\n",
    "    reconstruction_loss = binary_crossentropy(K.flatten(true), K.flatten(pred)) * reshaped_train.shape[1] * reshaped_train.shape[2]   # KL divergence loss\n",
    "    # KL loss\n",
    "    kl_loss = - 0.5 * K.sum(1 + sigma - K.square(mu) - K.square(K.exp(sigma)), axis=-1)\n",
    "    total = K.mean(reconstruction_loss + kl_loss)    \n",
    "    return total\n",
    "\n",
    "#computes rmse \n",
    "def RMSE(v1, v2):\n",
    "    return np.sqrt(np.mean((v1 - v2) ** 2, axis=1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "365cd27f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vae_model():\n",
    "    \n",
    "    input_shape = (30 , frameSize[0], frameSize[1])\n",
    "    latent_dim = 2\n",
    "    # Encoder model\n",
    "    encoder_inputs = Input(shape=input_shape, name='encoder_input')\n",
    "\n",
    "    encoder_conv = Conv2D(filters=8, kernel_size=3, strides=2, \\\n",
    "                    padding='same',activation='relu')(encoder_inputs)\n",
    "    encoder_conv = Conv2D(filters=16, kernel_size=3, strides=2, \\\n",
    "                    padding='same',activation='relu')(encoder_inputs)\n",
    "    encoder = Flatten()(encoder_conv)\n",
    "\n",
    "    mu = Dense(latent_dim, name='mu')(encoder)\n",
    "    sigma = Dense(latent_dim, name='sigma')(encoder)\n",
    "    latent_space = Lambda(latent, output_shape=(latent_dim,))([mu, sigma])\n",
    "    encoder = Model(encoder_inputs, latent_space, name='encoder')\n",
    "    \n",
    "    conv_shape = K.int_shape(encoder_conv)\n",
    "    # decoder model\n",
    "    inputs = Input(shape=(latent_dim,))\n",
    "\n",
    "    decoder = Dense(conv_shape[1]*conv_shape[2]*conv_shape[3], activation='relu')(inputs)\n",
    "    decoder = Reshape((conv_shape[1], conv_shape[2], conv_shape[3]))(decoder)\n",
    "    decoder_conv = Conv2DTranspose(filters=16, kernel_size=3, strides=2, \n",
    "                               padding='same',activation='relu')(decoder)\n",
    "    decoder_conv = Conv2DTranspose(filters=8, kernel_size=3, strides=2, \n",
    "                               padding='same',activation='relu')(decoder)\n",
    "    decoder_conv =  Conv2DTranspose(filters=reshaped_train.shape[3], kernel_size=3, \n",
    "                              padding='same',activation='sigmoid')(decoder_conv)\n",
    "\n",
    "    decoder = Model(inputs, decoder_conv, name='decoder')\n",
    "    #combines the encodr and decoder to define the model\n",
    "    outputs = decoder(encoder(encoder_inputs))\n",
    "    vae = Model(encoder_inputs, outputs, name='vae_model')\n",
    "    \n",
    "    opt = optimizers.Adam(learning_rate=0.001)\n",
    "    vae.compile(optimizer=opt, loss=total_loss)\n",
    "    return vae"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c3dfbf7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "core_model = vae_model()\n",
    "model_json = core_model.to_json()\n",
    "with open(\"model.json\", \"w\") as f:\n",
    "    f.write(model_json)\n",
    "    \n",
    "\n",
    "broadcast_model = sc.broadcast(model_json)\n",
    "\n",
    "def learn_infer(data):\n",
    "    loaded_model = model_from_json(broadcast_model.value)\n",
    "    loaded_model.load_weights(\"./myWeights.h5\")\n",
    "    \n",
    "\n",
    "instances = kafkaStream.map(lambda x: x[1])\n",
    "\n",
    "def myMap(time, rdd):\n",
    "    spark = getSparkSessionInstance(rdd.context.getConf())\n",
    "\n",
    "    rowRdd = rdd.map(lambda w: Row(sen=learn_infer(w)))\n",
    "    dfrows = spark.createDataFrame(rowRdd)\n",
    "    dfrows.show()\n",
    "\n",
    "\n",
    "\n",
    "instances.foreachRDD(myMap)\n",
    "ssc.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7d76128d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dbutils.notebook.exit(\"The distributed training has started successfully\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}