train-RAEwALL-clf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from keras import losses\n",
    "import keras\n",
    "\n",
    "from src.data import get_generators, format_bytes\n",
    "from src.model import get_2losses_model, add_clf_layer, add_clf_layer_sparse, get_model, get_RAEwSC_compiled, get_RAEwSC_and_WS_compiled\n",
    "from src.losses import cross_entropy_with_axis, model_evaluate, write_submission_files\n",
    "\n",
    "from keras.utils.vis_utils import plot_model\n",
    "# plot_model(model, show_shapes=True, show_layer_names=True)\n",
    "# start tensorboard: tensorboard --logdir ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define global names\n",
    "EXTRA_DATA_MODEL, EXTRA_DATA_MODELwIN, RAEwSCwWSwINwCLF = \"RAE_w_SC_WS\", \"RAE_w_SC_WS_wIN\", \"RAEwSCwWSwINwCLF\"\n",
    "\n",
    "# define model and city\n",
    "MODEL_TYPES = {0: \"ConvLSTM\", 1: \"ConvLSTM+Clf\", 2: \"RAE_w_SC\", 3: EXTRA_DATA_MODEL, 4: EXTRA_DATA_MODELwIN, \n",
    "               5: RAEwSCwWSwINwCLF}\n",
    "CITIES = {0: 'Moscow', 1: 'Istanbul', 2: 'Berlin'}\n",
    "DATA_SPLIT = {0: \"non-overlapping\", 1: \"all_possible_slots\", 2: \"like-test\"}\n",
    "\n",
    "############################################# DEFINE HYPERPARAMETERS\n",
    "DATA_SPLIT_IDX = 2 # choose how to split data\n",
    "params = {\n",
    "    # data params\n",
    "    \"batch_size\": 4,\n",
    "    \"length_seq_in\": 3, #3, # Set up when calling a model\n",
    "    \"length_seq_out\": 3,\n",
    "    \"batch_size_validation\": 48,\n",
    "    \"data_split\": DATA_SPLIT[DATA_SPLIT_IDX], # Set up when calling a model\n",
    "    \"batch_size_test\": 5,\n",
    "    # loss params\n",
    "    'optimizer': 'adam',  # adam or sgd\n",
    "    'learning_rate': 0.001,\n",
    "    'decay_rate': 0.9999,  # Learning rate decay per minibatch.\n",
    "    'min_learning_rate': .0000001,  # Minimum learning rate.\n",
    "}\n",
    "\n",
    "############################################# DEFINE MODEL WEIGHTS\n",
    "models_weight = {\n",
    "    \"ConvLSTM\": { # _weights_e04_valoss0.009_relu \n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_model_ep_2.h5\".format(\"Moscow\"), \n",
    "        \"Istanbul\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_weights_e03_valoss0.009_relu3.hdf5\".format(\"Istanbul\"), \n",
    "        \"Berlin\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_model_ep_5.h5\".format(\"Berlin\"),\n",
    "        \"Model_definition_file\": \"{}/projects/nips_traffic/src/models.py\".format(\"/home/pherruzo\"),\n",
    "        \"Model_definition_function\": \"build_model\" \n",
    "    },\n",
    "    \n",
    "    \"ConvLSTM+Clf\": { # _weights_e04_valoss0.025_relu.hdf5, _weights_e04_valoss0.025_relu2.hdf\n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_weights_e04_valoss0.025_relu2.hdf5\".format(\"Moscow\"),\n",
    "        \"Istanbul\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_weights_06-1.04.hdf5\".format(\"Istanbul\"), # weights_Istanbul_06-1.04.hdf5, weights_Istanbul_04-1.05.hdf5, weights_Istanbul_03-1.05.hdf5  \n",
    "        \"Berlin\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_model_ep_1_reg_clf_cont.h5\".format(\"Berlin\")\n",
    "    },\n",
    "    \n",
    "    \"RAE_w_SC\": { # best: _weights_e03_valoss0.012_raewsc_sgd | prev: _weights_e02_valoss0.014_raewsc_sgd_emb\n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_weights_e03_valoss0.012_raewsc_sgd.hdf5\".format(\"Moscow\"),\n",
    "        \"Istanbul\": None, \n",
    "        \"Berlin\": None\n",
    "    },\n",
    "    \n",
    "    EXTRA_DATA_MODEL: { # best: _e02_valoss0.012_exo_vars, best in test-like: _e91_valoss0.012_just_try \n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_e91_valoss0.012_just_try.hdf5\".format(\"Moscow\"), # works good for num_seq_in=3\n",
    "        \"Istanbul\": None,\n",
    "        \"Berlin\": None\n",
    "    },\n",
    "    \n",
    "    EXTRA_DATA_MODELwIN: { # best (all slots): _e04_valoss0.012_RAEwSCwWTwIN_new_all_slots, # like-test: _e18_valoss0.012_RAEwSCwWTwINz_all_slots_like_test\n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_e18_valoss0.012_RAEwSCwWTwINz_all_slots_like_test.hdf5\".format(\"Moscow\"),\n",
    "        \"Istanbul\": None,\n",
    "        \"Berlin\": None\n",
    "    },\n",
    "    \n",
    "    RAEwSCwWSwINwCLF: { # best (all slots): _e19_valoss0.502_RAEwSCwWSwINwCLF #  _e12_valoss1.294_RAEwSCwWSwINwCLF\n",
    "        \"Moscow\": \"/home/pherruzo/projects/nips_traffic/checkpoints/{}_e44_valoss0.026_RAEwSCwWSwINwCLF.hdf5\".format(\"Moscow\"),\n",
    "        \"Istanbul\": None,\n",
    "        \"Berlin\": None,\n",
    "        \"loss_type\": [None, 'Sparse'][1]#'Categorical'\n",
    "    }\n",
    "}\n",
    "\n",
    "############################################# PREPARE POSSIBLE MODELS\n",
    "def get_a_model(model_type, training_params):\n",
    "    \"\"\" model_type in [\"ConvLSTM+Clf\", \"ConvLSTM\", \"RAE_w_SC\"] \"\"\"\n",
    "    \n",
    "    if model_type==\"ConvLSTM+Clf\":\n",
    "        return get_2losses_model(model_path=models_weight[training_params[\"previus_model\"]][\"Model_definition_file\"], \n",
    "                                model_name=models_weight[training_params[\"previus_model\"]][\"Model_definition_function\"], \n",
    "                                weights_initial_model=None, \n",
    "                                weights_current_model=models_weight[training_params[\"current_model\"]][training_params[\"city\"]], \n",
    "                                add_model=add_clf_layer_sparse, \n",
    "                                add_loss=losses.sparse_categorical_crossentropy,\n",
    "                                sample_weight_mode=\"temporal\")\n",
    "    \n",
    "    elif model_type==\"ConvLSTM\":\n",
    "        from keras import optimizers\n",
    "        return get_model(model_path=models_weight[training_params[\"previus_model\"]][\"Model_definition_file\"], \n",
    "                         function_name=models_weight[training_params[\"previus_model\"]][\"Model_definition_function\"], \n",
    "                         weights=models_weight[training_params[\"previus_model\"]][training_params[\"city\"]], \n",
    "                         opt=optimizers.Adam(lr=0.0001), loss=losses.mean_squared_error)\n",
    "    \n",
    "    elif model_type==\"RAE_w_SC\":\n",
    "        return get_RAEwSC_compiled(weights=models_weight[model_type][training_params[\"city\"]], \n",
    "                                   lr=params['learning_rate'])\n",
    "    elif model_type==EXTRA_DATA_MODEL:\n",
    "        dropout_enc, b_norm_enc = 0.04, True\n",
    "        dropout_dec, b_norm_dec = 0.03, True\n",
    "        print(\"dropout_enc:\", dropout_enc, \"dropout_dec:\", dropout_dec)\n",
    "        # get basic model def get_RAEwSC_and_WS_compiled(weights=None, lr = 0.001, grad_clip=1., loss_weights={'predicted_frames':1., 'predicted_emb':1.}):\n",
    "        return get_RAEwSC_and_WS_compiled(weights=models_weight[model_type][training_params[\"city\"]], lr=params['learning_rate'],\n",
    "                                          loss_weights={'predicted_frames':1., 'predicted_emb':1.},\n",
    "                                           dropout_enc=dropout_enc, dropout_dec=dropout_dec, \n",
    "                                           b_norm_enc=b_norm_enc, b_norm_dec=b_norm_dec)\n",
    "    elif model_type==EXTRA_DATA_MODELwIN:\n",
    "        from src.models.RAEwSCwWSwIN import get_RAEwSC_and_WS_compiled_w_input\n",
    "        \n",
    "        loss_w = {'predicted_frames':1., 'predicted_emb':0.9}\n",
    "        \n",
    "        dropout_enc = 0.05\n",
    "        dropout_dec = 0.03\n",
    "        \n",
    "        print(\"loading {} ...\".format(model_type))\n",
    "        print(\"dropout_enc:\", dropout_enc, \"dropout_dec:\", dropout_dec)\n",
    "        print(\"loss weights:\", loss_w)\n",
    "        return get_RAEwSC_and_WS_compiled_w_input(weights=models_weight[model_type][training_params[\"city\"]], length_seq_in=params['length_seq_in'], \n",
    "                                                  lr=params['learning_rate'], loss_weights=loss_w,\n",
    "                                                  dropout_enc=dropout_enc, dropout_dec=dropout_dec, for_new_model=True) ######### This is True only to create the new model\n",
    "    \n",
    "    elif model_type==RAEwSCwWSwINwCLF:\n",
    "        from src.models.RAEwSCwWSwINxCLF import get_RAEwSCwWSwIwCLF_compiled as get_model\n",
    "        \n",
    "        loss_w = {'predicted_frames':0., 'predicted_emb':1.}\n",
    "        \n",
    "        dropout_enc = 0.05\n",
    "        dropout_dec = 0.03\n",
    "        \n",
    "        print(\"loading {} ...\".format(model_type))\n",
    "        print(\"dropout_enc:\", dropout_enc, \"dropout_dec:\", dropout_dec)\n",
    "        print(\"loss weights:\", loss_w)\n",
    "        return get_model(weights=models_weight[model_type][training_params[\"city\"]], length_seq_in=params['length_seq_in'], \n",
    "                         lr=params['learning_rate'], loss_weights=loss_w,\n",
    "                         dropout_enc=dropout_enc, dropout_dec=dropout_dec)\n",
    "    # get_RAEwSCwWSwIwCLF_compiled\n",
    "\n",
    "############################################# DEFINE TRAINING PARAMS\n",
    "def get_training_params(model_type, city):\n",
    "    \"\"\" model_type in [\"ConvLSTM\", \"ConvLSTM+Clf\", \"RAE_w_SC\", \"RAE_w_SC_WS\"] \n",
    "        city in ['Moscow', 'Istanbul', 'Berlin']\n",
    "    \"\"\"\n",
    "    training_params = {\n",
    "        \"city\": city,\n",
    "        \"previus_model\": \"ConvLSTM\",     # used in old methods, NOT TO CHANGE todo: replace\n",
    "        \"current_model\": \"ConvLSTM+Clf\", # used in old methods, NOT TO CHANGE todo: replace\n",
    "        \"tensorboard_path\": \"/home/pherruzo/projects/nips_traffic/tensorboard/\",\n",
    "        \"log_path\": \"/home/pherruzo/projects/nips_traffic/log_files/\",\n",
    "        \"model_type\": model_type,\n",
    "        \"output_dir\": \"/home/pherruzo/data/nips_traffic_submissions/conv_clf_mixt\"\n",
    "    }\n",
    "    return training_params\n",
    "\n",
    "############################################# LOAD DATA GENERATOR\n",
    "def load_model_and_data(model_type, city, params, length_seq_in=None, data_split=None, debug=False):\n",
    "    \n",
    "    if length_seq_in is not None:\n",
    "        params[\"length_seq_in\"] = length_seq_in\n",
    "        \n",
    "    if data_split is not None:\n",
    "        params[\"data_split\"] = DATA_SPLIT[data_split]\n",
    "        \n",
    "    # load training params\n",
    "    training_params = get_training_params(model_type, city)\n",
    "    \n",
    "    # load data\n",
    "    if \"loss_type\" in models_weight[model_type]:\n",
    "        loss_type = models_weight[model_type][\"loss_type\"]\n",
    "        print(\"---> loss type for clf:\", loss_type)\n",
    "    else:\n",
    "        loss_type = None\n",
    "    training_ds, val_ds, test_ds = get_generators(training_params[\"model_type\"],   params[\"batch_size\"], \n",
    "                                                  training_params[\"city\"],         params[\"length_seq_in\"], \n",
    "                                                  params[\"batch_size_validation\"], data_split=params[\"data_split\"], debug=debug, loss_type=loss_type)\n",
    "    # load model\n",
    "    model = get_a_model(model_type, training_params)\n",
    "    \n",
    "    # print info\n",
    "    print(\"\")\n",
    "    print(\"· length_seq_in:\", params[\"length_seq_in\"])\n",
    "    print(\"· data_split:\", params[\"data_split\"])\n",
    "    print(\"· model_type:\", training_params[\"model_type\"])\n",
    "    print(\"· city      :\", training_params[\"city\"])\n",
    "    print(\" 1 epoch training is {} batches, validation is {}, and test is {}.\".format(len(training_ds), len(val_ds), len(test_ds)))\n",
    "    \n",
    "    return training_ds, val_ds, test_ds, model, training_params\n",
    "\n",
    "def plot(keys_list, history):\n",
    "\n",
    "    leg = []\n",
    "    for k in keys_list:\n",
    "        plt.plot(history.history[k], '-o')\n",
    "        leg.append(k)\n",
    "    \n",
    "    plt.title(leg)\n",
    "    plt.xlabel('epoch')\n",
    "    plt.legend(leg, loc='center')#'upper left')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the new model\n",
    "\n",
    "### 1. Load model RAE_w_SC_WS_wIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-->> Data has been shuffled in: training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W1022 09:13:58.139183 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:517: The name tf.placeholder is deprecated. Please use tf.compat.v1.placeholder instead.\n",
      "\n",
      "W1022 09:13:58.147670 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:74: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
      "\n",
      "W1022 09:13:58.150118 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:4185: The name tf.truncated_normal is deprecated. Please use tf.random.truncated_normal instead.\n",
      "\n",
      "W1022 09:13:58.166394 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:174: The name tf.get_default_session is deprecated. Please use tf.compat.v1.get_default_session instead.\n",
      "\n",
      "W1022 09:13:58.166708 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:190: The name tf.global_variables is deprecated. Please use tf.compat.v1.global_variables instead.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading RAE_w_SC_WS_wIN ...\n",
      "dropout_enc: 0.05 dropout_dec: 0.03\n",
      "loss weights: {'predicted_frames': 1.0, 'predicted_emb': 0.9}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W1022 09:13:58.648825 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:1834: The name tf.nn.fused_batch_norm is deprecated. Please use tf.compat.v1.nn.fused_batch_norm instead.\n",
      "\n",
      "W1022 09:13:58.723539 140354142054208 deprecation_wrapper.py:119] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:3976: The name tf.nn.max_pool is deprecated. Please use tf.nn.max_pool2d instead.\n",
      "\n",
      "W1022 09:13:58.726788 140354142054208 deprecation.py:506] From /home/pherruzo/anaconda3/envs/nips/lib/python3.7/site-packages/keras/backend/tensorflow_backend.py:3445: calling dropout (from tensorflow.python.ops.nn_ops) with keep_prob is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `rate` instead of `keep_prob`. Rate should be set to `rate = 1 - keep_prob`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoder output shape: (?, 8, 8, 32) reshaped: (?, ?)\n",
      "Decoder input shape: (?, 2048) reshaped: (?, 8, 8, 32)\n",
      "embeddings.shape: (?, 6, 2048)\n",
      "future_embeddings.shape: (?, 3, 2048)\n",
      "concatenation of all inputs: (?, 6, 2388)\n",
      "FC before recurrent embeddings.shape: (?, 6, 2048)\n",
      "recurrent embeddings.shape: (?, ?, 2048)\n",
      "prediced_frames.shape: (?, 3, 495, 436, 3)\n",
      "loading weitghs: /home/pherruzo/projects/nips_traffic/checkpoints/Moscow_e18_valoss0.012_RAEwSCwWTwINz_all_slots_like_test.hdf5\n",
      "Freezing layers...\n",
      "\n",
      "· length_seq_in: 6\n",
      "· data_split: like-test\n",
      "· model_type: RAE_w_SC_WS_wIN\n",
      "· city      : Moscow\n",
      " 1 epoch training is 357 batches, validation is 5, and test is 72.\n"
     ]
    }
   ],
   "source": [
    "#### load data and model --> \n",
    "# choose model          {0: \"ConvLSTM\", 1: \"ConvLSTM+Clf\", 2: \"RAE_w_SC\", 3: \"RAE_w_SC_WS\", 4: \"RAE_w_SC_WS_wIN\", 5: \"RAEwSCwWSwINwCLF\"}\n",
    "# choose city           {0: 'Moscow', 1: 'Istanbul', 2: 'Berlin'}\n",
    "# choose data-split     {0: \"non-overlapping\", 1: \"all_possible_slots\", 2: \"like-test\"}\n",
    "# choose length_seq_in  3, 4, 5, 6, 7, 8, 9, 10, 11, 12\n",
    "\n",
    "params['learning_rate'] = 0.01\n",
    "debug = False\n",
    "\n",
    "m_type, city = 4, 0\n",
    "data_split_idx = 2\n",
    "length_seq_in = 6\n",
    "\n",
    "training_ds, val_ds, test_ds, model, training_params = load_model_and_data(MODEL_TYPES[m_type], CITIES[city], params, \n",
    "                                                                           length_seq_in=length_seq_in, data_split=data_split_idx, debug=debug)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Freezing Layers & Compiling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "###################################### change layer name\n",
    "# model.get_layer(\"Concat_predicted_frames\").name = \"Concat_predicted_frames_2\"\n",
    "\n",
    "###################################### freeze layers\n",
    "# once we load the weight we freeze all layers\n",
    "for i, layer in enumerate(model.layers):\n",
    "#     if 'norm'  in layer.name:\n",
    "        layer.trainable = False\n",
    "\n",
    "# freeze also nested layers\n",
    "for i, layer in enumerate(model.get_layer(\"encoder\").layers):\n",
    "#     if 'norm' not in layer.name:\n",
    "        layer.trainable = False\n",
    "        \n",
    "for i, layer in enumerate(model.get_layer(\"decoder\").layers):\n",
    "#     if 'norm' not in layer.name:\n",
    "        layer.trainable = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lr: 0.001 optimizer: <keras.optimizers.Adam object at 0x7fa64ea56908>\n"
     ]
    }
   ],
   "source": [
    "from keras import models, activations\n",
    "from keras import layers\n",
    "from keras import optimizers\n",
    "from keras import losses\n",
    "\n",
    "# compile with a custom loss that handle both regression and classification\n",
    "lr = 0.001\n",
    "grad_clip=1.\n",
    "optimizer = [optimizers.SGD(lr=lr, decay=1e-6, momentum=0.9, nesterov=True), optimizers.Adam(lr=lr, clipvalue=grad_clip)][-1]\n",
    "print(\"lr:\", lr, 'optimizer:', optimizer)\n",
    "    \n",
    "model.compile(optimizer=optimizer,\n",
    "              loss=[losses.mean_squared_error],\n",
    "              metrics=[losses.mean_squared_error])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.01236429757305554, 0.012364300579896995]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check evaluation\n",
    "model.evaluate_generator(val_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Building new layers on top of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lr: (0.1,) optimizer: <keras.optimizers.Adam object at 0x7f9ee9350a90>\n",
      "---> loss type for clf: Sparse\n",
      "-->> Data has been shuffled in: training\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "sess = tf.compat.v1.Session()\n",
    "\n",
    "from keras import backend as K\n",
    "K.set_session(sess)\n",
    "\n",
    "from keras import models, activations\n",
    "from keras import layers\n",
    "from keras import optimizers\n",
    "from keras import losses\n",
    "# from keras import backend as K \n",
    "from tensorflow import device\n",
    "\n",
    "from keras.layers.convolutional import Convolution2D, MaxPooling2D\n",
    "from keras.utils.np_utils import to_categorical  \n",
    "\n",
    "# imports for RAE_w_SC\n",
    "from keras import layers\n",
    "from keras.layers import Input, BatchNormalization, Activation, Dense, Dropout, ZeroPadding2D, Cropping2D, TimeDistributed\n",
    "from keras.layers.convolutional import Conv2D, Conv2DTranspose\n",
    "from keras.layers.pooling import MaxPooling2D\n",
    "from keras.layers import concatenate\n",
    "from keras.optimizers import Adam\n",
    "from keras.models import Model, load_model\n",
    "from keras.layers.core import Lambda, RepeatVector, Reshape\n",
    "from keras import losses, optimizers\n",
    "\n",
    "# own libraries\n",
    "from src.losses import *\n",
    "\n",
    "ACTIVATION = ['relu', 'elu'][-1]\n",
    "\n",
    "# ###################################### change layer name\n",
    "# model.get_layer(\"Concat_predicted_frames\").name = \"Concat_predicted_frames_2\"\n",
    "\n",
    "# ###################################### freeze layers\n",
    "# # once we load the weight we freeze all layers\n",
    "# for i, layer in enumerate(model.layers):\n",
    "#     if 'norm' not in layer.name:\n",
    "#         layer.trainable = False\n",
    "\n",
    "# # freeze also nested layers\n",
    "# for i, layer in enumerate(model.get_layer(\"encoder\").layers):\n",
    "#     if 'norm' not in layer.name:\n",
    "#         layer.trainable = False\n",
    "        \n",
    "# for i, layer in enumerate(model.get_layer(\"decoder\").layers):\n",
    "#     if 'norm' not in layer.name:\n",
    "#         layer.trainable = False\n",
    "        \n",
    "# get needed layers (batch, seq_out, width, height, channel)\n",
    "new_frames = model.output\n",
    "\n",
    "###################################### DEFINE NEW MODEL\n",
    "out_clf = 5\n",
    "channels_out = 3\n",
    "################################# add clf\n",
    "name = 'clf'\n",
    "n_filters = out_clf\n",
    "kernel_size = 3\n",
    "# first layer\n",
    "x = TimeDistributed( Conv2D(filters=n_filters, kernel_size=(kernel_size, kernel_size), kernel_initializer=\"he_normal\", padding=\"same\", name=name+'_conv1') ) (new_frames)\n",
    "x = TimeDistributed( BatchNormalization(trainable=True, name=name+'_BN1') ) (x)\n",
    "x = TimeDistributed( Activation(ACTIVATION) ) (x)\n",
    "# second layer\n",
    "x = TimeDistributed( Conv2D(filters=n_filters, kernel_size=(kernel_size, kernel_size), kernel_initializer=\"he_normal\", padding=\"same\", name=name+'_conv2') )(x)\n",
    "x = TimeDistributed( BatchNormalization(trainable=True, name=name+'_BN2') ) (x)\n",
    "x = TimeDistributed( Activation(ACTIVATION) ) (x)\n",
    "\n",
    "# prepare the clfs for the loss\n",
    "clfs = layers.Reshape((-1, out_clf))(x)      # vectorize all frames\n",
    "clfs = layers.Softmax(name='softmax_clf')(clfs) # apply softmax\n",
    "\n",
    "################################# add clf to decoder output\n",
    "\n",
    "# concat decoder output with clf\n",
    "concat_0 = concatenate([new_frames, x])\n",
    "concat = TimeDistributed( BatchNormalization(trainable=True, name='concat_1_bn') ) (concat_0)\n",
    "concat_1 = TimeDistributed( Activation(ACTIVATION) ) (concat)\n",
    "\n",
    "# spatial branch\n",
    "# conv 1\n",
    "c1 = TimeDistributed( Conv2D(channels_out*3, (kernel_size, kernel_size), activation='elu', padding=\"same\"), name='c1_1' ) (concat_1)\n",
    "c1 = TimeDistributed( BatchNormalization(trainable=True, name='bn_c1_1') ) (c1)\n",
    "c1 = TimeDistributed( Activation(ACTIVATION) ) (c1)\n",
    "# conv 2\n",
    "c1 = TimeDistributed( Conv2D(channels_out*2, (kernel_size, kernel_size), activation='elu', padding=\"same\"), name='c1_2' ) (c1)\n",
    "c1 = TimeDistributed( BatchNormalization(trainable=True, name='bn_c1_1') ) (c1)\n",
    "c1 = TimeDistributed( Activation(ACTIVATION) ) (c1)\n",
    "\n",
    "# depth branch\n",
    "# conv 1\n",
    "c2 = TimeDistributed( Conv2D(channels_out*3, (1, 1), activation='elu', padding=\"same\"), name='c2_1' ) (concat_1)\n",
    "c2 = TimeDistributed( BatchNormalization(trainable=True, name='bn_c2_1') ) (c2)\n",
    "c2 = TimeDistributed( Activation(ACTIVATION) ) (c2)\n",
    "# conv 2\n",
    "c2 = TimeDistributed( Conv2D(channels_out*2, (1, 1), activation='elu', padding=\"same\"), name='c2_2' ) (concat_1)\n",
    "c2 = TimeDistributed( BatchNormalization(trainable=True, name='bn_c2_2') ) (c2)\n",
    "c2 = TimeDistributed( Activation(ACTIVATION) ) (c2)\n",
    "\n",
    "# concat depth convs & spatial convs\n",
    "concat = concatenate([concat_0, c1, c2])\n",
    "concat = TimeDistributed( BatchNormalization(trainable=True, name='concat_2_bn') ) (concat)\n",
    "concat = TimeDistributed( Activation(ACTIVATION) ) (concat)\n",
    "\n",
    "# final conv\n",
    "prediced_frames = TimeDistributed( Conv2D(channels_out, (1, 1), activation='elu'), name='Concat_predicted_frames' ) (concat)\n",
    "\n",
    "model = models.Model(inputs=model.input, outputs=[prediced_frames, clfs], name='final_model')\n",
    "\n",
    "###################################### COMPILE MODEL\n",
    "general_loss_weights={'Concat_predicted_frames':1., 'softmax_clf':.1}\n",
    "\n",
    "\n",
    "# compile with a custom loss that handle both regression and classification\n",
    "lr = 0.1, \n",
    "grad_clip=1.\n",
    "optimizer = [optimizers.SGD(lr=lr, decay=1e-6, momentum=0.9, nesterov=True), optimizers.Adam(lr=lr, clipvalue=grad_clip)][-1]\n",
    "print(\"lr:\", lr, 'optimizer:', optimizer)\n",
    "\n",
    "# # load weights\n",
    "# name = [ 'Moscow_e02_valoss0.025_RAEwSCwWSwINwCLF_end.hdf5', 'Moscow_e06_valoss0.024_RAEwSCwWSwINwCLF_end.hdf5'][-1]\n",
    "# weights = '/home/pherruzo/projects/nips_traffic/checkpoints/' + name\n",
    "# if weights is not None:\n",
    "#     print('loading weitghs:', weights)\n",
    "#     model.load_weights(weights, by_name=True, skip_mismatch=True) #l24)#, by_name=True)\n",
    "\n",
    "model.compile(optimizer=optimizer,\n",
    "              loss={'Concat_predicted_frames': 'mse',\n",
    "                    'softmax_clf': losses.sparse_categorical_crossentropy},\n",
    "              metrics={'Concat_predicted_frames': ['mse'], \n",
    "                       'softmax_clf':             [losses.sparse_categorical_crossentropy]},#'accuracy']},\n",
    "              loss_weights=general_loss_weights, \n",
    "              sample_weight_mode=\"temporal\")\n",
    "\n",
    "###################################### GET DATA FOR THIS MODEL\n",
    "training_params[\"model_type\"] = RAEwSCwWSwINwCLF\n",
    "if \"loss_type\" in models_weight[training_params[\"model_type\"]]:\n",
    "    loss_type = models_weight[training_params[\"model_type\"]][\"loss_type\"]\n",
    "    print(\"---> loss type for clf:\", loss_type)\n",
    "else:\n",
    "    loss_type = None\n",
    "    \n",
    "training_ds, val_ds, test_ds = get_generators(training_params[\"model_type\"],   params[\"batch_size\"], \n",
    "                                                  training_params[\"city\"],         params[\"length_seq_in\"], \n",
    "                                                  params[\"batch_size_validation\"], data_split=params[\"data_split\"], debug=debug, loss_type=loss_type)\n",
    "\n",
    "# Init vars\n",
    "keras.backend.get_session().run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.18832453872476304,\n",
       " 0.027383473834821155,\n",
       " 1.6094106197357179,\n",
       " 0.027383483733449662,\n",
       " 1.609486481121608]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "keras.backend.get_session().run(tf.global_variables_initializer())\n",
    "model.evaluate_generator(val_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.011967233182596309, 0.011967231120382036, 2.1739711776800975e-09]"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# evaluate\n",
    "model.evaluate_generator(val_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Load new weights & compile it again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lr: (0.001,) optimizer: <keras.optimizers.Adam object at 0x7f9ebaea9ac8>\n",
      "loading weitghs: /home/pherruzo/projects/nips_traffic/checkpoints/Moscow_e02_valoss0.212_RAEwSCwWSwINwCLF_frozenUnet.hdf5\n"
     ]
    }
   ],
   "source": [
    "general_loss_weights={'Concat_predicted_frames':1., 'softmax_clf':.1}\n",
    "\n",
    "\n",
    "# compile with a custom loss that handle both regression and classification\n",
    "lr = 0.001, \n",
    "grad_clip=1.\n",
    "optimizer = [optimizers.SGD(lr=lr, decay=1e-6, momentum=0.9, nesterov=True), optimizers.Adam(lr=lr, clipvalue=grad_clip)][-1]\n",
    "print(\"lr:\", lr, 'optimizer:', optimizer)\n",
    "\n",
    "# load weights\n",
    "name = ['Moscow_e02_valoss0.212_RAEwSCwWSwINwCLF_frozenUnet.hdf5'][-1]\n",
    "weights = '/home/pherruzo/projects/nips_traffic/checkpoints/' + name\n",
    "if weights is not None:\n",
    "    print('loading weitghs:', weights)\n",
    "    model.load_weights(weights, by_name=True, skip_mismatch=True) #l24)#, by_name=True)\n",
    "\n",
    "model.compile(optimizer=optimizer,\n",
    "              loss={'Concat_predicted_frames': 'mse',\n",
    "                    'softmax_clf': losses.sparse_categorical_crossentropy},\n",
    "              metrics={'Concat_predicted_frames': ['mse'], \n",
    "                       'softmax_clf':             [losses.sparse_categorical_crossentropy]},#'accuracy']},\n",
    "              loss_weights=general_loss_weights, \n",
    "              sample_weight_mode=\"temporal\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Define callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('/home/pherruzo/projects/nips_traffic/checkpoints/Moscow_e{epoch:02d}_valoss{val_loss:.3f}_tttest.hdf5',\n",
       " '/home/pherruzo/projects/nips_traffic/tensorboard/RAE_w_SC_WS_wIN/tb_Moscow_1/run_tttest_1')"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from keras import backend as K \n",
    "\n",
    "class LRTensorBoard(keras.callbacks.TensorBoard):\n",
    "    def __init__(self, log_dir, **kwargs):  # add other arguments to __init__ if you need\n",
    "        super().__init__(log_dir=log_dir, **kwargs)\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        logs.update({'lr': K.eval(self.model.optimizer.lr)})\n",
    "        super().on_epoch_end(epoch, logs)\n",
    "        \n",
    "############################################# PREPARE CALLBACKS\n",
    "run_name = 'RAEwALLwCLF'#'RAEwSCaWT_w_IN'\n",
    "run_num = '1'\n",
    "## tensorboard\n",
    "tb_folder = \"/tb_{}_{}/run_{}\".format(training_params[\"city\"],\"1\", run_name+'_'+run_num)\n",
    "tb_folder = training_params[\"tensorboard_path\"]+MODEL_TYPES[m_type]+tb_folder\n",
    "# tb_update_freq = 1000 * params[\"batch_size\"]\n",
    "# print(\"Saving training progress in {} each {} samples.\".format(tb_folder, tb_update_freq))\n",
    "\n",
    "# save best model\n",
    "# checkpoint\n",
    "root= \"/home/pherruzo\" \n",
    "model_dir=root+\"/projects/nips_traffic/checkpoints/\"\n",
    "filepath=model_dir+training_params[\"city\"]+\"_e{epoch:02d}_valoss{val_loss:.3f}_\"+run_name+\".hdf5\"\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(filepath, monitor='val_loss', save_weights_only=True, save_best_only=True)\n",
    "\n",
    "# ep 1-4 lr: 0.001\n",
    "# ep 5- lr: 0.0001\n",
    "\n",
    "my_callbacks = [keras.callbacks.TensorBoard(log_dir=tb_folder), checkpoint, LRTensorBoard(log_dir=tb_folder)] #, update_freq=tb_update_freq)]\n",
    "# my_callbacks = [keras.callbacks.TensorBoard(log_dir=tb_folder, histogram_freq=1, write_grads=True), checkpoint, LRTensorBoard(log_dir=tb_folder]\n",
    "# my_callbacks = None\n",
    "filepath, tb_folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TRAIN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "learning rate: 0.01\n",
      "Epoch 2/20\n",
      "2279/2280 [============================>.] - ETA: 0s - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0445e-08-->> Data has been shuffled in: training\n",
      "2280/2280 [==============================] - 1806s 792ms/step - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0438e-08 - val_loss: 0.0120 - val_mean_squared_error: 0.0120 - val_r_emb_loss: 3.2361e-09\n",
      "Epoch 3/20\n",
      "2279/2280 [============================>.] - ETA: 0s - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0804e-08-->> Data has been shuffled in: training\n",
      "2280/2280 [==============================] - 1801s 790ms/step - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0797e-08 - val_loss: 0.0120 - val_mean_squared_error: 0.0120 - val_r_emb_loss: 6.6008e-09\n",
      "Epoch 4/20\n",
      "2279/2280 [============================>.] - ETA: 0s - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.1589e-08-->> Data has been shuffled in: training\n",
      "2280/2280 [==============================] - 1808s 793ms/step - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.1595e-08 - val_loss: 0.0119 - val_mean_squared_error: 0.0119 - val_r_emb_loss: 3.9253e-08\n",
      "Epoch 5/20\n",
      "2279/2280 [============================>.] - ETA: 0s - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0835e-08-->> Data has been shuffled in: training\n",
      "2280/2280 [==============================] - 1804s 791ms/step - loss: 0.0121 - mean_squared_error: 0.0121 - r_emb_loss: 2.0826e-08 - val_loss: 0.0119 - val_mean_squared_error: 0.0119 - val_r_emb_loss: 3.6410e-10\n",
      "Epoch 6/20\n",
      "1169/2280 [==============>...............] - ETA: 14:36 - loss: 0.0120 - mean_squared_error: 0.0120 - r_emb_loss: 2.0431e-08"
     ]
    }
   ],
   "source": [
    "DEBUG = False\n",
    "\n",
    "if DEBUG:\n",
    "    print('validation data to memory')\n",
    "    valid_data = ([x_val, y_val], y_val)\n",
    "    my_callbacks = [keras.callbacks.TensorBoard(log_dir=tb_folder, histogram_freq=1, write_grads=True), \n",
    "                    checkpoint, \n",
    "                    LRTensorBoard(log_dir=tb_folder)]\n",
    "else:\n",
    "    valid_data = val_ds\n",
    "    my_callbacks = [keras.callbacks.TensorBoard(log_dir=tb_folder), checkpoint, LRTensorBoard(log_dir=tb_folder)]\n",
    "\n",
    "print(\"learning rate:\", K.get_value(model.optimizer.lr))\n",
    "hist = model.fit_generator(training_ds, epochs=20, verbose=1, callbacks=my_callbacks, \n",
    "                           validation_data=valid_data,\n",
    "                           max_queue_size=10, workers=4, shuffle=True, initial_epoch=1)#, steps_per_epoch=2000) # starts with next epoch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "performance in validations:\n",
      "val_loss: 1.2941249370574952\n",
      "val_Concat_predicted_frames_loss: 0.07352456067289626\n",
      "val_softmax_clf_loss: 1.2206004006522042\n",
      "val_Concat_predicted_frames_mean_squared_error: 0.02733775063284806\n",
      "val_Concat_predicted_frames_r_emb_loss: 0.04618680945464543\n",
      "val_softmax_clf_sparse_categorical_crossentropy: 1.013511712210519\n"
     ]
    }
   ],
   "source": [
    "desc = ['val_loss', 'val_Concat_predicted_frames_loss', 'val_softmax_clf_loss', 'val_Concat_predicted_frames_mean_squared_error', 'val_Concat_predicted_frames_r_emb_loss', 'val_softmax_clf_sparse_categorical_crossentropy']\n",
    "\n",
    "# with clf\n",
    "metrics = model.evaluate_generator(val_ds)\n",
    "\n",
    "print(\"performance in validations:\")\n",
    "for d, v in zip(desc, metrics):\n",
    "    print('{}: {}'.format(d, v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "performance in validations:\n",
      "val_loss: 0.025962504638092857\n",
      "val_Concat_predicted_frames_loss: 0.025962504638092857\n",
      "val_softmax_clf_loss: 1.668187713623047\n",
      "val_Concat_predicted_frames_mean_squared_error: 0.015087313205003738\n",
      "val_Concat_predicted_frames_r_emb_loss: 0.010875191566135202\n",
      "val_softmax_clf_sparse_categorical_crossentropy: 0.8352041755403791\n"
     ]
    }
   ],
   "source": [
    "desc = ['val_loss', 'val_Concat_predicted_frames_loss', 'val_softmax_clf_loss', 'val_Concat_predicted_frames_mean_squared_error', 'val_Concat_predicted_frames_r_emb_loss', 'val_softmax_clf_sparse_categorical_crossentropy']\n",
    "\n",
    "# with clf\n",
    "metrics = model.evaluate_generator(val_ds)\n",
    "\n",
    "print(\"performance in validations:\")\n",
    "for d, v in zip(desc, metrics):\n",
    "    print('{}: {}'.format(d, v))"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}