AE_Classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Classification Code\n",
    "\n",
    "## Should be run after the weights determined from the Transformer autoencoder.\n",
    "\n",
    "Author: Kara Ponder (SLAC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "from transformer import Encoder, Decoder, create_decoder_masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set parameters\n",
    "d_model = 128 # input vector must have length d_model\n",
    "target_vocab_size = 6  # possible results to choose from\n",
    "\n",
    "lc_length = 50 + 1 # light curve length\n",
    "input_vocab_size = lc_length\n",
    "\n",
    "## hyperparameters:\n",
    "num_layers = 8\n",
    "dropout_rate = 0.0\n",
    "dff = 64 # hidden layer size of the feed forward network, needs to be larger than 24\n",
    "num_heads = 8 # d_model % num_heads == 0 \n",
    "\n",
    "# LC stuff\n",
    "N = 500 # number of objects\n",
    "N_days = 100 + 1\n",
    "Nf = 6 # number of filters\n",
    "num_class = 4 # set to 4 for the test data. Will be more for PLAsTiCC\n",
    "\n",
    "batch_size=64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function does the classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_ffn(nclass, dff, rate=0.0):\n",
    "    model = tf.keras.models.Sequential()\n",
    "    model.add(tf.keras.layers.Dense(dff, activation='relu')) \n",
    "    model.add(tf.keras.layers.Dropout(rate))\n",
    "    model.add(tf.keras.layers.Dense(dff, activation='relu'))\n",
    "    model.add(tf.keras.layers.GlobalMaxPool1D())\n",
    "    model.add(tf.keras.layers.Dense(nclass, activation='softmax'))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize transformer model in order to extract the encoder layers that will be used to encode before classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = Encoder(num_layers, d_model, num_heads, dff,\n",
    "                       lc_length, dropout_rate, embed=True)\n",
    "\n",
    "decoder = Decoder(num_layers, d_model, num_heads, dff,\n",
    "                       lc_length, dropout_rate, embed=True)\n",
    "\n",
    "final_layer = tf.keras.layers.Dense(target_vocab_size)\n",
    "\n",
    "inp = tf.keras.layers.Input(shape=(None,6))#shape=(None,None))#\n",
    "target = tf.keras.layers.Input(shape=(None,6))#shape=(None,None))#\n",
    "\n",
    "x = encoder(inp)\n",
    "x = decoder(target, x, mask=create_decoder_masks(inp, target))\n",
    "\n",
    "x = final_layer(x)\n",
    "\n",
    "model = tf.keras.models.Model(inputs=[inp, target], outputs=x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the previous weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('TRANSFORMER_WEIGHTS.h5')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`extras` corresponds to the MIN/MAX normalization constants per objects that are saved when preprocessing the lightcurves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extras = tf.keras.layers.Input(shape=(None, 2), dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the classifier by taking the Transformer encoder layer and not training those weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_inp = tf.keras.layers.Input(shape=(None,6))\n",
    "\n",
    "classify_encoder = Encoder(num_layers, d_model, num_heads, dff,\n",
    "                       lc_length, dropout_rate, embed=True)\n",
    "classify_encoder(cl_inp)\n",
    "classify_encoder.set_weights(model.get_layer(name='encoder').get_weights())\n",
    "\n",
    "classify_encoder.trainable = False\n",
    "\n",
    "class_ffn = classify_ffn(num_class, dff, rate=dropout_rate)\n",
    "\n",
    "cl_x = classify_encoder(inp)\n",
    "\n",
    "# Add in normalization constants\n",
    "cl_x = tf.keras.layers.Concatenate(axis=-1)([cl_x, extras])\n",
    "cl_x = class_ffn(cl_x)\n",
    "\n",
    "aeclass = tf.keras.models.Model(inputs=[inp, extras], outputs=cl_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define optimizer and compile the AEClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(0.0001)\n",
    "\n",
    "loss_object = tf.keras.losses.CategoricalCrossentropy(\n",
    "\n",
    "train_loss = tf.keras.metrics.Mean(name='train_loss')\n",
    "\n",
    "train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')\n",
    "\n",
    "aeclass.compile(loss=loss_object,\n",
    "                optimizer=optimizer, \n",
    "                metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load simplified data for illustration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = np.load('label.npy')\n",
    "lc_data = np.load('lc_data.npy')\n",
    "norm_data = np.loadtxt('min_max.txt')\n",
    "\n",
    "labels = tf.keras.utils.to_categorical(label, num_classes=num_class, dtype=\"float64\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = tf.data.Dataset.from_tensor_slices((lc_data, labels, norm_data))\n",
    "batch_ds = dataset.batch(batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define number of batch steps and "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_batches = 0\n",
    "for (batch, _) in enumerate(batch_ds):\n",
    "    num_batches = batch\n",
    "\n",
    "def generator(data_set):\n",
    "    while True:\n",
    "        for in_batch, tar_batch, norm_batch in data_set:\n",
    "            yield ( [in_batch, norm_batch] , tar_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit for the AEClassifier weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history = aeclass.fit(x = generator(batch_ds),\n",
    "                    #validation_data = generator(batch_ds_VALID),\n",
    "                    epochs=15,\n",
    "                    steps_per_epoch = num_batches,\n",
    "                    #validation_steps = num_batches_VALID,\n",
    "                    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#aeclass.save_weights('save_aeclass_weights.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}