Training_script.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Weak Lensing Maps / Cosmic Web Slices:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the script for training the cosmoGAN algorithm. If you want to train the algorithm on your own dataset be sure to do the following:\n",
    "\n",
    "* Set the output size accordingly. The \"output_size\" variable should be equal to the input dataset image dimensions (e.g. 256x256 px).\n",
    "* Currently the algorithm only works for images of the same height and width dimensions (e.g. 512x512 px etc).\n",
    "* The training dataset should be a python list/array of the following shape (No. of images x height x width (NHWC). The get_data() function will reshape the input the array by adding an extra dimension corresponding to the color dimension of the training images.\n",
    "* cosmoGAN can also be trained on RGB images (multidimensional arrays). For more information see the notebook on preparing Illustris data. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** Importing the libraries: **"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /mnt/c/Users/tamos/Desktop/PhD/for_github/functions.py:91: The name tf.random_normal is deprecated. Please use tf.random.normal instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "import sys\n",
    "import time\n",
    "import numpy as np\n",
    "import pprint\n",
    "import functions\n",
    "from functions import train_dcgan\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[name: \"/device:CPU:0\"\n",
      "device_type: \"CPU\"\n",
      "memory_limit: 268435456\n",
      "locality {\n",
      "}\n",
      "incarnation: 1305401620340962448\n",
      ", name: \"/device:XLA_CPU:0\"\n",
      "device_type: \"XLA_CPU\"\n",
      "memory_limit: 17179869184\n",
      "locality {\n",
      "}\n",
      "incarnation: 2944228838210510154\n",
      "physical_device_desc: \"device: XLA_CPU device\"\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.python.client import device_lib\n",
    "print(device_lib.list_local_devices())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**GPU settings:**\n",
    "\n",
    "The training script was tested with both GPU and CPU-only access. If you encounter any problems when training with a GPU, most likely, it's a memory issue, with the dataset being too big. If you need to disable the GPU, the following command can be used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "## If you encounter any problems running the training script with the GPU enable, you can disable it using:\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"-1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** A function to load the training data:**\n",
    "\n",
    "Here we choose the function to load the training data. If training on greyscale images (i.e. cosmic-web slices and weak lensing convergence maps), the data array is reshaped according to the user choice into No. x height x width x color format (NHWC) or no. x color dimesion x height x width (NCHW) data format. In case of Illustris data (i.e. 3-D RGB arrays) this is not required. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#dataset_type = \"Illustris\"\n",
    "dataset_type = \"CW_WL\"\n",
    "\n",
    "def get_data():\n",
    "    data = np.load(config.datafile, mmap_mode='r')\n",
    "\n",
    "    if dataset_type != \"Illustris\":\n",
    "        if config.data_format == 'NHWC':  ## This is the data format: Number of images x height x width x color dimension\n",
    "            data = np.expand_dims(data, axis=-1)\n",
    "        else: # 'NCHW'\n",
    "            data = np.expand_dims(data, axis=1)\n",
    "\n",
    "    return data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** Training settings: **"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "## The important parameters to get right here are the batch_size, c_dim & data_format. The batch_size here is \n",
    "## the number of dataset samples used per iteration of training. Also, it controls the number of samples\n",
    "## produced by the generator neural network after training. The c_dim parameter refers to the dimensionality \n",
    "## of the training dataset, i.e. for greyscale images arrays it's c_dim = 1 and for full-color RGB\n",
    "## arrays, it's c_dim = 3 (e.g. the Illustris dataset). The data_format also has to be set correctly, given the\n",
    "## dataset you are using. It is highly recommended to reshape your data into  NHWC format, if you are\n",
    "## training the GAN on your own data. \n",
    "\n",
    "flags = tf.app.flags\n",
    "flags.DEFINE_string(\"dataset\", \"cosmic_web\", \"The name of the dataset\")\n",
    "#flags.DEFINE_string(\"datafile\", \"./data/cosmogan_maps_256_8k_1.npy\", \"Training dataset file location\")\n",
    "flags.DEFINE_string(\"datafile\", \"./data/stack_z0p0_fR_f1_f7_a_250_test.npy\", \"Training dataset file location\")\n",
    "flags.DEFINE_integer(\"epoch\", 5000, \"Epochs to train before stopping\")\n",
    "flags.DEFINE_float(\"learning_rate\", 0.00009, \"The learning rate parameter\")\n",
    "flags.DEFINE_float(\"beta1\", 0.5, \"Momentum term of adam (default value: 0.5)\")\n",
    "flags.DEFINE_float(\"flip_labels\", 0.01, \"Probability of flipping labels (default value: 0.01)\")\n",
    "flags.DEFINE_integer(\"z_dim\", 256, \"Dimension of noise vector z\")\n",
    "flags.DEFINE_integer(\"nd_layers\", 4, \"Number of discriminator convolutional layers (default value: 4)\")\n",
    "flags.DEFINE_integer(\"ng_layers\", 4, \"Number of generator conv_T layers (default value: 4)\")\n",
    "flags.DEFINE_integer(\"gf_dim\", 64, \"Dimension of generator filters in last conv layer (default value: 64)\")\n",
    "flags.DEFINE_integer(\"df_dim\", 64, \"Dimension of discriminator filters in first conv layer (default value: 64)\")\n",
    "flags.DEFINE_integer(\"batch_size\", 64, \"The size of batch images (default value: 64)\")\n",
    "flags.DEFINE_integer(\"output_size\", 256, \"The size of the output images to produce (default value: 64 for weak lensing maps and 256 for CW slices)\")\n",
    "flags.DEFINE_integer(\"c_dim\", 1, \"Dimension of image color. 1 = greyscale image, 3 = RGB image\")\n",
    "flags.DEFINE_string(\"data_format\", \"NHWC\", \"data format (NHWC = No. x height x width x color dimension while NCHW = no. x color dimesion x height x width)\")\n",
    "flags.DEFINE_boolean(\"transpose_matmul_b\", False, \"Transpose matmul B matrix for performance [False]\")\n",
    "flags.DEFINE_string(\"checkpoint_dir\", \"./checkpoints/checkpoint_name\", \"Directory name to save the checkpoints (default value: checkpoint)\")\n",
    "flags.DEFINE_string(\"experiment\", \"run_0\", \"Tensorboard run directory name (run_0)\")\n",
    "flags.DEFINE_boolean(\"save_every_step\", False, \"Save a checkpoint after every step (default value: False)\")\n",
    "flags.DEFINE_boolean(\"verbose\", True, \"print loss on every step (default value: False)\")\n",
    "config = flags.FLAGS\n",
    "\n",
    "tf.app.flags.DEFINE_string('f', '', 'kernel')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint.PrettyPrinter().pprint(config.__flags)\n",
    "train_dcgan(get_data(), config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test_env",
   "language": "python",
   "name": "test_env"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}