Train_3D_Unet_Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# matplotlib plots within notebook\n",
    "%matplotlib inline\n",
    "\n",
    "import platform\n",
    "print(\"python: \"+platform.python_version())\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time\n",
    "\n",
    "import os, shutil, sys\n",
    "\n",
    "\n",
    "sys.path.insert(0, '.')\n",
    "from DeepAttCorr_lib import GAN_3D_lib as GAN\n",
    "from DeepAttCorr_lib import data_handling as DH\n",
    "from DeepAttCorr_lib import file_manage_utils as File_mng\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "import tensorflow.keras.backend as K\n",
    "print('Using TensorFlow version: '+tf.__version__)\n",
    "\n",
    "from tensorflow.python.client import device_lib\n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Paths and Definitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Basics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Network name\n",
    "NETWORK_NAME = 'DeepAttCorr_3D_Unet_Network'\n",
    "\n",
    "# Dataset location\n",
    "DATASET_PATH = './datasets/'\n",
    "\n",
    "# Checkpoint location\n",
    "CHECKPOINT_PATH = \"./Outputs/\"+NETWORK_NAME+\"/\"\n",
    "\n",
    "# Clear outputs before running\n",
    "CLEAR_OUTS = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imput volume size\n",
    "voxels_X = 128\n",
    "voxels_Y = 128\n",
    "voxels_Z = 32\n",
    "input_size = (voxels_X,voxels_Y,voxels_Z)\n",
    "\n",
    "# Network convolutional channels by resolution level\n",
    "USE_GEN_net_conv_Channels = [10, 20, 40, 80, 160]\n",
    "# Convolutional layers by resolution level\n",
    "USE_GEN_net_conv_Layers = 2\n",
    "\n",
    "# Hyperbolic tangent or Sigmoid output for generator\n",
    "USE_GEN_TANH_OUT = False\n",
    "USE_GEN_SIGMOID_OUT = True\n",
    "# Use or not segmentation path\n",
    "USE_GEN_SEGMENTATION = True\n",
    "# Number of objective clases\n",
    "USE_GEN_OBJECTIVE_SEGMENTATION_CLASES = 4\n",
    "# Number of fully connected segmentation layers\n",
    "USE_GEN_OBJECTIVE_SEGMENTATION_LAYERS = 2\n",
    "# Number of convolutional segmentation layers\n",
    "USE_GEN_OBJECTIVE_CONV_SEGMENTATION_LAYERS = 4\n",
    "# Segmentation kernel size\n",
    "USE_GEN_SEGM_KERNEL_SIZE = 3\n",
    "# Use batch normalization for generator training\n",
    "USE_GEN_BATCH_NORM = False\n",
    "# Use pixel normalization for generator training\n",
    "USE_GEN_PIXEL_NORM = True\n",
    "# Use He scalling of weights for generator\n",
    "USE_GEN_HE_SCALLING = True\n",
    "# Wheight initialization standard deviation\n",
    "USE_GEN_INI_STD = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset full size (without slicing)\n",
    "DATASET_X_size = 128\n",
    "DATASET_Y_size = 128\n",
    "DATASET_Z_size = 256\n",
    "\n",
    "# Train cicles and checkpoints\n",
    "CICLES_TRAIN = 100\n",
    "CICLES_PER_SAVE = 5\n",
    "EPOCHS_PER_PLOTS = 10\n",
    "STEPS_PER_EPOCH = 100\n",
    "\n",
    "\n",
    "# Mini-batch size\n",
    "BATCH_SIZE_TRAIN = 4\n",
    "BUFFER_SIZE_TRAIN = 4\n",
    "BATCH_SIZE_VALIDATION = 4\n",
    "BUFFER_SIZE_VALIDATION = 4\n",
    "\n",
    "# Initial step size\n",
    "step_size = 0.001\n",
    "\n",
    "# Uniform or custom sampling of the input FOV\n",
    "# If True the input sample is sliced with uniform probability\n",
    "# If False, the Cumulative Density Function in CDF_PATH will control the sampling\n",
    "UNIFORM_SAMPLING = False\n",
    "CDF_PATH = \"./DeepAttCorr_lib/cdf_coef.npy\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "File_mng.check_create_path('CHECKPOINT_PATH', CHECKPOINT_PATH, clear_folder=CLEAR_OUTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cdf_coef = [1.0]\n",
    "if not UNIFORM_SAMPLING:\n",
    "    cdf_coef = np.load(CDF_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset Reader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Create dataset list\n",
    "shape_X = int(voxels_X)\n",
    "shape_Y = int(voxels_Y)\n",
    "shape_Z = int(voxels_Z)\n",
    "\n",
    "data_size = np.array([int(DATASET_X_size), \n",
    "                      int(DATASET_Y_size), \n",
    "                      int(DATASET_Z_size)])\n",
    "\n",
    "input_size_this = (shape_X,shape_Y,shape_Z)\n",
    "\n",
    "# Get dataset name\n",
    "train_dataset_name = 'Train_Dataset_%dx%dx%d.tfrecord'%(data_size[0],data_size[1],data_size[2])\n",
    "validation_dataset_name = 'Validation_Dataset_%dx%dx%d.tfrecord'%(data_size[0],data_size[1],data_size[2])\n",
    "\n",
    "# Create dataset reading pipelines\n",
    "PATH_TFRECORD_TRAIN = os.path.join(DATASET_PATH, train_dataset_name)\n",
    "PATH_TFRECORD_VALIDATION = os.path.join(DATASET_PATH, validation_dataset_name)\n",
    "\n",
    "dataset_train = tf.data.TFRecordDataset(PATH_TFRECORD_TRAIN)\n",
    "dataset_validation = tf.data.TFRecordDataset(PATH_TFRECORD_VALIDATION)\n",
    "\n",
    "if shape_X <= 32: \n",
    "    dataset_train = dataset_train.cache()\n",
    "    dataset_validation = dataset_validation.cache()\n",
    "    print('Using cache for dataset: %s'%train_dataset_name)\n",
    "\n",
    "# Create train dataset with transformations\n",
    "dataset_train = dataset_train.map(lambda x: DH.tf_get_keras_sample(x,\n",
    "                                                                   data_size, \n",
    "                                                                   input_size_this, \n",
    "                                                                   not_transformed = True,\n",
    "                                                                   cdf_sampler_coef=cdf_coef))\n",
    "\n",
    "# Create validation dataset, unmodified\n",
    "dataset_validation = dataset_validation.map(lambda x: DH.tf_get_keras_sample(x,\n",
    "                                                                             data_size, \n",
    "                                                                             input_size_this, \n",
    "                                                                             not_transformed = True,\n",
    "                                                                             cdf_sampler_coef=cdf_coef))\n",
    "\n",
    "\n",
    "# Shuffle the train dataset\n",
    "dataset_train= dataset_train.shuffle(buffer_size=BUFFER_SIZE_TRAIN, reshuffle_each_iteration=True).repeat(-1)\n",
    "\n",
    "\n",
    "# Set batch size\n",
    "dataset_train = dataset_train.batch(batch_size=BATCH_SIZE_TRAIN)\n",
    "dataset_validation = dataset_validation.batch(batch_size=BATCH_SIZE_VALIDATION)\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Creation -- Keras API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "\n",
    "    param_gen = GAN.topologies.Gen_param_structure()\n",
    "\n",
    "    param_gen.block_conv_channels = USE_GEN_net_conv_Channels\n",
    "    param_gen.block_conv_layers = USE_GEN_net_conv_Layers\n",
    "    param_gen.n_blocks = len(param_gen.block_conv_channels)\n",
    "    param_gen.latent_dim = (voxels_X,voxels_Y,voxels_Z,1)\n",
    "\n",
    "    param_gen.use_tanh_out = USE_GEN_TANH_OUT\n",
    "    param_gen.use_sigmoid_out = USE_GEN_SIGMOID_OUT\n",
    "    param_gen.segmentation_output = USE_GEN_SEGMENTATION\n",
    "    param_gen.segmentation_classes = USE_GEN_OBJECTIVE_SEGMENTATION_CLASES\n",
    "    param_gen.segmentation_layers = USE_GEN_OBJECTIVE_SEGMENTATION_LAYERS\n",
    "    param_gen.conv_segmentation_channels = USE_GEN_OBJECTIVE_CONV_SEGMENTATION_LAYERS\n",
    "    param_gen.segmentation_kernel_size = USE_GEN_SEGM_KERNEL_SIZE\n",
    "\n",
    "    param_gen.use_BatchNorm = USE_GEN_BATCH_NORM\n",
    "    param_gen.use_PixelNorm = USE_GEN_PIXEL_NORM\n",
    "    param_gen.use_He_scale = USE_GEN_HE_SCALLING\n",
    "    param_gen.initializer_std = USE_GEN_INI_STD\n",
    "\n",
    "    # Crea una instancia del modelo\n",
    "    gen_model = GAN.topologies.define_3D_Vnet_generator(param_gen)\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gen_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(gen_model, to_file=os.path.join(CHECKPOINT_PATH,'generator_model.png'), show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "\n",
    "    gen_model.compile(optimizer=keras.optimizers.Adam(lr=step_size, beta_1=0.9, beta_2=0.99, epsilon=1.0e-8),\n",
    "                  loss=GAN.losses.Vnet_compound_loss,\n",
    "                  metrics=[keras.metrics.mse])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx_cicle in range(CICLES_TRAIN):\n",
    "    \n",
    "    gen_model.fit(dataset_train, steps_per_epoch=STEPS_PER_EPOCH, epochs=EPOCHS_PER_PLOTS)\n",
    "    \n",
    "    show_images(dataset_validation)\n",
    "    \n",
    "    plt.close('all')\n",
    "    print('Cicle %d/%d'%((idx_cicle+1), CICLES_TRAIN))\n",
    "\n",
    "    if idx_cicle%CICLES_PER_SAVE == 0:\n",
    "        FILENAME_SAVE_USE = NETWORK_NAME+'_%dx%dx%d'%(voxels_X,voxels_Y,voxels_Z)\n",
    "        \n",
    "        GAN.train_support.save_model(gen_model, CHECKPOINT_PATH, FILENAME_SAVE_USE)\n",
    "    "
   ]
  }
 ],
 "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
}