custom_loops.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd87782d",
   "metadata": {},
   "source": [
    "## These are cells that I am no longer using because I have adopted Keras a little better\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "351fb360",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CosAnnealWarmRestarts():\n",
    "    def __init__(self, T_0: float, T_mult: float):\n",
    "        '''\n",
    "        Cosine Annealing with Warm Restarts\n",
    "        Returns a new learning rate based on the call method\n",
    "        \n",
    "        Parameters:\n",
    "        `T_0` int\n",
    "            the number of iterations for the first restart to occur\n",
    "        `T_mult` int\n",
    "            the factor to increase T_i by after a restart, where T_i is the i^th restart.\n",
    "        '''\n",
    "        super(CosAnnealWarmRestarts, self).__init__()\n",
    "        \n",
    "        assert isinstance(T_0, float) and isinstance(T_mult, float)\n",
    "        assert T_0 > 0.0\n",
    "        self.mu_max = G.learning_rate #initial and max learning_rate\n",
    "        self.mu_min = G.min_learning_rate #minimum learning_rate\n",
    "        self.T_i = T_0\n",
    "        self.T_mult = T_mult\n",
    "        self.T_cur = 0.0\n",
    "        \n",
    "        self.learning_rate = G.learning_rate\n",
    "        \n",
    "    def step(self, increment:float, optimizer):\n",
    "        '''\n",
    "        Cosine Annealing with Warm Restarts\n",
    "        Returns a new learning rate based on the schedule described below\n",
    "        \n",
    "        Call after every batch\n",
    "\n",
    "        Parameters:\n",
    "        `increment` float\n",
    "            1 batch / total number of batches\n",
    "            !!!!! Not the current batch number, that would be a series summation\n",
    "            every epoch, a total of 1.0 will be added to self.T_cur\n",
    "        `optimizer`\n",
    "            the optimizer for the neural network\n",
    "        '''\n",
    "        try:\n",
    "            optimizer.learning_rate\n",
    "        except AttributeError:\n",
    "            print(\"Error: optimizer does not have a learning_rate parameter\")\n",
    "        \n",
    "        mu_i = self.mu_min + 0.5 * (\n",
    "                self.mu_max - self.mu_min) * (\n",
    "                    1 + tf.math.cos(np.pi * self.T_cur / self.T_i))\n",
    "        \n",
    "        self.T_cur += increment\n",
    "        \n",
    "        if np.isclose(self.T_cur, self.T_i):\n",
    "            self.T_i *= self.T_mult\n",
    "            self.T_cur = 0.0\n",
    "        \n",
    "        #update the learning_rate accordingly:\n",
    "        optimizer.learning_rate.assign(tf.cast(mu_i,tf.float32))\n",
    "        \n",
    "        self.learning_rate = mu_i\n",
    "        #this is just so that you can find the current learning rate from the scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304898cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.LogCosh()\n",
    "\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate = G.learning_rate,\n",
    "                                     beta_1 = 0.9,\n",
    "                                     beta_2 = 0.999\n",
    "                                    )\n",
    "scheduler = CosAnnealWarmRestarts(T_0 = 1.0, T_mult = 3.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3152bce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(model, x, y, training):\n",
    "    y_hat = model(x, training=training)\n",
    "    return y_hat, loss_object(y_true = y, y_pred = y_hat)\n",
    "\n",
    "def grad(model, inputs, targets):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_hat, attn_weights, loss = loss_fn(model, inputs, targets, training=True)\n",
    "    return y_hat, loss, tape.gradient(loss, model.trainable_variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3667a492",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_loop(dataloader, model, loss_fn, optimizer, scheduler, training = True):\n",
    "    size = len(dataloader)\n",
    "    perc_error = 0.0\n",
    "    epoch_loss_avg = tf.keras.metrics.Mean()\n",
    "    \n",
    "    for batch, (x,y) in enumerate(dataloader):\n",
    "        \n",
    "        predict, loss, grads = grad(model, x, y) # assert(loss.shape == [])\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_variables))\n",
    "        scheduler.step(1 / size, optimizer)\n",
    "        \n",
    "        perc_error += tf.math.reduce_mean(tf.abs(predict - y) / (y + 1e-2) * 100, [0,1])\n",
    "        epoch_loss_avg.update_state(loss)\n",
    "        if batch % (size // 15) == 0:\n",
    "            print(f\"Mean loss: {epoch_loss_avg.result():>7f}  [{batch:4d}/{size:4d}]\")\n",
    "\n",
    "    perc_error /= size\n",
    "    print(f\"Train Error: \\nAverage Accuracy: {100 - perc_error}%\")\n",
    "    return epoch_loss_avg.result(), (100. - perc_error), attn_weights\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn, training = False):\n",
    "    size = len(dataloader)\n",
    "    perc_error = 0.0\n",
    "    counter = 0\n",
    "    epoch_loss_avg = tf.keras.metrics.Mean()\n",
    "    \n",
    "    for x,y in dataloader:\n",
    "        predict, test_loss = loss_fn(model, x, y, training)\n",
    "        \n",
    "        if np.isnan(test_loss).any():\n",
    "            print(\"Test Loss had a np.nan value\")\n",
    "            break\n",
    "        \n",
    "        epoch_loss_avg.update_state(test_loss)\n",
    "        perc_error += tf.math.reduce_mean(tf.abs(predict - y) / (y + 1e-2) * 100, [0,1])\n",
    "\n",
    "        counter += 1\n",
    "        if counter % (size // 2) == 0:\n",
    "            print(f\"{counter} / {size} tested\")\n",
    "\n",
    "            \n",
    "    perc_error /= size\n",
    "    print(f\"Test Error: \\nAverage Accuracy: {100 - perc_error}%, Avg Loss: {epoch_loss_avg.result():>8f}\\n\")\n",
    "    return epoch_loss_avg.result(), 100. - perc_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd4aa9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = PP(plot_names = [\"Mean Log Loss\", \"% Accuracy\"],\n",
    "        line_names = [\"Train Loop\", \"Test Loop\"],\n",
    "        x_label = \"epochs\"\n",
    "       )\n",
    "## Note the y-axis gets cut off with numbers longer than three digits because the source code has a bug\n",
    "## I checked the github repo for lr-curve and the issue has been raised but not closed\n",
    "\n",
    "for epoch in range(1, 10):\n",
    "    print(f\"Epoch {epoch}/{G.epochs}\\n--------------------------------------\")\n",
    "    train_loss, train_acc, attn_weights = train_loop(train_dataloader, model, loss_fn, optimizer, scheduler)\n",
    "    test_loss, test_acc = test_loop(test_dataloader, model, loss_fn)\n",
    "    pp.update([[train_loss.numpy(), test_loss.numpy()], [train_acc, test_acc]])\n",
    "    \n",
    "#     if epoch % 15:\n",
    "#         model.save_weights(\"/content/drive/MyDrive/transformer_soc/decoder/model_weights.tf\", overwrite = True)\n",
    "    \n",
    "print(\"Completed!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2239e5a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights(\"/content/drive/MyDrive/transformer_soc/decoder/model_weights.tf\", overwrite = True)\n",
    "\n",
    "filehandler = open(\"/content/drive/MyDrive/transformer_soc/decoder/attn_weights\", \"wb\") #write in binary\n",
    "pickle.dump(attn_weights, filehandler)\n",
    "\n",
    "np.save(\n",
    "    \"/content/drive/MyDrive/transformer_soc/decoder/scheduler_state.npy\",\n",
    "    np.array([scheduler.learning_rate.numpy(), scheduler.T_cur, scheduler.T_i])\n",
    "       )\n",
    "\n",
    "print(f'''\n",
    "lr: {scheduler.learning_rate.numpy()}\n",
    "T_cur: {scheduler.T_cur}\n",
    "T_i: {scheduler.T_i}\n",
    "''')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}