Train_GAN_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_GAN_Network'\n",
    "\n",
    "# Dataset location\n",
    "DATASET_PATH = './datasets/'\n",
    "\n",
    "# Checkpoint location\n",
    "CHECKPOINT_PATH = \"./Outputs/\"+NETWORK_NAME+\"/\"\n",
    "\n",
    "# Path to tensorboard desired output\n",
    "TENSORBOARD_BASE_PATH = \"./TensorBoard_output\"\n",
    "TENSORBOARD_OUT_PATH = os.path.join(TENSORBOARD_BASE_PATH,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": [
    "# Pre-Train: If True the generator is loaded from a previously \n",
    "# trained generator, trained with the \"./Train_Standard_Model.ipynb\" notebook.\n",
    "# The checkpoint is loaded from \"PRE_TRAINED_GENERATOR_PATH\" using name \"PRE_TRAINED_GENERATOR_NAME\"\n",
    "Pre_Train = True\n",
    "PRE_TRAINED_GENERATOR_PATH = './trained_models/Sup_Pre_Train_No_GAN_Loss/'\n",
    "PRE_TRAINED_GENERATOR_NAME = 'Sup_Pre_Train_No_GAN_Loss'\n",
    "\n",
    "# If True the adversarial gradient is restricted to the conditional generator network\n",
    "RestrictedGrad = True\n",
    "# If True the segementator networks losses (DICE and L2) are used to regularize training\n",
    "SupervisedLoss = True\n",
    "\n",
    "\n",
    "\n",
    "# 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",
    "# --------------------------------------------------------------------------------\n",
    "# ---------------------- GENERATOR / SEGMENTATION NET ----------------------------\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",
    "# 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\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------------------------\n",
    "# ---------------------- CONDITIONAL GENERATOR -----------------------------------\n",
    "# --------------------------------------------------------------------------------\n",
    "\n",
    "# Network convolutional channels by resolution level\n",
    "USE_COND_GEN_net_conv_Channels = [8,8,8,8,8]\n",
    "# Convolutional layers\n",
    "USE_COND_GEN_net_conv_Layers = 5\n",
    "# Hyperbolic tangent or Sigmoid output for generator\n",
    "USE_COND_GEN_TANH_OUT = False\n",
    "USE_COND_GEN_SIGMOID_OUT = False\n",
    "# Use batch normalization for generator training\n",
    "USE_COND_GEN_BATCH_NORM = False\n",
    "# Use pixel normalization for generator training\n",
    "USE_COND_GEN_PIXEL_NORM = True\n",
    "# Use He scalling of weights for generator\n",
    "USE_COND_GEN_HE_SCALLING = True\n",
    "# Wheight initialization standard deviation\n",
    "USE_COND_GEN_INI_STD = 1.0\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------------------------\n",
    "# ---------------------- DISCRIMNATOR --------------------------------------------\n",
    "# --------------------------------------------------------------------------------\n",
    "# If true the discriminator recieves the generator output AND input latent space\n",
    "IS_CONDITIONAL_DISC = True\n",
    "# Network convolutional channels by resolution level\n",
    "USE_DISC_net_conv_Channels = [4, 8, 16, 32, 64, 128, 256]\n",
    "# Use weight norm constraint\n",
    "USE_DISC_NORM_CONSTRAINT_SCALE = False\n",
    "# Use mini batch std\n",
    "USE_DISC_MINI_BATCH_STD = False\n",
    "\n"
   ]
  },
  {
   "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",
    "# Volume 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_gen = 0.0001 \n",
    "step_size_disc = 0.0005 \n",
    "step_size_segm = 0.001 \n",
    "\n",
    "# Set training steps for each network\n",
    "TRAINING_FUNCTION_DISC = GAN.train_support.train_step_discriminator_conditional_3D_GAN_tf\n",
    "TRAINING_FUNCTION_GEN = GAN.train_support.train_step_generator_and_segmentator_conditional_3D_GAN_tf\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\"\n",
    "\n",
    "\n",
    "# Total number of training steps to perform\n",
    "STEPS_RUN = 100001\n",
    "# Number of steps per information print\n",
    "STEPS_PER_PRINT = 100\n",
    "# Number of steps per plotting of progress\n",
    "STEPS_PER_PLOTS = 500\n",
    "# Number of steps per checkpoint\n",
    "STEPS_PER_SAVE = 500\n",
    "# Initial number of discriminator training steps\n",
    "steps_disc_per_gen_ini = 500\n",
    "# Number of discriminator training steps per generator training steps\n",
    "steps_disc_per_gen_loop = 5\n",
    "\n"
   ]
  },
  {
   "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)\n",
    "File_mng.check_create_path('TENSORBOARD_BASE_PATH', TENSORBOARD_BASE_PATH)\n",
    "File_mng.check_create_path('TENSORBOARD_OUT_PATH', TENSORBOARD_OUT_PATH, clear_folder=CLEAR_OUTS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load sampling CDF\n",
    "cdf_coef = [1.0]\n",
    "if not UNIFORM_SAMPLING:\n",
    "    cdf_coef = np.load(CDF_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set multip-GPU mirror strategy\n",
    "strategy = tf.distribute.MirroredStrategy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset reader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset list\n",
    "\n",
    "shape_X = int(voxels_X)\n",
    "shape_Y = int(voxels_Y)\n",
    "shape_Z = int(voxels_Z)\n",
    "\n",
    "data_size = np.array([DATASET_X_size, \n",
    "                      DATASET_Y_size, \n",
    "                      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_GAN = tf.data.TFRecordDataset(PATH_TFRECORD_TRAIN)\n",
    "dataset_validation_GAN = tf.data.TFRecordDataset(PATH_TFRECORD_VALIDATION)\n",
    "\n",
    "if shape_X <= 32: \n",
    "    dataset_train_GAN = dataset_train_GAN.cache()\n",
    "    dataset_validation_GAN = dataset_validation_GAN.cache()\n",
    "    print('Using cache for dataset: %s'%train_dataset_name)\n",
    "\n",
    "# Create train dataset with transformations\n",
    "dataset_train_GAN = dataset_train_GAN.map(lambda x: DH.tf_read_sample_file(x, \n",
    "                                                                              data_size, \n",
    "                                                                              input_size_this, \n",
    "                                                                              not_transformed = True,\n",
    "                                                                              cdf_sampler_coef=cdf_coef))\n",
    "# Create validation dataset, whole image\n",
    "dataset_validation_GAN = dataset_validation_GAN.map(lambda x: DH.tf_read_raw_sample(x, data_size))\n",
    "\n",
    "# Shuffle the train dataset\n",
    "dataset_train_GAN = dataset_train_GAN.shuffle(buffer_size=BUFFER_SIZE_TRAIN, reshuffle_each_iteration=True).repeat(-1)\n",
    "\n",
    "\n",
    "# Set batch size\n",
    "dataset_train_GAN = dataset_train_GAN.batch(batch_size=BATCH_SIZE_TRAIN)\n",
    "dataset_validation_GAN = dataset_validation_GAN.batch(batch_size=BATCH_SIZE_VALIDATION)\n",
    "\n",
    "# Create distributed datasets\n",
    "dist_dataset_train_GAN = strategy.experimental_distribute_dataset(dataset_train_GAN)\n",
    "dist_dataset_validation_GAN = strategy.experimental_distribute_dataset(dataset_validation_GAN)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Creation -- Keras API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Segmentation V-Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "\n",
    "    if Pre_Train:\n",
    "        custom_layers_dict = {'HeScale': GAN.layers.HeScale,\n",
    "                              'BiasLayer': GAN.layers.BiasLayer,\n",
    "                              'PixelNormalization': GAN.layers.PixelNormalization}\n",
    "        segm_model = GAN.train_support.load_model(PRE_TRAINED_GENERATOR_PATH, \n",
    "                                                  PRE_TRAINED_GENERATOR_NAME, \n",
    "                                                  custom_obj_dict = custom_layers_dict)\n",
    "        \n",
    "    else:\n",
    "        param_segm = GAN.topologies.Gen_param_structure()\n",
    "\n",
    "        param_segm.block_conv_layers = USE_GEN_net_conv_Layers\n",
    "        param_segm.block_conv_channels = USE_GEN_net_conv_Channels\n",
    "        param_segm.n_blocks = len(param_segm.block_conv_channels)\n",
    "        param_segm.latent_dim = (voxels_X,voxels_Y,voxels_Z,1)\n",
    "\n",
    "        param_segm.use_tanh_out = USE_GEN_TANH_OUT\n",
    "        param_segm.use_sigmoid_out = USE_GEN_SIGMOID_OUT\n",
    "        param_segm.segmentation_output = USE_GEN_SEGMENTATION\n",
    "        param_segm.segmentation_classes = USE_GEN_OBJECTIVE_SEGMENTATION_CLASES\n",
    "        param_segm.segmentation_layers = USE_GEN_OBJECTIVE_SEGMENTATION_LAYERS\n",
    "        param_segm.conv_segmentation_channels = USE_GEN_OBJECTIVE_CONV_SEGMENTATION_LAYERS\n",
    "        param_segm.segmentation_kernel_size = USE_GEN_SEGM_KERNEL_SIZE\n",
    "\n",
    "        param_segm.use_BatchNorm = USE_GEN_BATCH_NORM\n",
    "        param_segm.use_PixelNorm = USE_GEN_PIXEL_NORM\n",
    "        param_segm.use_He_scale = USE_GEN_HE_SCALLING\n",
    "        param_segm.initializer_std = USE_GEN_INI_STD\n",
    "\n",
    "        # Crea una instancia del modelo\n",
    "        segm_model = GAN.topologies.define_3D_Vnet_generator(param_segm)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "segm_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(segm_model, to_file=os.path.join(CHECKPOINT_PATH,'segmentator_model.png'), show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conditional Generator Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    param_gen = GAN.topologies.Gen_param_structure()\n",
    "\n",
    "    param_gen.block_conv_layers = USE_COND_GEN_net_conv_Layers\n",
    "    param_gen.block_conv_channels = USE_COND_GEN_net_conv_Channels\n",
    "    param_gen.n_blocks = len(param_gen.block_conv_channels)\n",
    "    \n",
    "    if SupervisedLoss:\n",
    "        # The input is the synth. CT and the segmentation output of previous network\n",
    "        param_gen.latent_dim = (voxels_X,voxels_Y,voxels_Z,USE_GEN_OBJECTIVE_SEGMENTATION_CLASES+1)\n",
    "    else:\n",
    "        # The input is only the synth. CT output of previous network\n",
    "        param_gen.latent_dim = (voxels_X,voxels_Y,voxels_Z,1)\n",
    "        \n",
    "\n",
    "    param_gen.use_tanh_out = USE_COND_GEN_TANH_OUT\n",
    "    param_gen.use_sigmoid_out = USE_COND_GEN_SIGMOID_OUT\n",
    "\n",
    "    param_gen.use_BatchNorm = USE_COND_GEN_BATCH_NORM\n",
    "    param_gen.use_PixelNorm = USE_COND_GEN_PIXEL_NORM\n",
    "    param_gen.use_He_scale = USE_COND_GEN_HE_SCALLING\n",
    "    param_gen.initializer_std = USE_COND_GEN_INI_STD\n",
    "\n",
    "    # Create model instance\n",
    "    gen_model = GAN.topologies.define_3D_Convolutional_generator(param_gen, name_appendix='_Texturator')\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,'texturator_model.png'), show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Composed Generator Network\n",
    "\n",
    "This model is just the union of the previous networks. This way is easier to control the gradient flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    \n",
    "    inputRefImage = keras.Input(shape=segm_model.input.shape[1:], name=\"Input_PET_GAN\")\n",
    "\n",
    "    # Set to segm range...\n",
    "    t_nada = (inputRefImage/2.0)+0.5\n",
    "    \n",
    "    # Apply 3D U-Net\n",
    "    out_v_net_segm = segm_model(t_nada)\n",
    "    \n",
    "    # Apply texture network\n",
    "    if SupervisedLoss:\n",
    "        out_texture_gan = gen_model(out_v_net_segm)\n",
    "    else:\n",
    "        out_texture, out_segm = tf.split(out_v_net_segm, [1,4], axis=-1)\n",
    "        out_texture_gan = gen_model(out_texture)\n",
    "        \n",
    "\n",
    "    comp_gen_model = keras.Model(inputRefImage, out_texture_gan, name='Composed_Generator')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "comp_gen_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(comp_gen_model, to_file=os.path.join(CHECKPOINT_PATH,'generator_model.png'), show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conditional Discriminator Net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with strategy.scope():\n",
    "    \n",
    "    \n",
    "    \n",
    "    param_disc = GAN.topologies.Disc_param_structure()\n",
    "    param_disc.conditional = IS_CONDITIONAL_DISC\n",
    "    param_disc.block_conv_channels = USE_DISC_net_conv_Channels\n",
    "    \n",
    "    # Nomber of pooling operations\n",
    "    LOW_RES_POW = len(USE_DISC_net_conv_Channels)-3\n",
    "    # Lowest resolution inside the convolutional discriminator befores using fully connected layers\n",
    "    voxels_X_low_res = int(voxels_X/(2**LOW_RES_POW))\n",
    "    voxels_Y_low_res = int(voxels_Y/(2**LOW_RES_POW))\n",
    "    voxels_Z_low_res = int(voxels_Z/(2**LOW_RES_POW))\n",
    "    # This is the input of the frst block. This is only to simplify contruction code.\n",
    "    param_disc.input_shape = (voxels_X_low_res,voxels_Y_low_res,voxels_Z_low_res,1)\n",
    "    \n",
    "    param_disc.use_norm_contrain_scale = USE_DISC_NORM_CONSTRAINT_SCALE\n",
    "    param_disc.use_minibatch_stdev = USE_DISC_MINI_BATCH_STD\n",
    "    \n",
    "    param_disc.n_blocks = LOW_RES_POW+1\n",
    "    \n",
    "    param_disc.downSampling_layers  = len(param_disc.block_conv_channels) \n",
    "\n",
    "    # Create a model instance\n",
    "    disc_model_list = GAN.topologies.define_discriminator_3D(param_disc)\n",
    "    # Keep the last model, since the previous function creates a list of\n",
    "    # models for different resolutions (used in ProGAN, in an other notebook)\n",
    "    disc_model = disc_model_list[-1][0]\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "disc_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(disc_model, to_file=os.path.join(CHECKPOINT_PATH,'discrimnator_model.png'), show_shapes=True, show_layer_names=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create optimizers and metrics for training \n",
    "with strategy.scope():\n",
    "    \n",
    "    optimizer_segm = keras.optimizers.Adam(lr=step_size_segm, beta_1=0.9, beta_2=0.99, epsilon=1.0e-8)    \n",
    "    optimizer_gen = keras.optimizers.Adam(lr=step_size_gen, beta_1=0, beta_2=0.99, epsilon=1.0e-8)\n",
    "    optimizer_disc = keras.optimizers.RMSprop(step_size_disc)\n",
    "    \n",
    "    train_d1_loss = keras.metrics.Mean(name='train_d1_loss')\n",
    "    train_d2_loss = keras.metrics.Mean(name='train_d2_loss')\n",
    "    train_dgrad_loss = keras.metrics.Mean(name='train_dgrad_loss')\n",
    "    \n",
    "    train_g_loss = keras.metrics.Mean(name='train_g_loss')\n",
    "    train_g_comp_loss = keras.metrics.Mean(name='train_g_comp_loss')\n",
    "    \n",
    "    train_s1_loss = keras.metrics.Mean(name='train_s1_loss')\n",
    "    train_s2_loss = keras.metrics.Mean(name='train_s2_loss')\n",
    "    \n",
    "loss_plot = np.zeros((STEPS_RUN*STEPS_PER_PRINT,6))      \n",
    "loss_val_plot = np.zeros((STEPS_RUN*STEPS_PER_PLOTS,5))\n",
    "loss_std_val_plot = np.zeros((STEPS_RUN*STEPS_PER_PLOTS,5))\n",
    "axis_val_plot = np.zeros((STEPS_RUN*STEPS_PER_PLOTS,1))\n",
    "idx_plot = 0\n",
    "\n",
    "# train normal or straight-through models\n",
    "train_step_gen_tf_this = tf.function(TRAINING_FUNCTION_GEN)\n",
    "train_step_disc_tf_this = tf.function(TRAINING_FUNCTION_DISC)\n",
    "\n",
    "# Plot titles... decoration...\n",
    "titulos_use = ['disc. Real','disc. Fake','gen. Wasserstein','gen. MSE','Segm. DICE','Segm. L2']\n",
    "\n",
    "\n",
    "step_mean_d1_loss = 0\n",
    "step_mean_d2_loss = 0\n",
    "step_mean_dgrad_loss = 0\n",
    "step_mean_g_loss = 0\n",
    "step_mean_g_comp_loss = 0\n",
    "step_mean_s1_loss = 0\n",
    "step_mean_s2_loss = 0\n",
    "\n",
    "# Create tensorboard writer\n",
    "tb_writer = tf.summary.create_file_writer(TENSORBOARD_OUT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('Discrimnator metrics:\\n\\t d1:\\t  Real volume critic score.\\n\\t d2:\\t  Fake volume critic score.\\n\\t dg:\\t  Gradient penalty.\\n\\t d:\\t  Total loss.')\n",
    "print('Generator metrics:\\n\\t g:\\t  Critic loss of generated sample (WGAN loss).\\n\\t g_comp:   L2 loss of generated sample.\\n\\t s1:\\t  DICE segementation loss.\\n\\t s2:\\t  L2 loss of segmentator net.')\n",
    "\n",
    "with tb_writer.as_default():\n",
    "    for step in range(STEPS_RUN):\n",
    "\n",
    "\n",
    "        # Train Discriminator\n",
    "        idx_disc_train = 0\n",
    "        disc_loss_d1_aux = 0\n",
    "        disc_loss_d2_aux = 0\n",
    "        disc_loss_dgrad_aux = 0\n",
    "\n",
    "        if step == 0:\n",
    "            steps_disc_per_gen = steps_disc_per_gen_ini\n",
    "        else:\n",
    "            steps_disc_per_gen = steps_disc_per_gen_loop\n",
    "\n",
    "        # ---------------------------------------------------------------------------------\n",
    "        # --------------------------- TRAIN DISCRIMINATOR ---------------------------------\n",
    "        # ---------------------------------------------------------------------------------\n",
    "        for in_NAC_PET, in_LABELS, in_CT in dist_dataset_train_GAN:\n",
    "\n",
    "            train_step_disc_tf_this(in_NAC_PET,\n",
    "                                    in_CT,\n",
    "                                    comp_gen_model, \n",
    "                                    disc_model, \n",
    "                                    train_d1_loss,\n",
    "                                    train_d2_loss,\n",
    "                                    train_dgrad_loss,\n",
    "                                    optimizer_disc,\n",
    "                                    strategy,\n",
    "                                    K_grad=10.0,\n",
    "                                   cross_sample_loss = False)\n",
    "\n",
    "            disc_loss_d1_aux += train_d1_loss.result()\n",
    "            disc_loss_d2_aux += train_d2_loss.result()\n",
    "            disc_loss_dgrad_aux += train_dgrad_loss.result()\n",
    "\n",
    "\n",
    "            print_l_real = disc_loss_d1_aux/(idx_disc_train+1)\n",
    "            print_l_fake = disc_loss_d2_aux/(idx_disc_train+1)\n",
    "            print_l_grad = disc_loss_dgrad_aux/(idx_disc_train+1)\n",
    "            print_l_tot = print_l_fake+print_l_real+print_l_grad\n",
    "            print('(Steps: %d) Training Discriminator: %d/%d\\tLoss: %.4f (Real:%.4f Fake:%.4f Grad.:%.4f)      '%(step,\n",
    "                                                                                                                  idx_disc_train+1,\n",
    "                                                                                                                  steps_disc_per_gen,\n",
    "                                                                                                                  print_l_tot,\n",
    "                                                                                                                  print_l_real,\n",
    "                                                                                                                  print_l_fake,\n",
    "                                                                                                                  print_l_grad), end='')\n",
    "            print('', end='\\r')\n",
    "\n",
    "            idx_disc_train += 1\n",
    "            if idx_disc_train == steps_disc_per_gen:\n",
    "                last_fake_score_module = tf.abs(print_l_fake)\n",
    "                break\n",
    "\n",
    "\n",
    "        # ---------------------------------------------------------------------------------\n",
    "        # --------------------------- TRAIN GENERATOR -------------------------------------\n",
    "        # ---------------------------------------------------------------------------------\n",
    "        for in_NAC_PET, in_LABELS, in_CT in dist_dataset_train_GAN:\n",
    "\n",
    "            print('(Steps: %d) Training Generator...                                                                             '%(step), end='')\n",
    "            print('', end='\\r')\n",
    "\n",
    "            train_step_gen_tf_this(in_NAC_PET,\n",
    "                                   in_CT,\n",
    "                                   in_LABELS,\n",
    "                                   comp_gen_model, \n",
    "                                   disc_model, \n",
    "                                   segm_model,\n",
    "                                   train_g_loss,\n",
    "                                   train_g_comp_loss,\n",
    "                                   train_s1_loss,\n",
    "                                   train_s2_loss,\n",
    "                                   optimizer_gen,\n",
    "                                   strategy,\n",
    "                                   K_comp_loss = last_fake_score_module*10.0,\n",
    "                                   K_comp_segm_loss = GAN.losses.k_coupling, \n",
    "                                   norm_by_size = False,\n",
    "                                   split_train = RestrictedGrad,\n",
    "                                   train_segm = SupervisedLoss)\n",
    "\n",
    "\n",
    "\n",
    "            break\n",
    "\n",
    "\n",
    "        # ---------------------------------------------------------------------------------\n",
    "        # --------------------------- TRAIN METRICS ---------------------------------------\n",
    "        # ---------------------------------------------------------------------------------\n",
    "\n",
    "        loss_plot[step,0] = disc_loss_d1_aux/steps_disc_per_gen\n",
    "        loss_plot[step,1] = disc_loss_d2_aux/steps_disc_per_gen\n",
    "        loss_plot[step,2] = train_g_loss.result()\n",
    "        loss_plot[step,3] = train_g_comp_loss.result()    \n",
    "        loss_plot[step,4] = train_s1_loss.result()\n",
    "        loss_plot[step,5] = train_s2_loss.result()\n",
    "        \n",
    "        tf.summary.scalar(\"disc_loss_d1\", loss_plot[step,0], step=step)\n",
    "        tf.summary.scalar(\"disc_loss_d2\", loss_plot[step,1], step=step)\n",
    "        tf.summary.scalar(\"train_g_loss\", loss_plot[step,2], step=step)\n",
    "        tf.summary.scalar(\"train_g_comp_loss\", loss_plot[step,3], step=step)\n",
    "        tf.summary.scalar(\"train_s1_loss\", loss_plot[step,4], step=step)\n",
    "        tf.summary.scalar(\"train_s2_loss\", loss_plot[step,5], step=step)\n",
    "\n",
    "        train_d1_loss.reset_states()\n",
    "        train_d2_loss.reset_states()\n",
    "        train_dgrad_loss.reset_states()\n",
    "        train_g_loss.reset_states()\n",
    "        train_g_comp_loss.reset_states()\n",
    "        train_s1_loss.reset_states()\n",
    "        train_s2_loss.reset_states()\n",
    "\n",
    "        step_mean_d1_loss += loss_plot[step,0]\n",
    "        step_mean_d2_loss += loss_plot[step,1]\n",
    "        step_mean_dgrad_loss += disc_loss_dgrad_aux/steps_disc_per_gen\n",
    "        step_mean_g_loss += loss_plot[step,2]\n",
    "        step_mean_g_comp_loss += loss_plot[step,3]\n",
    "        step_mean_s1_loss += loss_plot[step,4]\n",
    "        step_mean_s2_loss += loss_plot[step,5]\n",
    "\n",
    "        if step%STEPS_PER_PRINT == 0:\n",
    "\n",
    "            if step != 0:\n",
    "                step_mean_d1_loss = step_mean_d1_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_d2_loss = step_mean_d2_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_g_loss = step_mean_g_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_g_comp_loss = step_mean_g_comp_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_s1_loss = step_mean_s1_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_s2_loss = step_mean_s2_loss/float(STEPS_PER_PRINT)\n",
    "                step_mean_dgrad_loss = step_mean_dgrad_loss/float(STEPS_PER_PRINT)\n",
    "\n",
    "            step_mean_dtot_loss = step_mean_d1_loss+step_mean_d2_loss+step_mean_dgrad_loss\n",
    "\n",
    "            print('>%d, d1=%.4f, d2=%.4f, dg=%.4f (d=%.4f) ; g=%.4f g_comp=%.4f s_1=%.4f s_2=%.4f ' % (step, \n",
    "                                                                                                       step_mean_d1_loss, \n",
    "                                                                                                       step_mean_d2_loss,\n",
    "                                                                                                       step_mean_dgrad_loss,\n",
    "                                                                                                       step_mean_dtot_loss,\n",
    "                                                                                                       step_mean_g_loss,\n",
    "                                                                                                       step_mean_g_comp_loss,\n",
    "                                                                                                       step_mean_s1_loss,\n",
    "                                                                                                       step_mean_s2_loss))\n",
    "\n",
    "            step_mean_d1_loss = 0\n",
    "            step_mean_d2_loss = 0\n",
    "            step_mean_dgrad_loss = 0\n",
    "            step_mean_g_loss = 0\n",
    "            step_mean_g_comp_loss = 0\n",
    "            step_mean_s1_loss = 0\n",
    "            step_mean_s2_loss = 0\n",
    "\n",
    "        if step%STEPS_PER_PLOTS == 0 and step > 0:\n",
    "\n",
    "            GAN.train_support.plot_loss(loss_plot, step, forma='2_cols', dpi_use=100, titulos=titulos_use, skip_n_initial=10)\n",
    "\n",
    "\n",
    "            (PET_INPUT_images, \n",
    "             CT_INPUT_images, \n",
    "             LABELS_INPUT_images, \n",
    "             CT_SYNTH_images, \n",
    "             SEGMENTED_images, \n",
    "             SCORE_images, \n",
    "             PSNR_images, \n",
    "             ME_images, \n",
    "             NMSE_images, \n",
    "             NCC_images) = GAN.train_support.validate_whole_volume(comp_gen_model, \n",
    "                                                                   disc_model, \n",
    "                                                                   dataset_validation_GAN, \n",
    "                                                                   [shape_X,shape_Y,shape_Z], \n",
    "                                                                   segm_net = SupervisedLoss, \n",
    "                                                                   s_model=segm_model,\n",
    "                                                                   single_image = False)\n",
    "\n",
    "            GAN.train_support.add_tensorboard_3Dimage(PET_INPUT_images, \"PET_INPUT\", 0, channel=0, min_val=-1.0, max_val=1.0)\n",
    "            GAN.train_support.add_tensorboard_3Dimage(CT_INPUT_images, \"CT_INPUT\", 0, channel=0, min_val=-1.0, max_val=1.0)\n",
    "            GAN.train_support.add_tensorboard_3Dimage(CT_SYNTH_images, \"CT_SYNTH\", 0, channel=0, min_val=-1.0, max_val=1.0)\n",
    "\n",
    "            if SupervisedLoss:\n",
    "                LABELS_INPUT = np.argmax(LABELS_INPUT_images, axis=-1)\n",
    "                LABELS_INPUT = np.expand_dims(LABELS_INPUT, axis=-1)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(LABELS_INPUT, \"LABELS\", 0, channel=0, min_val=0.0, max_val=3.0)\n",
    "\n",
    "                LABELS_SYNTH = np.argmax(SEGMENTED_images, axis=-1)\n",
    "                SEGMENTED_images = np.expand_dims(SEGMENTED_images, axis=-1)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(SEGMENTED_images, \"LABELS_SYNTH\", 0, channel=0, min_val=-1.0, max_val=1.0)\n",
    "\n",
    "                GAN.train_support.add_tensorboard_3Dimage(LABELS_INPUT_images, \"LABELS_INPUT_0\", 0, channel=0)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(LABELS_INPUT_images, \"LABELS_INPUT_1\", 0, channel=1)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(LABELS_INPUT_images, \"LABELS_INPUT_2\", 0, channel=2)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(LABELS_INPUT_images, \"LABELS_INPUT_3\", 0, channel=3)\n",
    "\n",
    "                GAN.train_support.add_tensorboard_3Dimage(SEGMENTED_images, \"SEGMENTED_0\", 0, channel=0)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(SEGMENTED_images, \"SEGMENTED_1\", 0, channel=1)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(SEGMENTED_images, \"SEGMENTED_2\", 0, channel=2)\n",
    "                GAN.train_support.add_tensorboard_3Dimage(SEGMENTED_images, \"SEGMENTED_3\", 0, channel=3)\n",
    "\n",
    "\n",
    "            print('')\n",
    "            GAN.train_support.plot_images3D_conditional(PET_INPUT_images,\n",
    "                                                       CT_INPUT_images, \n",
    "                                                       CT_SYNTH_images, \n",
    "                                                       segm_net = SupervisedLoss, \n",
    "                                                       X_segm=SEGMENTED_images, \n",
    "                                                       X_real_segm=LABELS_INPUT_images,\n",
    "                                                       dpi_use = 150)\n",
    "\n",
    "            print('Metrics:                              ')\n",
    "            print('\\tCritic Score:\\t mean=%0.5f ; std=%0.5f'%(np.array(SCORE_images).mean(),np.array(SCORE_images).std()))\n",
    "            print('\\tPSNR:\\t\\t mean=%0.5f ; std=%0.5f'%(np.array(PSNR_images).mean(),np.array(PSNR_images).std()))\n",
    "            print('\\tME:\\t\\t mean=%0.5f ; std=%0.5f'%(np.array(ME_images).mean(),np.array(ME_images).std()))\n",
    "            print('\\tNMSE:\\t\\t mean=%0.5f ; std=%0.5f'%(np.array(NMSE_images).mean(),np.array(NMSE_images).std()))\n",
    "            print('\\tNCC:\\t\\t mean=%0.5f ; std=%0.5f'%(np.array(NCC_images).mean(),np.array(NCC_images).std()))\n",
    "\n",
    "            tf.summary.scalar(\"valid_disc_score\", np.array(SCORE_images).mean(), step=0)\n",
    "            tf.summary.scalar(\"valid_PSNR\", np.array(PSNR_images).mean(), step=0)\n",
    "            tf.summary.scalar(\"valid_ME\", np.array(ME_images).mean(), step=0)\n",
    "            tf.summary.scalar(\"valid_NMSE\", np.array(NMSE_images).mean(), step=0)\n",
    "            tf.summary.scalar(\"valid_NCC\", np.array(NCC_images).mean(), step=0)\n",
    "\n",
    "\n",
    "            loss_val_plot[idx_plot, 0] = np.array(SCORE_images).mean()\n",
    "            loss_val_plot[idx_plot, 1] = np.array(PSNR_images).mean()\n",
    "            loss_val_plot[idx_plot, 2] = np.array(ME_images).mean()\n",
    "            loss_val_plot[idx_plot, 3] = np.array(NMSE_images).mean()\n",
    "            loss_val_plot[idx_plot, 4] = np.array(NCC_images).mean()\n",
    "\n",
    "            loss_std_val_plot[idx_plot, 0] = np.array(SCORE_images).std()\n",
    "            loss_std_val_plot[idx_plot, 1] = np.array(PSNR_images).std()\n",
    "            loss_std_val_plot[idx_plot, 2] = np.array(ME_images).std()\n",
    "            loss_std_val_plot[idx_plot, 3] = np.array(NMSE_images).std()\n",
    "            loss_std_val_plot[idx_plot, 4] = np.array(NCC_images).std()\n",
    "\n",
    "            axis_val_plot[idx_plot, 0] = step\n",
    "            idx_plot += 1\n",
    "\n",
    "\n",
    "            plt.figure(dpi=100)\n",
    "            plt.subplot(3,2,1)\n",
    "            plt.errorbar(axis_val_plot[:idx_plot,0], loss_val_plot[:idx_plot,0], yerr=loss_std_val_plot[:idx_plot,0], fmt='-o')\n",
    "            plt.grid('on')\n",
    "            plt.title('Critic score')\n",
    "            plt.subplot(3,2,2)\n",
    "            plt.errorbar(axis_val_plot[:idx_plot,0], loss_val_plot[:idx_plot,1], yerr=loss_std_val_plot[:idx_plot,1], fmt='-o')\n",
    "            plt.grid('on')\n",
    "            plt.title('PSNR')\n",
    "            plt.subplot(3,2,3)\n",
    "            plt.errorbar(axis_val_plot[:idx_plot,0], loss_val_plot[:idx_plot,2], yerr=loss_std_val_plot[:idx_plot,2], fmt='-o')\n",
    "            plt.grid('on')\n",
    "            plt.title('ME')\n",
    "            plt.subplot(3,2,4)\n",
    "            plt.errorbar(axis_val_plot[:idx_plot,0], loss_val_plot[:idx_plot,3], yerr=loss_std_val_plot[:idx_plot,3], fmt='-o')\n",
    "            plt.grid('on')\n",
    "            plt.title('NMSE')\n",
    "            plt.subplot(3,2,5)\n",
    "            plt.errorbar(axis_val_plot[:idx_plot,0], loss_val_plot[:idx_plot,4], yerr=loss_std_val_plot[:idx_plot,4], fmt='-o')\n",
    "            plt.grid('on')\n",
    "            plt.title('NCC')\n",
    "            plt.tight_layout()\n",
    "            plt.show()\n",
    "\n",
    "\n",
    "            plt.close('all')\n",
    "            \n",
    "            tb_writer.flush()\n",
    "\n",
    "        if step%STEPS_PER_SAVE == 0 and step > 0:\n",
    "            GAN.train_support.save_multiple_models([disc_model,comp_gen_model,segm_model], \n",
    "                                                         ['discrimnator','compound_generator','segmentatior'],\n",
    "                                                         CHECKPOINT_PATH,\n",
    "                                                         NETWORK_NAME,\n",
    "                                                         name_prefix='_%d'%step)"
   ]
  }
 ],
 "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
}