{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Ann Optax Flax" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## ANN, Flax, Optax\n", "\n", "So far, you've seen how to define a model using jax, train the model, do backpropagation, and test it. Let's take it a bit further. In this notebook we'll be writing a simple ANN to classify penguin species using the [palmers penguin dataset](https://github.com/mcnakhaee/palmerpenguins).\n", "\n", "Since it's the first notebook on something NN, it'll be a good opportunity to introduce the Neural Net library in the Jax ecosystem, Flax. Also, like the linear regression example, we'll have to update the model parameters during backpropagation, for which there is also a library, Optax. In this notebook, I'll show you how to use both with a simple classifier.\n", "\n", "Let's load the data and process it to get started. (apparently the main focus of this notebook is to show you the ANN and not how you can load data, so, I'm going to short circuit the data related process a bit)\n", "\n", "### Data" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2026-05-07 19:28:54.684\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mutils.device\u001b[0m:\u001b[36msetup_backend\u001b[0m:\u001b[36m24\u001b[0m - \u001b[1mSetting up JAX to use the `mps` backend.\u001b[0m\n", "\u001b[32m2026-05-07 19:28:54.685\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[36mutils.device\u001b[0m:\u001b[36msetup_backend\u001b[0m:\u001b[36m26\u001b[0m - \u001b[32m\u001b[1mJAX Platform set to use the `mps` backend.\u001b[0m\n", "Platform 'mps' is experimental and not all JAX functionality may be correctly supported!\n", "\u001b[32m2026-05-07 19:28:54.816\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mutils.device\u001b[0m:\u001b[36mlog_device\u001b[0m:\u001b[36m14\u001b[0m - \u001b[1mXLA DEVICE: mps\u001b[0m\n" ] } ], "source": [ "# setting up device backend for apple silicon (mps)\n", "# if you have it that is.\n", "\n", "from utils import setup_backend, log_device\n", "\n", "setup_backend()\n", "log_device()" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import pandas as pd\n", "from palmerpenguins import load_penguins" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " species island bill_length_mm bill_depth_mm flipper_length_mm \\\n", "0 Adelie Torgersen 39.1 18.7 181.0 \n", "1 Adelie Torgersen 39.5 17.4 186.0 \n", "2 Adelie Torgersen 40.3 18.0 195.0 \n", "4 Adelie Torgersen 36.7 19.3 193.0 \n", "5 Adelie Torgersen 39.3 20.6 190.0 \n", "\n", " body_mass_g sex year \n", "0 3750.0 male 2007 \n", "1 3800.0 female 2007 \n", "2 3250.0 female 2007 \n", "4 3450.0 female 2007 \n", "5 3650.0 male 2007 \n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "33a02deb36934a4eb1ac1b054ef04970", "version_major": 2, "version_minor": 0 }, "text/plain": [ "0it [00:00, ?it/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "(333, 4)\n", "(333,)\n" ] } ], "source": [ "import jax\n", "import jax.numpy as jnp\n", "\n", "from tqdm.auto import tqdm\n", "from sklearn import preprocessing\n", "\n", "\n", "def load_data():\n", " penguins = load_penguins() # penguins is a dataframe\n", " penguins = penguins.dropna() # type: ignore\n", " \n", " # print the head of the dataframe to give some view\n", " print(penguins.head()) # type: ignore\n", " \n", " # collect the feature columns\n", " feature_columns = ['bill_length_mm', 'bill_depth_mm', 'flipper_length_mm', 'body_mass_g']\n", " # classification target\n", " target_column = \"species\"\n", " \n", " # features and targets\n", " features = penguins[feature_columns].values # type: ignore\n", " targets = penguins[target_column].values # type: ignore\n", " \n", " # but here's a catch\n", " # the targets are categorical, so we have two two options here\n", " # one hot encode them, or, assign a numeric value to them and keep a dictionary\n", " # with the target label to int mapping\n", " # the second approach is easier xD\n", " \n", " target_ids_dict = dict()\n", " unique_target_labels = set(targets)\n", " _id = 0\n", " \n", " for ul in unique_target_labels:\n", " target_ids_dict[ul] = _id\n", " _id += 1\n", " \n", " # convert target labels to integers using the same dict\n", " def convert_label_to_ids(targets, id_dict):\n", " converted_targets = np.zeros(shape=(len(targets, )), dtype=np.int32)\n", " for idx, target in tqdm(enumerate(targets)):\n", " converted_targets[idx] = id_dict[target]\n", "\n", " return converted_targets\n", " \n", " targets_converted = convert_label_to_ids(\n", " targets=targets, id_dict=target_ids_dict)\n", " \n", " assert features.shape[0] == targets_converted.shape[0]\n", " \n", " # the features from the dataset are not normalised and\n", " # this can cause probblems during training, such as \n", " # gradients getting stuck in a local minima\n", " # there's a lot of literature which talks about the \n", " # necessity of normalisation, this is a good starter\n", " # https://machinelearningmastery.com/how-to-improve-neural-network-stability-and-modeling-performance-with-data-scaling/\n", " \n", " features_norm = preprocessing.normalize(features, norm=\"l2\")\n", " \n", " \n", " return (features_norm, targets_converted)\n", " \n", " \n", "X, y = load_data()\n", "print(X.shape) # type: ignore\n", "print(y.shape) # type: ignore" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Okay splendid. Now to create the data split and also convert these numpy arrays to jax arrays. " ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Train Size : 233\n", "Test Size: 100\n" ] } ], "source": [ "from sklearn.model_selection import train_test_split\n", "\n", "features_train, features_test, targets_train, targets_test = train_test_split(\n", " X, y, test_size=0.3, random_state=42)\n", "\n", "\n", "features_train = jnp.array(features_train)\n", "features_test = jnp.array(features_test)\n", "targets_train = jnp.array(targets_train)\n", "targets_test = jnp.array(targets_test)\n", "\n", "\n", "print(f\"Train Size : {features_train.shape[0]}\")\n", "print(f\"Test Size: {features_test.shape[0]}\")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### ANN\n", "\n", "This is going to be a 2 layer ANN with Leaky ReLU activation, 3 target classes and cross entropy as the loss function." ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### PRNG" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "# old ritual of generating prngs\n", "key = jax.random.key(3047)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### ANN definition\n", "\n", "In Flax, `flax.linen` provides all the necessary bells and whistles to implement neural networks. If you're coming from pytorch, this is equivalent to `torch.nn`.\n", "\n", "Another intresting thing about Flax is that, despite Jax being functional, Flax does make use of classes (as PyTorch does!). Let's see the class definition. " ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "import flax.linen as nn\n", "\n", "\n", "class PenguinFinder(nn.Module):\n", " in_dim: int\n", " hidden_dim: int\n", " out_dim: int\n", " \n", " def setup(self): \n", " # layer definition goes here\n", " self.layer1 = nn.Dense(features=self.in_dim)\n", " self.hidden = nn.Dense(features=self.hidden_dim)\n", " self.layer2 = nn.Dense(features=self.out_dim)\n", "\n", " def __call__(self, x):\n", " out = self.layer1(x)\n", " # out = nn.relu(out)\n", " \n", " out = self.hidden(out)\n", " # out = nn.relu(out)\n", " \n", " out = self.layer2(out)\n", " out = nn.leaky_relu(out)\n", "\n", " # apply softmax to convert to \n", " # probability distribution\n", " logits = jax.nn.softmax(out)\n", "\n", " return logits" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "Let's take a look at the class definition here. First we declare the variables / hyperparameters to init the model as static python vars. Next we have a `setup` method which defines the layers in the NN. The `__call__` method is similar to what you see as `forward` in Pytorch, it applies all the layer operations on your data. \n", "\n", "\n", "The `setup` method comes in handy if you have to modify the default behavior of some linen class or want to define your own custom layers. If this is not the case (you want to use linen classes as they are), you can reduce the code by getting rid of setup and using `compact`. The class definition will then look like this\n", "\n", "```python\n", "class PenguinFinder(nn.Module):\n", " in_dim: int\n", " hidden_dim: int\n", " out_dim: int\n", "\n", " @nn.compact\n", " def __call__(self, x):\n", " out = nn.Dense(features=self.in_dim)(x)\n", " out = nn.Dense(features=self.hidden_dim)(out)\n", " out = nn.Dense(features=self.out_dim)(out)\n", " out = nn.leaky_relu(out)\n", " \n", " # apply softmax\n", " logits = jax.nn.softmax(out)\n", " \n", " return logits\n", "```" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "# init model\n", "model = PenguinFinder(in_dim=4, hidden_dim=4, out_dim=3)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "params from a linen module are returned as a Frozen dict, which is immutable. During training, when we won't be updating the params in place, rather Flax has some nifty trick with training state which we'll see pretty soon." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "# model params\n", "# requires a PRNG key and an array with the same shape of unbatched input\n", "params = model.init(key, features_train[0])\n" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Array([0.33331168, 0.333406 , 0.3332823 ], dtype=float32)" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# example forward pass\n", "model.apply(params, features_train[0])" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Loss function and grad\n", "\n", "The formal definition of cross entropy loss for a multiclass classification is this: \n", "$$\n", "ce = -\\sum_{c=1}^My_{t}\\log(p_{t})\n", "$$\n", "\n", "where, $t$ stands for the correct class\n", "\n", "$y_t$ is the correct label and $p_t$ is what a model predicted for $t$\n", "\n", "There is a nifty trick to it if you represent your classes with int ids as I have done above. The ids start from 0, so you can basically treat them as indexes. Using this, the cross entropy for an instance basically becomes\n", "\n", "$$\n", "ce = -ln(p_t)\n", "$$\n", "\n", "This trick works fine for single dimension multi class probabilities. I have never verified it outside course assignments or simple experiments. Then again, this notebook is just here to show you how jax works. In practice, it'll be buckwild to write everything from scratch. *Don't violate the DRY principle!*" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Array(1.0986773, dtype=float32)" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# loss function\n", "# for a single instance\n", "# will vmap for batches\n", "\n", "@jax.jit\n", "def cross_entropy(params, x, y):\n", " # forward pass .....\n", " logits = model.apply(params, x)\n", " \n", " return -jnp.log(logits[y]) # type: ignore\n", "\n", "\n", "cross_entropy(params, features_train[0], targets_train[0])\n" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Array(1.0986663, dtype=float32)" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "@jax.jit\n", "def calculate_loss(params, x_batched, y_batched):\n", " batch_loss = jax.vmap(cross_entropy, in_axes=(None, 0, 0))(params, x_batched, y_batched)\n", " return jnp.mean(batch_loss)\n", "\n", "\n", "calculate_loss(params, features_train, targets_train)\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Optimisation\n", "\n", "This is where we introduce Optax. For this notebook, I'm using the Adam optimiser. \n", "\n", "There are other optimizers in Optax which you can use: https://optax.readthedocs.io/en/latest/api.html\n", "\n", "Also, more on Adam: https://arxiv.org/abs/1412.6980\n" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "import optax\n", "\n", "optimiser = optax.adam(learning_rate=0.01)\n", "loss_grad_fn = jax.value_and_grad(calculate_loss)\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Flax Train state\n", "\n", "Okay so we've an optimiser but we need to tell it how it should be updating the parameters. We can do it using optax, go through the frozen dict (after making it immutable and froze it agian) but that's too much hassle and opens up more avenues for user error. Instead we'll be using `train_state` in Flax." ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [], "source": [ "from flax.training import train_state\n", "\n", "model_state = train_state.TrainState.create(\n", " apply_fn=model.apply, # the forward function\n", " params=params,\n", " tx=optimiser\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "During training, all we have to do is update this model state, and Flax will do the rest for us. (It'll be easier to understand the whole process if we consider params as states for models.)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "#### Train-ing\n", "Finally! Der Zug ist da! (Oh wait..........no, it's late again.)" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "# batched\n", "@jax.jit\n", "def train_step(state, xs, ys):\n", " loss_value, grads = loss_grad_fn(state.params, xs, ys)\n", " state = state.apply_gradients(grads=grads)\n", " return loss_value, state\n", " " ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [], "source": [ "from tqdm.auto import trange\n", "\n", "\n", "def train(state, xs, ys, epochs, log_every_n_step):\n", " losses = list() # keeps track of losses every n steps\n", " steps = list()\n", "\n", " step_count = 0\n", "\n", " for e in trange(epochs):\n", " loss, state = train_step(state, xs, ys)\n", " if step_count % log_every_n_step == 0:\n", " losses.append(loss)\n", " steps.append(step_count)\n", " \n", " print(f\"Epoch = {e + 1} / {epochs} \\t :: Train Loss = {loss}\")\n", " step_count += 1\n", "\n", " return state, losses, steps\n" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "65a5369304c3434abeb8d7ae6c3cc58e", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/2000 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "%matplotlib inline\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "\n", "\n", "sns.lineplot(x=np.array(steps), y=np.array(losses))\n", "plt.show()" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### Evaluation" ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "def test_step(state, x):\n", " logits = model.apply(state.params, x)\n", " out = jnp.argmax(logits, axis=-1)\n", " \n", " return out\n", "\n", "\n", "@jax.jit\n", "def test(state, xs):\n", " return jax.vmap(test_step, in_axes=(None, 0))(state, xs)\n" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "predictions = test(trained_model_state, features_test)\n" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " precision recall f1-score support\n", "\n", " 0 0.72 0.92 0.81 48\n", " 1 0.60 0.26 0.36 23\n", " 2 1.00 1.00 1.00 29\n", "\n", " accuracy 0.79 100\n", " macro avg 0.77 0.73 0.72 100\n", "weighted avg 0.77 0.79 0.76 100\n", "\n" ] } ], "source": [ "from sklearn.metrics import classification_report\n", "\n", "report = classification_report(y_pred=predictions, y_true=targets_test)\n", "print(report)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "### A fun little experiment\n", "Comment out all the `@jax.jit` annotations, restart the jupyter server and run all the cells again. Check how long it takes to train the network. \n", "\n", "tsk tsk ;) " ] } ], "metadata": { "kernelspec": { "display_name": ".venv", "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.13.13" } }, "nbformat": 4, "nbformat_minor": 4 }