MODEL00_VQ_VAE_main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "93b514f8-2474-4936-a38f-38b39644f2bf",
   "metadata": {},
   "source": [
    "# VQ-VAE for CCPA Embedding\n",
    "\n",
    "This notebook contains information for the design, hyperparameters, and training of a VQ-VAE that projects the CCPA precipitation field into a latent space\n",
    "\n",
    "* Precipitation has discontinuities between zero and nonzero values. The VQ-VAE creates a regularized space (similar to other VAEs) where such discontinuity will be normalized.\n",
    "\n",
    "* \n",
    "By embedding precipitation fields into a latent space, it also reduces the overall size of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "64b842bb-810f-4bef-9858-b6ac22842cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "import numpy as np\n",
    "from glob import glob\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Turn-off warnings\n",
    "import logging\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "logging.getLogger(\"tensorflow\").setLevel(logging.ERROR)\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Turn-off tensoflow-specific warnings\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "tf.autograph.set_verbosity(0)\n",
    "tf.get_logger().setLevel('ERROR')\n",
    "\n",
    "# ------------------------------------------------------- #\n",
    "# Import customized modules and settings\n",
    "sys.path.insert(0, '/glade/u/home/ksha/GAN_proj/')\n",
    "sys.path.insert(0, '/glade/u/home/ksha/GAN_proj/libs/')\n",
    "\n",
    "from namelist import *\n",
    "import data_utils as du\n",
    "import model_utils as mu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "621fbd3e-b1dd-40fc-984b-0102d62e65cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0caeb6a-e0b2-4e73-af21-03eb93001128",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8e13b863-2e65-4182-ae87-3ff6d697f4ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Hyperparameters\n",
    "filter_nums = [64, 128] # number of convolution kernels per down-/upsampling layer \n",
    "latent_dim = 4 # number of latent feature channels\n",
    "activation = 'gelu' # activation function\n",
    "num_embeddings = 128 #128 # number of the VQ codes\n",
    "\n",
    "input_size = (224, 464, 1) # size of MRMS input\n",
    "latent_size = (14, 29, latent_dim) # size of compressed latent features\n",
    "\n",
    "load_weights = True\n",
    "\n",
    "model_name_load = model_dir+'baseline/VQ_VAE_stack1_pretrain/'\n",
    "model_name_save = model_dir+'baseline/VQ_VAE_stack1_pretrain/'\n",
    "\n",
    "# separate save encoder and decoder\n",
    "model_name_encoder_save = model_dir+'models/VQ_VAE_encoder_stack1_tune0'\n",
    "model_name_decoder_save = model_dir+'models/VQ_VAE_decoder_stack1_tune0'\n",
    "\n",
    "drop_encode = True\n",
    "drop_decode = True\n",
    "\n",
    "lr = 1e-5 # learning rate\n",
    "# samples per epoch = N_batch * batch_size\n",
    "epochs = 99999\n",
    "N_batch = 4\n",
    "batch_size = 32*16\n",
    "batch_size_ = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e44c8541-9c0b-437e-94a7-d93781afd44c",
   "metadata": {},
   "source": [
    "## Validation set preparation\n",
    "\n",
    "* Training set: 23646 samples CCPA from 20020101 to 20191231 \n",
    "* Validation set: 2627 samples (~10%) from the training set\n",
    "* Reproducible selection:\n",
    "```python\n",
    "filename_valid = filenames[::8][:2627]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "824d94ff-3ee6-4ebb-b7f9-d835ae8bcffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Validation set prep\n",
    "# location of training data\n",
    "BATCH_dir = camp_dir+'BATCH_CCPA_full/'\n",
    "# validation set size\n",
    "\n",
    "# collect validation set sampales\n",
    "filenames = sorted(glob(BATCH_dir+'*.npy'))\n",
    "filenames = filenames[:-3648]\n",
    "L = len(filenames)\n",
    "\n",
    "filename_valid = filenames[::8][:2627]\n",
    "#filename_valid = filenames[::8][:500] # samller validation set size\n",
    "L_valid = len(filename_valid)\n",
    "\n",
    "Y_valid = np.empty((L_valid, 224, 464, 1))\n",
    "Y_valid[...] = np.nan\n",
    "\n",
    "for i, name in enumerate(filename_valid):\n",
    "    Y_valid[i, ..., 0] = np.load(name)\n",
    "    \n",
    "# # make sure the validation set contains no NaNs\n",
    "# print('NaN grids = {}'.format(np.sum(np.isnan(Y_valid))))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9521893d-81e1-43ef-96ff-76827a1674bf",
   "metadata": {},
   "source": [
    "## Multi-GPU training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "645f9570-cc99-404e-b0d8-6065ee093ae5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Training data prep\n",
    "# Capture all the training set\n",
    "filenames = sorted(glob(BATCH_dir+'*.npy'))\n",
    "filenames = filenames[:-3648]\n",
    "filename_train = list(set(filenames) - set(filename_valid))\n",
    "\n",
    "L_train = len(filename_train)\n",
    "\n",
    "min_del = 0.0\n",
    "max_tol = 3 # early stopping with 2-epoch patience\n",
    "tol = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "78725eea-73a8-459d-9154-5e584d1dbaa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Distributed training\n",
    "\n",
    "strategy = tf.distribute.MirroredStrategy()\n",
    "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))\n",
    "\n",
    "with strategy.scope():\n",
    "    # ------------------------------------------------------- #\n",
    "    # VQ-VAE\n",
    "    # ---------------- encoder ----------------- #\n",
    "    encoder = mu.VQ_VAE_encoder(input_size, filter_nums, latent_dim, num_embeddings, activation, drop_encode)\n",
    "    decoder = mu.VQ_VAE_decoder(latent_size, filter_nums, activation, drop_decode)\n",
    "\n",
    "    # Connect the encoder and decoder\n",
    "    X = keras.Input(shape=input_size)\n",
    "    X_encode = encoder(X)\n",
    "    X_decode = decoder(X_encode)\n",
    "    model_vqvae = keras.Model(X, X_decode)\n",
    "    \n",
    "    # subclass to VAE training\n",
    "    vqvae_trainer = mu.VQVAETrainer(model_vqvae, 1.0, latent_dim, num_embeddings)\n",
    "    \n",
    "    # load weights\n",
    "    if load_weights:\n",
    "        W_old = mu.dummy_loader(model_name_load)\n",
    "        vqvae_trainer.vqvae.set_weights(W_old)\n",
    "    \n",
    "    # compile\n",
    "    vqvae_trainer.compile(optimizer=keras.optimizers.Adam(learning_rate=lr))\n",
    "    \n",
    "    # ----------------------------------------------- #\n",
    "    # Major training loop + training batch generation\n",
    "    \n",
    "    Y_batch = np.empty((batch_size, 224, 464, 1))\n",
    "    Y_batch[...] = np.nan\n",
    "    \n",
    "    for i in range(epochs):\n",
    "        \n",
    "        print('epoch = {}'.format(i))\n",
    "        if i == 0:\n",
    "            model_ = vqvae_trainer.vqvae\n",
    "            Y_pred = model_.predict(Y_valid)\n",
    "            Y_pred[Y_pred<0] = 0\n",
    "            record = du.mean_absolute_error(Y_valid, Y_pred)\n",
    "            print('Initial validation loss: {}'.format(record))\n",
    "        \n",
    "        start_time = time.time()\n",
    "        for j in range(N_batch):\n",
    "            \n",
    "            inds_rnd = du.shuffle_ind(L_train)\n",
    "            inds_ = inds_rnd[:batch_size]\n",
    "    \n",
    "            for k, ind in enumerate(inds_):\n",
    "                # import batch data\n",
    "                name = filename_train[ind]\n",
    "                Y_batch[k, ..., 0] = np.load(name)\n",
    "\n",
    "            Y_batch += np.random.normal(0, 0.01, size=Y_batch.shape)\n",
    "            vqvae_trainer.fit(Y_batch, epochs=1, batch_size=batch_size_, verbose=0)\n",
    "            \n",
    "        # on epoch-end\n",
    "        model_ = vqvae_trainer.vqvae\n",
    "        Y_pred = model_.predict(Y_valid)\n",
    "        Y_pred[Y_pred<0] = 0\n",
    "        record_temp = du.mean_absolute_error(Y_valid, Y_pred)\n",
    "    \n",
    "        if record - record_temp > min_del:\n",
    "            print('Validation loss improved from {} to {}'.format(record, record_temp))\n",
    "            record = record_temp\n",
    "            model_ = vqvae_trainer.vqvae\n",
    "            print(\"Save to {}\".format(model_name_save))\n",
    "            model_.save(model_name_save)\n",
    "            \n",
    "        else:\n",
    "            print('Validation loss {} NOT improved'.format(record_temp))\n",
    "        \n",
    "        print(\"--- %s seconds ---\" % (time.time() - start_time))\n",
    "        # mannual callbacks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe226de-41d1-4e16-bd0d-8617f65dcb76",
   "metadata": {},
   "source": [
    "## Load model after trainnig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8ef9506c-bb54-43fa-8e6e-8e9590be7d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Make sure the model can load\n",
    "encoder = mu.VQ_VAE_encoder(input_size, filter_nums, latent_dim, num_embeddings, activation, drop_encode)\n",
    "decoder = mu.VQ_VAE_decoder(latent_size, filter_nums, activation, drop_decode)\n",
    "\n",
    "# Connect the encoder and decoder\n",
    "X = keras.Input(shape=input_size)\n",
    "X_encode = encoder(X)\n",
    "X_decode = decoder(X_encode)\n",
    "model_vqvae = keras.Model(X, X_decode)\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_load)\n",
    "model_vqvae.set_weights(W_old)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3efe8b52-c725-4681-a542-5763f9606ffd",
   "metadata": {},
   "source": [
    "## Separate encoder and decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8a92685f-2889-45bd-87f1-2a8d3f5951d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16/16 [==============================] - 8s 314ms/step\n",
      "Initial validation loss: 0.0065533267885267334\n"
     ]
    }
   ],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Save encoder and decoder separately\n",
    "Y_pred = model_vqvae.predict(Y_valid)\n",
    "Y_pred[Y_pred<0] = 0\n",
    "record = du.mean_absolute_error(Y_valid, Y_pred)\n",
    "print('Initial validation loss: {}'.format(record))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ffdaffb-3f8c-45db-9842-278a4be0158d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model_encoder.save(model_name_encoder_save)\n",
    "# model_decoder.save(model_name_decoder_save)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86cb3998-59ad-4667-9ddf-f55952ad070a",
   "metadata": {},
   "source": [
    "## Validate saved encoder-decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4a01879c-c7ad-40b6-bd23-9549a08e4672",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------- #\n",
    "# Make sure that the save encoder and decoder can connect\n",
    "model_encoder = mu.VQ_VAE_encoder(input_size, filter_nums, latent_dim, num_embeddings, activation, drop_encode)\n",
    "model_decoder = mu.VQ_VAE_decoder(latent_size, filter_nums, activation, drop_decode)\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_encoder_save)\n",
    "model_encoder.set_weights(W_old)\n",
    "\n",
    "W_old = mu.dummy_loader(model_name_decoder_save)\n",
    "model_decoder.set_weights(W_old)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1d5779bb-2aa2-4e1b-9152-2e86ed60f165",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "83/83 [==============================] - 8s 94ms/step\n",
      "83/83 [==============================] - 12s 144ms/step\n"
     ]
    }
   ],
   "source": [
    "Y_encoder = model_encoder.predict(Y_valid)\n",
    "Y_pred_valid = model_decoder.predict(Y_encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5b178a46-7de2-4494-8885-fc144cc53730",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial validation loss: 0.006887156543567399\n"
     ]
    }
   ],
   "source": [
    "Y_pred_valid[Y_pred_valid<0] = 0\n",
    "record = du.mean_absolute_error(Y_valid, Y_pred_valid)\n",
    "print('Initial validation loss: {}'.format(record))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d2ba3e-fa95-4cec-86a4-62892961efa7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}