train_full_data_set.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import runet as vae_util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from keras import backend as K\n",
    "from keras import layers\n",
    "from keras.models import Model\n",
    "from keras.optimizers import Adam\n",
    "import h5py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "image (InputLayer)              (None, 128, 128, 3)  0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_1 (Conv2D)               (None, 64, 64, 16)   448         image[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_1 (BatchNor (None, 64, 64, 16)   64          conv2d_1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_1 (Activation)       (None, 64, 64, 16)   0           batch_normalization_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_2 (Conv2D)               (None, 64, 64, 32)   4640        activation_1[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_2 (BatchNor (None, 64, 64, 32)   128         conv2d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_2 (Activation)       (None, 64, 64, 32)   0           batch_normalization_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_3 (Conv2D)               (None, 32, 32, 64)   18496       activation_2[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_3 (BatchNor (None, 32, 32, 64)   256         conv2d_3[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_3 (Activation)       (None, 32, 32, 64)   0           batch_normalization_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_4 (Conv2D)               (None, 32, 32, 64)   36928       activation_3[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_4 (BatchNor (None, 32, 32, 64)   256         conv2d_4[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_4 (Activation)       (None, 32, 32, 64)   0           batch_normalization_4[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_5 (Conv2D)               (None, 16, 16, 128)  73856       activation_4[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_5 (BatchNor (None, 16, 16, 128)  512         conv2d_5[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_5 (Activation)       (None, 16, 16, 128)  0           batch_normalization_5[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_6 (Conv2D)               (None, 16, 16, 128)  147584      activation_5[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_6 (BatchNor (None, 16, 16, 128)  512         conv2d_6[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_6 (Activation)       (None, 16, 16, 128)  0           batch_normalization_6[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_7 (Conv2D)               (None, 16, 16, 128)  147584      activation_6[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_7 (BatchNor (None, 16, 16, 128)  512         conv2d_7[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_7 (Activation)       (None, 16, 16, 128)  0           batch_normalization_7[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_8 (Conv2D)               (None, 16, 16, 128)  147584      activation_7[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_8 (BatchNor (None, 16, 16, 128)  512         conv2d_8[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "add_1 (Add)                     (None, 16, 16, 128)  0           activation_6[0][0]               \n",
      "                                                                 batch_normalization_8[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_9 (Conv2D)               (None, 16, 16, 128)  147584      add_1[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_9 (BatchNor (None, 16, 16, 128)  512         conv2d_9[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "activation_8 (Activation)       (None, 16, 16, 128)  0           batch_normalization_9[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_10 (Conv2D)              (None, 16, 16, 128)  147584      activation_8[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_10 (BatchNo (None, 16, 16, 128)  512         conv2d_10[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_2 (Add)                     (None, 16, 16, 128)  0           add_1[0][0]                      \n",
      "                                                                 batch_normalization_10[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_11 (Conv2D)              (None, 16, 16, 128)  147584      add_2[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_11 (BatchNo (None, 16, 16, 128)  512         conv2d_11[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_9 (Activation)       (None, 16, 16, 128)  0           batch_normalization_11[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_12 (Conv2D)              (None, 16, 16, 128)  147584      activation_9[0][0]               \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_12 (BatchNo (None, 16, 16, 128)  512         conv2d_12[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_3 (Add)                     (None, 16, 16, 128)  0           add_2[0][0]                      \n",
      "                                                                 batch_normalization_12[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_13 (Conv2D)              (None, 16, 16, 128)  147584      add_3[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_13 (BatchNo (None, 16, 16, 128)  512         conv2d_13[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_10 (Activation)      (None, 16, 16, 128)  0           batch_normalization_13[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_14 (Conv2D)              (None, 16, 16, 128)  147584      activation_10[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_14 (BatchNo (None, 16, 16, 128)  512         conv2d_14[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_4 (Add)                     (None, 16, 16, 128)  0           add_3[0][0]                      \n",
      "                                                                 batch_normalization_14[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_15 (Conv2D)              (None, 16, 16, 128)  147584      add_4[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_15 (BatchNo (None, 16, 16, 128)  512         conv2d_15[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_11 (Activation)      (None, 16, 16, 128)  0           batch_normalization_15[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_16 (Conv2D)              (None, 16, 16, 128)  147584      activation_11[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_16 (BatchNo (None, 16, 16, 128)  512         conv2d_16[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_5 (Add)                     (None, 16, 16, 128)  0           add_4[0][0]                      \n",
      "                                                                 batch_normalization_16[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_17 (Conv2D)              (None, 16, 16, 128)  147584      add_5[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_17 (BatchNo (None, 16, 16, 128)  512         conv2d_17[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_12 (Activation)      (None, 16, 16, 128)  0           batch_normalization_17[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_18 (Conv2D)              (None, 16, 16, 128)  147584      activation_12[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_18 (BatchNo (None, 16, 16, 128)  512         conv2d_18[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_6 (Add)                     (None, 16, 16, 128)  0           add_5[0][0]                      \n",
      "                                                                 batch_normalization_18[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_1 (Concatenate)     (None, 16, 16, 256)  0           activation_6[0][0]               \n",
      "                                                                 add_6[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_1 (UnPooling2D)    (None, 16, 16, 256)  0           concatenate_1[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_1 (Reflect (None, 18, 18, 256)  0           un_pooling2d_1[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_19 (Conv2D)              (None, 16, 16, 128)  295040      reflection_padding2d_1[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_19 (BatchNo (None, 16, 16, 128)  512         conv2d_19[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_13 (Activation)      (None, 16, 16, 128)  0           batch_normalization_19[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_2 (Concatenate)     (None, 16, 16, 256)  0           activation_5[0][0]               \n",
      "                                                                 activation_13[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_2 (UnPooling2D)    (None, 32, 32, 256)  0           concatenate_2[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_2 (Reflect (None, 34, 34, 256)  0           un_pooling2d_2[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_20 (Conv2D)              (None, 32, 32, 128)  295040      reflection_padding2d_2[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_20 (BatchNo (None, 32, 32, 128)  512         conv2d_20[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_14 (Activation)      (None, 32, 32, 128)  0           batch_normalization_20[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_3 (Concatenate)     (None, 32, 32, 192)  0           activation_4[0][0]               \n",
      "                                                                 activation_14[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_3 (UnPooling2D)    (None, 32, 32, 192)  0           concatenate_3[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_3 (Reflect (None, 34, 34, 192)  0           un_pooling2d_3[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_21 (Conv2D)              (None, 32, 32, 64)   110656      reflection_padding2d_3[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_21 (BatchNo (None, 32, 32, 64)   256         conv2d_21[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_15 (Activation)      (None, 32, 32, 64)   0           batch_normalization_21[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_4 (Concatenate)     (None, 32, 32, 128)  0           activation_3[0][0]               \n",
      "                                                                 activation_15[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_4 (UnPooling2D)    (None, 64, 64, 128)  0           concatenate_4[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_4 (Reflect (None, 66, 66, 128)  0           un_pooling2d_4[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_22 (Conv2D)              (None, 64, 64, 64)   73792       reflection_padding2d_4[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_22 (BatchNo (None, 64, 64, 64)   256         conv2d_22[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_16 (Activation)      (None, 64, 64, 64)   0           batch_normalization_22[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_5 (Concatenate)     (None, 64, 64, 96)   0           activation_2[0][0]               \n",
      "                                                                 activation_16[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_5 (UnPooling2D)    (None, 64, 64, 96)   0           concatenate_5[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_5 (Reflect (None, 66, 66, 96)   0           un_pooling2d_5[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_23 (Conv2D)              (None, 64, 64, 32)   27680       reflection_padding2d_5[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_23 (BatchNo (None, 64, 64, 32)   128         conv2d_23[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_17 (Activation)      (None, 64, 64, 32)   0           batch_normalization_23[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_6 (Concatenate)     (None, 64, 64, 48)   0           activation_1[0][0]               \n",
      "                                                                 activation_17[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "un_pooling2d_6 (UnPooling2D)    (None, 128, 128, 48) 0           concatenate_6[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "reflection_padding2d_6 (Reflect (None, 130, 130, 48) 0           un_pooling2d_6[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_24 (Conv2D)              (None, 128, 128, 16) 6928        reflection_padding2d_6[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_24 (BatchNo (None, 128, 128, 16) 64          conv2d_24[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "activation_18 (Activation)      (None, 128, 128, 16) 0           batch_normalization_24[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_25 (Conv2D)              (None, 128, 128, 1)  145         activation_18[0][0]              \n",
      "==================================================================================================\n",
      "Total params: 2,871,841\n",
      "Trainable params: 2,867,041\n",
      "Non-trainable params: 4,800\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "input_shape = (128, 128, 3)\n",
    "vae_model = vae_util.create_vae(input_shape)\n",
    "vae_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please download the data at https://drive.google.com/drive/folders/1nHpXhKjOw_nBk8GvoYfzBj5tcBNaZ6Hh?usp=sharing\n",
    "data_dir = '/your/data/dir/'\n",
    "hf_r = h5py.File(data_dir + 'train_200000.hdf5', 'r')\n",
    "train_x = np.array(hf_r.get('k'))\n",
    "train_y = np.array(hf_r.get('S'))\n",
    "hf_r.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_w = h5py.File(data_dir + 'test_200000.hdf5', 'r')\n",
    "test_x = np.array(hf_w.get('k'))\n",
    "test_y = np.array(hf_w.get('S'))\n",
    "hf_w.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('train_x shape is ', train_x.shape)\n",
    "print('train_y shape is ', train_y.shape)\n",
    "print('test_x shape is ', test_x.shape)\n",
    "print('test_y shape is ', test_y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x = train_x.transpose(2, 0, 1, 3)\n",
    "train_y = train_y.transpose(2, 0, 1, 3)\n",
    "test_x = test_x.transpose(2, 0, 1, 3)\n",
    "test_y = test_y.transpose(2, 0, 1, 3)\n",
    "print('train_x shape is ', train_x.shape)\n",
    "print('train_y shape is ', train_y.shape)\n",
    "print('test_x shape is ', test_x.shape)\n",
    "print('test_y shape is ', test_y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Loss\n",
    "def vae_loss(x, t_decoded):\n",
    "    '''Total loss for the plain UAE'''\n",
    "    return K.mean(reconstruction_loss(x, t_decoded))\n",
    "\n",
    "def reconstruction_loss(x, t_decoded):\n",
    "    '''Reconstruction loss for the plain UAE'''\n",
    "    return K.sum((K.batch_flatten(x) - K.batch_flatten(t_decoded)) ** 2, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training specification\n",
    "epoch = 50\n",
    "train_nr = train_x.shape[0]\n",
    "batch_size = 16\n",
    "test_nr = 80\n",
    "num_batch = int(train_nr/batch_size) \n",
    "learning_rate = 1e-4\n",
    "\n",
    "opt = Adam(lr = learning_rate)\n",
    "\n",
    "train_target = K.placeholder(shape=(batch_size, 128, 128, 1))\n",
    "test_target = K.placeholder(shape=(test_nr, 128, 128, 1))\n",
    "\n",
    "rec_loss = vae_loss(vae_model.output, train_target)\n",
    "vae_model.compile(optimizer=opt, loss='mse')\n",
    "\n",
    "total_loss = rec_loss\n",
    "\n",
    "updates = opt.get_updates(total_loss, vae_model.trainable_weights)\n",
    "\n",
    "iterate = K.function(vae_model.inputs + [train_target], [rec_loss], updates=updates)\n",
    "\n",
    "eval_rec_loss = vae_loss(vae_model.output, test_target)\n",
    "\n",
    "evaluate = K.function(vae_model.inputs + [test_target], [eval_rec_loss])\n",
    "\n",
    "output_dir = '/your/model/dir/saved_models/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train\n",
    "for e in range(epoch):\n",
    "    for ib in range(num_batch):\n",
    "        ind0 = ib * batch_size\n",
    "        n_itr = e * train_nr + ind0 + batch_size # for tensorboard output\n",
    "        x_batch = train_x[ind0:ind0+batch_size, ...]\n",
    "        y_batch = train_y[ind0:ind0+batch_size, ...]\n",
    "        rec_loss_val = iterate([x_batch] + [y_batch])\n",
    "        eval_loss_val = evaluate([test_x[:100,...]] + [test_y[:100,...]])\n",
    "        \n",
    "        if ib % 100 == 0:\n",
    "            print('Epoch %d/%d, Batch %d/%d, Rec Loss %f' % (e+1, epoch, ib+1, num_batch, rec_loss_val[0]))\n",
    "    print('Epoch %d/%d, Train Rec loss %f, Eval Rec loss %f' % (e + 1, epoch, rec_loss_val[0], eval_loss_val[0]))\n",
    "\n",
    "    if (e+1) % 10 == 0:\n",
    "        vae_model.save_weights(output_dir +  'trained_model_ep%d.h5' % ((e+1)))\n",
    "\n",
    "vae_model.save_weights(output_dir +  'trained_model_ep%d.h5' % ((e+1)))"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}