04_03_wgangp_faces_train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# WGAN-GP Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from models.WGANGP import WGANGP\n",
    "from utils.loaders import load_celeb\n",
    "\n",
    "import pickle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run params\n",
    "SECTION = 'gan'\n",
    "RUN_ID = '0003'\n",
    "DATA_NAME = 'celeb'\n",
    "RUN_FOLDER = 'run/{}/'.format(SECTION)\n",
    "RUN_FOLDER += '_'.join([RUN_ID, DATA_NAME])\n",
    "\n",
    "if not os.path.exists(RUN_FOLDER):\n",
    "    os.mkdir(RUN_FOLDER)\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'viz'))\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'images'))\n",
    "    os.mkdir(os.path.join(RUN_FOLDER, 'weights'))\n",
    "\n",
    "mode =  'build' #'load' #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 64\n",
    "IMAGE_SIZE = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train = load_celeb(DATA_NAME, IMAGE_SIZE, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train[0][0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow((x_train[0][0][0]+1)/2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gan = WGANGP(input_dim = (IMAGE_SIZE,IMAGE_SIZE,3)\n",
    "        , critic_conv_filters = [64,128,256,512]\n",
    "        , critic_conv_kernel_size = [5,5,5,5]\n",
    "        , critic_conv_strides = [2,2,2,2]\n",
    "        , critic_batch_norm_momentum = None\n",
    "        , critic_activation = 'leaky_relu'\n",
    "        , critic_dropout_rate = None\n",
    "        , critic_learning_rate = 0.0002\n",
    "        , generator_initial_dense_layer_size = (4, 4, 512)\n",
    "        , generator_upsample = [1,1,1,1]\n",
    "        , generator_conv_filters = [256,128,64,3]\n",
    "        , generator_conv_kernel_size = [5,5,5,5]\n",
    "        , generator_conv_strides = [2,2,2,2]\n",
    "        , generator_batch_norm_momentum = 0.9\n",
    "        , generator_activation = 'leaky_relu'\n",
    "        , generator_dropout_rate = None\n",
    "        , generator_learning_rate = 0.0002\n",
    "        , optimiser = 'adam'\n",
    "        , grad_weight = 10\n",
    "        , z_dim = 100\n",
    "        , batch_size = BATCH_SIZE\n",
    "        )\n",
    "\n",
    "if mode == 'build':\n",
    "    gan.save(RUN_FOLDER)\n",
    "\n",
    "else:\n",
    "    gan.load_weights(os.path.join(RUN_FOLDER, 'weights/weights.h5'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gan.critic.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gan.generator.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 6000\n",
    "PRINT_EVERY_N_BATCHES = 5\n",
    "N_CRITIC = 5\n",
    "BATCH_SIZE = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gan.train(     \n",
    "    x_train\n",
    "    , batch_size = BATCH_SIZE\n",
    "    , epochs = EPOCHS\n",
    "    , run_folder = RUN_FOLDER\n",
    "    , print_every_n_batches = PRINT_EVERY_N_BATCHES\n",
    "    , n_critic = N_CRITIC\n",
    "    , using_generator = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.plot([x[0] for x in gan.d_losses], color='black', linewidth=0.25)\n",
    "\n",
    "plt.plot([x[1] for x in gan.d_losses], color='green', linewidth=0.25)\n",
    "plt.plot([x[2] for x in gan.d_losses], color='red', linewidth=0.25)\n",
    "plt.plot(gan.g_losses, color='orange', linewidth=0.25)\n",
    "\n",
    "plt.xlabel('batch', fontsize=18)\n",
    "plt.ylabel('loss', fontsize=16)\n",
    "\n",
    "plt.xlim(0, 2000)\n",
    "# plt.ylim(0, 2)\n",
    "\n",
    "plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gdl_code",
   "language": "python",
   "name": "gdl_code"
  },
  "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": 2
}