examples/training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HIDRA model training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_VISIBLE_DEVICES=-1\n"
     ]
    }
   ],
   "source": [
    "# Comment out to enable GPU support (requires a compatible CUDA version)\n",
    "%env CUDA_VISIBLE_DEVICES -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import h5py\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hidra import HIDRA, compile_model\n",
    "from hidra.data import DataNormalization, hdf5_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 1**: Prepare the training (and validation) data and import it as a TensorFlow Dataset of samples with the following structure:\n",
    "```\n",
    "(atmospheric_data, sea_level_data), labels\n",
    "```\n",
    "\n",
    "In our setup, we prepare the data into an HDF5 file containing the samples. The HDF5 file contains the following fields:\n",
    "\n",
    "| Field name | Shape | Description |\n",
    "|---|----------------------|---|\n",
    "| `weather` | $$ N  \\times \\frac{ T_{max} + T_{min} }{4} \\times H \\times W \\times 4  $$ | Atmospheric input tensors subsampled (by average pooling) to a 4h temporal resolution. |\n",
    "| `ssh`, `tide` & `delta` | $$N \\times T_{min} \\times 1$$ | Sea level tensors (full, tidal component and residual component) |\n",
    "| `lbl_ssh`, `lbl_tide` & `lbl_delta` | $$N \\times T_{max} \\times 1$$ | Target (labels) sea level tensors (full, tidal component and residual component). |\n",
    "| `dates` (optional) | $$N \\times T_{max}$$ | Timestamps corresponding to prediction times of labeled data (`lbl_*`) |\n",
    "\n",
    "$N$ denotes the number of samples in the dataset, $H$ and $W$ define the spatial dimensions of the atmospheric data, $T_{min}$ and $T_{max}$ define the past and future horizons for HIDRA prediction. Refer to the provided sample data file for additional information about the structure of the data.\n",
    "\n",
    "In our setup we use the data for the last 24 hours ($T_{min} = 24$) and predict the values for the next 72 hours ($T_{max} = 72$). The atmospheric grid contains $29 \\times 37$ data points ($H=29$, $W=37$). Additionally, sea level input contains two signals - residual and tide. The shapes of a single sample of our data are thus as follows:\n",
    "\n",
    "| Data | Shape |\n",
    "|---|----------------------|\n",
    "| `atmospheric_data` | $$ 24 \\times 29 \\times 37 \\times 4  $$  |\n",
    "| `sea_level_data` | $$ 24 \\times 2 $$|\n",
    "| `labels` | $$ 72 $$ |\n",
    "\n",
    "\n",
    "\n",
    "We load the HDF5 file into a TensorFlow dataset then make use of a mapping function to select and prepare the data for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify which fields from the HDF5 file to load. \n",
    "# We'll use the residual and tide signals as input and predict residuals in this case.\n",
    "dataset = hdf5_dataset('../data/example_data.hdf5', ['weather', 'delta', 'tide', 'lbl_delta'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a mapping function that prepares the data\n",
    "def map_fn(weather, delta, tide, lbl_delta):\n",
    "    sea_level = tf.concat([delta, tide], axis=1)\n",
    "    return (weather, sea_level), lbl_delta[..., 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<MapDataset shapes: (((24, 29, 37, 4), (24, 2)), (72,)), types: ((tf.float32, tf.float64), tf.float64)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = dataset.map(map_fn)\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 2**: Create and compile the HIDRA model before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = HIDRA(num_predictions=72)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepares the loss function, metrics and the optimizer\n",
    "model = compile_model(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 3**: Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2000/2000 [==============================] - 802s 401ms/step - loss: -0.4123 - mean_absolute_error: 0.1972\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7ffa215cd210>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(dataset.repeat().batch(32), epochs=1, steps_per_epoch=2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 4**: Export the model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('../models/my_model.hdf5')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hidra",
   "language": "python",
   "name": "hidra"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}