{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3e818ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import autoreload\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e4aa92c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.8/site-packages/scipy/__init__.py:138: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.0)\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion} is required for this version of \"\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.layers import Input , Dense, Lambda, Normalization, Concatenate\n",
    "from tensorflow.keras.activations import relu, sigmoid, tanh\n",
    "from tensorflow.keras import backend as K\n",
    "from tensorflow.keras.optimizers import Adam, SGD, schedules\n",
    "from tensorflow_addons.optimizers import AdamW, extend_with_decoupled_weight_decay\n",
    "\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, TerminateOnNaN, TensorBoard, \\\n",
    "                                        EarlyStopping, ReduceLROnPlateau, CSVLogger, Callback\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "\n",
    "import gc\n",
    "\n",
    "from tqdm import tqdm "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "63cfa87b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import resource\n",
    "from tensorflow.compat.v1 import set_random_seed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa74ae9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelCheckpoint(Callback):\n",
    "    def __init__(self, model, path, pre_loss = np.inf):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.model = model\n",
    "        self.path = path\n",
    "        if not os.path.exists(path): os.makedirs(path)\n",
    "        self.best_loss = pre_loss\n",
    "        \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        loss = logs['val_loss']\n",
    "        if loss < self.best_loss:\n",
    "            print('\\nSaving model to {}'.format(self.path.format(epoch=epoch, loss=loss)))\n",
    "            print(', Validation loss decreased from {} to {}.\\n'.format(self.best_loss, loss))\n",
    "            self.model.save_weights(self.path.format(epoch=epoch, loss=loss), overwrite=True)\n",
    "            self.best_loss = loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "86dc4dc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClearMemory(Callback):\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        gc.collect()\n",
    "        print('Memory usage: ', resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)\n",
    "        \n",
    "    def on_batch_end(self, batch, logs=None):\n",
    "        if int(K.get_value(self.model.optimizer.iterations))%50 == 0:\n",
    "            gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "68251791",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchCallback(Callback):\n",
    "    def __init__(self, callbacks_tensorboard, logdir):\n",
    "        self.tb_callback = callbacks_tensorboard\n",
    "        self.logdir = logdir\n",
    "        self.train_writer = tf.summary.create_file_writer(os.path.join(self.logdir,\"training\"))\n",
    "        self.val_writer = tf.summary.create_file_writer(os.path.join(self.logdir,\"validation\"))\n",
    "        \n",
    "        self.loss_tag = 'Epoch Summary/'\n",
    "    \n",
    "    def on_train_begin(self, logs={}):\n",
    "        global_step = int(K.get_value(self.model.optimizer.iterations))\n",
    "        print('[INFO][BatchCallback] global step:{}'.format(global_step))\n",
    "        \n",
    "    def on_train_batch_end(self, batch, logs={}):\n",
    "        logs.update({'global_step': K.get_value(self.model.optimizer.iterations)})\n",
    "        logs.update({'loss': logs['loss']})\n",
    "        logs.update({'maskloss': logs['maskloss']})\n",
    "        logs.update({'spectralloss': logs['spectralloss']})\n",
    "\n",
    "        self._write_log(logs)\n",
    "        \n",
    "    def on_epoch_end(self, epoch, logs={}):\n",
    "        train_writer = self.train_writer\n",
    "        with train_writer.as_default():\n",
    "            for key, val in logs.items():\n",
    "                if key in ['loss']: tf.summary.scalar('loss', val, step=epoch)\n",
    "        train_writer.flush()\n",
    "        \n",
    "        val_writer = self.val_writer\n",
    "        with val_writer.as_default():\n",
    "            for key, val in logs.items():\n",
    "                if key in ['loss','val_loss']: tf.summary.scalar('loss', val, step=epoch)\n",
    "        val_writer.flush()\n",
    "        \n",
    "    def _write_log(self, logs):\n",
    "        writer = self.train_writer\n",
    "        with writer.as_default():\n",
    "            for key, val in logs.items():\n",
    "                if key in ['maskloss', 'dfalphaloss', 'spectralloss']:\n",
    "                    tag = 'Train(Metric)/' + key\n",
    "                    tf.summary.scalar(tag, val, step=logs['global_step'])\n",
    "                if key in ['loss']:\n",
    "                    tag = 'Train(Total)/' + key.upper()\n",
    "                    tf.summary.scalar(tag, val, step=logs['global_step'])\n",
    "                \n",
    "        writer.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbd6506",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e612cb6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cosine_decay_with_warmup(global_step,\n",
    "                             learning_rate_base,\n",
    "                             total_steps,\n",
    "                             warmup_learning_rate=0.0,\n",
    "                             warmup_steps=0,\n",
    "                             hold_base_rate_steps=0,\n",
    "                             steps_per_epoch=0):\n",
    "    \"\"\"Cosine decay schedule with warm up period.\n",
    "    Cosine annealing learning rate as described in:\n",
    "      Loshchilov and Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts.\n",
    "      ICLR 2017. https://arxiv.org/abs/1608.03983\n",
    "    In this schedule, the learning rate grows linearly from warmup_learning_rate\n",
    "    to learning_rate_base for warmup_steps, then transitions to a cosine decay\n",
    "    schedule.\n",
    "    Arguments:\n",
    "        global_step {int} -- global step.\n",
    "        learning_rate_base {float} -- base learning rate.\n",
    "        total_steps {int} -- total number of training steps.\n",
    "    Keyword Arguments:\n",
    "        warmup_learning_rate {float} -- initial learning rate for warm up. (default: {0.0})\n",
    "        warmup_steps {int} -- number of warmup steps. (default: {0})\n",
    "        hold_base_rate_steps {int} -- Optional number of steps to hold base learning rate\n",
    "                                    before decaying. (default: {0})\n",
    "    Returns:\n",
    "      a float representing learning rate.\n",
    "    Raises:\n",
    "      ValueError: if warmup_learning_rate is larger than learning_rate_base,\n",
    "        or if warmup_steps is larger than total_steps.\n",
    "    \"\"\"\n",
    "\n",
    "    if total_steps < warmup_steps:\n",
    "        raise ValueError('total_steps must be larger or equal to '\n",
    "                         'warmup_steps.')\n",
    "    learning_rate = 0.5 * learning_rate_base * (1 + np.cos(\n",
    "        np.pi *\n",
    "        (global_step - warmup_steps - hold_base_rate_steps\n",
    "         ) / float(total_steps - warmup_steps - hold_base_rate_steps)))\n",
    "#     learning_rate = learning_rate_base * 0.95 ** ((global_step - warmup_steps - hold_base_rate_steps) / (steps_per_epoch/25))\n",
    "    if hold_base_rate_steps > 0:\n",
    "        learning_rate = np.where(global_step > warmup_steps + hold_base_rate_steps,\n",
    "                                 learning_rate, learning_rate_base)\n",
    "    if warmup_steps > 0:\n",
    "        if learning_rate_base < warmup_learning_rate:\n",
    "            raise ValueError('learning_rate_base must be larger or equal to '\n",
    "                             'warmup_learning_rate.')\n",
    "        slope = (learning_rate_base - warmup_learning_rate) / warmup_steps\n",
    "        warmup_rate = slope * global_step + warmup_learning_rate\n",
    "        learning_rate = np.where(global_step < warmup_steps, warmup_rate,\n",
    "                                 learning_rate)\n",
    "    return np.where(global_step > total_steps, 0.0, learning_rate)\n",
    "\n",
    "\n",
    "class WarmUpCosineDecayScheduler(keras.callbacks.Callback):\n",
    "    \"\"\"Cosine decay with warmup learning rate scheduler\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 learning_rate_base,\n",
    "                 total_steps,\n",
    "                 global_step_init=0,\n",
    "                 warmup_learning_rate=0.0,\n",
    "                 warmup_steps=0,\n",
    "                 hold_base_rate_steps=0,\n",
    "                 mini_lr=0,\n",
    "                 steps_per_epoch=0,\n",
    "                 verbose=0):\n",
    "        \"\"\"Constructor for cosine decay with warmup learning rate scheduler.\n",
    "    Arguments:\n",
    "        learning_rate_base {float} -- base learning rate.\n",
    "        total_steps {int} -- total number of training steps.\n",
    "    Keyword Arguments:\n",
    "        global_step_init {int} -- initial global step, e.g. from previous checkpoint.\n",
    "        warmup_learning_rate {float} -- initial learning rate for warm up. (default: {0.0})\n",
    "        warmup_steps {int} -- number of warmup steps. (default: {0})\n",
    "        hold_base_rate_steps {int} -- Optional number of steps to hold base learning rate\n",
    "                                    before decaying. (default: {0})\n",
    "        verbose {int} -- 0: quiet, 1: update messages. (default: {0})\n",
    "        \"\"\"\n",
    "\n",
    "        super(WarmUpCosineDecayScheduler, self).__init__()\n",
    "        self.learning_rate_base = learning_rate_base\n",
    "        self.total_steps = total_steps\n",
    "        self.global_step = global_step_init\n",
    "        self.warmup_learning_rate = warmup_learning_rate\n",
    "        self.warmup_steps = warmup_steps\n",
    "        self.hold_base_rate_steps = hold_base_rate_steps\n",
    "        self.verbose = verbose\n",
    "        self.steps_per_epoch = steps_per_epoch\n",
    "        self.mini_lr = mini_lr\n",
    "#         self.learning_rates = []\n",
    "\n",
    "    def on_train_batch_end(self, batch, logs=None):\n",
    "        self.global_step = self.global_step + 1\n",
    "#         lr = K.get_value(self.model.optimizer.lr)\n",
    "#         self.learning_rates.append(lr)\n",
    "\n",
    "    def on_batch_begin(self, batch, logs=None):\n",
    "        lr = cosine_decay_with_warmup(global_step=self.global_step,\n",
    "                                      learning_rate_base=self.learning_rate_base,\n",
    "                                      total_steps=self.total_steps,\n",
    "                                      warmup_learning_rate=self.warmup_learning_rate,\n",
    "                                      warmup_steps=self.warmup_steps,\n",
    "                                      hold_base_rate_steps=self.hold_base_rate_steps,\n",
    "                                      steps_per_epoch=self.steps_per_epoch)\n",
    "        if lr < self.mini_lr: lr = self.mini_lr\n",
    "        K.set_value(self.model.optimizer.lr, lr)\n",
    "        if self.verbose > 0:\n",
    "            print('\\nBatch %05d: setting learning '\n",
    "                  'rate to %s.' % (self.global_step + 1, lr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "079e842f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainCallback(Callback):\n",
    "    def __init__(self, model):\n",
    "        self.model = model\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if epoch>=5:\n",
    "            for layer in self.model.layers:\n",
    "                layer.trainable=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2caaf51f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"   # see issue #152\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "# from tensorflow.keras.models import Sequential\n",
    "# from tensorflow.keras.layers import Dense\n",
    "# # Create a model.\n",
    "# model = Sequential()\n",
    "# model.add(Dense(32, activation='relu', input_dim=100))\n",
    "# model.add(Dense(10, activation='softmax'))\n",
    "# model.compile(optimizer='rmsprop',\n",
    "#             loss='categorical_crossentropy',\n",
    "#             metrics=['accuracy'])\n",
    "\n",
    "# # Number of training samples.\n",
    "# sample_count = 100*32\n",
    "\n",
    "# # Total epochs to train.\n",
    "# epochs = 30\n",
    "\n",
    "# # Number of warmup epochs.\n",
    "# warmup_epoch = 2\n",
    "\n",
    "# # Training batch size, set small value here for demonstration purpose.\n",
    "# batch_size = 32\n",
    "\n",
    "# # Base learning rate after warmup.\n",
    "# learning_rate_base = 0.001\n",
    "\n",
    "# total_steps = int(epochs * sample_count / batch_size)\n",
    "\n",
    "# # Compute the number of warmup batches.\n",
    "# warmup_steps = int(warmup_epoch * sample_count / batch_size)\n",
    "\n",
    "# # Generate dummy data.\n",
    "# data = np.random.random((sample_count, 100))\n",
    "# labels = np.random.randint(10, size=(sample_count, 1))\n",
    "\n",
    "# # Convert labels to categorical one-hot encoding.\n",
    "# one_hot_labels = keras.utils.to_categorical(labels, num_classes=10)\n",
    "\n",
    "# # Compute the number of warmup batches.\n",
    "# warmup_batches = warmup_epoch * sample_count / batch_size\n",
    "\n",
    "# # Create the Learning rate scheduler.\n",
    "# warm_up_lr = WarmUpCosineDecayScheduler(learning_rate_base=1e-4,\n",
    "#                                         total_steps=epochs*100,\n",
    "#                                         warmup_steps=1*100,\n",
    "#                                         warmup_learning_rate = 1e-4,\n",
    "#                                         mini_lr=1e-6,\n",
    "#                                         steps_per_epoch=100,\n",
    "#                                         global_step_init=0*100,\n",
    "#                                         verbose = 1)\n",
    "\n",
    "# # Train the model, iterating on the data in batches of 32 samples\n",
    "# model.fit(data, one_hot_labels, epochs=epochs, batch_size=batch_size,\n",
    "#         verbose=0, callbacks=[warm_up_lr])\n",
    "\n",
    "# import matplotlib.pyplot as plt\n",
    "# plt.plot(warm_up_lr.learning_rates)\n",
    "# plt.xlabel('Step', fontsize=20)\n",
    "# plt.ylabel('lr', fontsize=20)\n",
    "# plt.axis([0, total_steps, 0, learning_rate_base*1.1])\n",
    "# plt.xticks(np.arange(0, total_steps, 50))\n",
    "# plt.grid()\n",
    "# plt.title('Cosine decay with warmup', fontsize=20)\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "51f248f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0001000 0.0004000 0.0007000 0.0010000 0.0009966 0.0009865 0.0009698 0.0009468 0.0009177 0.0008830 \n",
      "0.0008431 0.0007986 0.0007500 0.0006980 0.0006434 0.0005868 0.0005291 0.0004709 0.0004132 0.0003566 \n",
      "0.0003020 0.0002500 0.0002014 0.0001569 0.0001170 0.0000823 0.0000532 0.0000302 0.0000135 0.0000034 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(3.38082113e-06)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# lr = []\n",
    "# for i in range(30):\n",
    "#     lr.append(cosine_decay_with_warmup(global_step=i,\n",
    "#                               learning_rate_base=1e-3,\n",
    "#                               total_steps=30*1,\n",
    "#                               warmup_learning_rate=1e-4,\n",
    "#                               warmup_steps=3*1,\n",
    "#                               hold_base_rate_steps=0,\n",
    "#                               steps_per_epoch=1))\n",
    "# for i in range(3):\n",
    "#     for j in range(10):\n",
    "#         print('%.7f'%lr[i*10+j], end=' ')\n",
    "#     print()\n",
    "# lr[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "70dadc3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.00e-04 4.00e-04 7.00e-04 1.00e-03 9.97e-04 9.87e-04 9.70e-04 9.47e-04\n",
      " 9.18e-04 8.83e-04 8.43e-04 7.99e-04 7.50e-04 6.98e-04 6.43e-04 5.87e-04\n",
      " 5.29e-04 4.71e-04 4.13e-04 3.57e-04 3.02e-04 2.50e-04 2.01e-04 1.57e-04\n",
      " 1.17e-04 8.20e-05 5.30e-05 3.00e-05 1.30e-05 3.00e-06]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f6c5757a430>]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 14928 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# print(np.round(lr,6))\n",
    "# import matplotlib.pyplot as plt\n",
    "# plt.plot(lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6de9b4af",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
