![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*WpEtFvcW7jxALJ5dXz44BQ.png\")
<\/figure>JAX<\/a> is a Python<\/a> library offering high performance in machine learning with XLA<\/a> and Just In Time (JIT) compilation. Its API is similar to NumPy\u2019s<\/a> with a few differences. JAX ships with functionalities that aim to improve and increase speed in machine learning research. These functionalities include:<\/p>\n Flax is a neural network library for JAX. This article will cover how to track JAX and Flax models with Comet.<\/p>\n Let\u2019s get started.<\/p>\n<\/div>\n<\/div>\n<\/div>\n\n\n\n Start by installing Comet<\/a>.<\/p>\n Next, install JAX and Flax.<\/p>\n Import the libraries you\u2019ll use in this project.<\/p>\n In this project, you\u2019ll build a simple Convolutional Neural Network<\/a> using the MNIST dataset. Define the network parameters and log them to Comet.<\/p>\n To achieve that, create a Comet experiment.<\/p>\n Next, use this experiment to log the network metrics.<\/p>\n In Flax, networks are defined using the Linen package<\/a>. Define a simple CNN network<\/a> using the parameters defined above.<\/p>\n Flax networks are defined explicitly using setup vs compact<\/a><\/strong><\/p>\n Next, define the metrics used to compute loss and accuracy during training. In JAX, we compute the loss using the Optax library<\/a>.<\/p>\n The JAX and Flax don\u2019t ship with data loaders. Therefore, you have to use data loaders from TensorFlow<\/a> and PyTorch<\/a>.<\/p>\n In this case, we load the dataset using TensorFlow.<\/p>\n In Flax, we create a training state to store all the training information, such as parameters and the optimizer state. This is achieved using The next step is to define the model training functions. Let\u2019s start by defining a function that will train the network for one step.<\/p>\n The function:<\/p>\n Applying Next, define a training function that applies the above training step. The function:<\/p>\n The evaluation step checks the performance of the network on the testing data.<\/p>\n Next, define a function that applies the evaluation step to the testing data. The function also obtains the evaluation metrics from the device.<\/p>\n Want to try Comet for yourself? Sign up for a free account today!<\/a><\/p><\/blockquote>\n<\/div>\n<\/div>\n<\/div>\n\n\n\n You now have all the required building blocks for training the Flax CNN.<\/p>\n Start by downloading the data.<\/p>\n JAX requires pure functions. This means that they should not have any side effects. Therefore, even the random seed generation should be reproducible. These numbers are generated using JAX PRNG Design<\/a><\/strong><\/p>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n Create a training state using the function defined earlier.<\/p>\n The next step is to apply the training function for the desired number of epochs. At each epoch, we:<\/p>\n Since the metrics are saved in a list, you can plot the data and log the chart to Comet.<\/p>\n Don\u2019t forget to end the experiment once you are done<\/p>\n Click the link generated when you end the experiment to view the experiment on Comet\u2019s UI.<\/p>\n The Charts dashboard shows plots of the metrics you logged.<\/p>\n The hyperparameters dashboard shows the logged parameters.<\/p>\n The Graphics dashboard shows all the logged charts.<\/p>\n In this article, you have seen how to track Flax experiments with Comet. Apart from that, you have also seen how to:<\/p>\n JAX is a Python library offering high performance in machine learning with XLA and Just In Time (JIT) compilation. Its API is similar to NumPy\u2019s with a few differences. JAX ships with functionalities that aim to improve and increase speed in machine learning research. These functionalities include: Automatic differentiation Vectorization JIT compilation Flax is a […]<\/p>\n","protected":false},"author":63,"featured_media":0,"comment_status":"closed","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"customer_name":"","customer_description":"","customer_industry":"","customer_technologies":"","customer_logo":"","_jetpack_memberships_contains_paid_content":false,"footnotes":""},"categories":[9,7],"tags":[],"coauthors":[163],"class_list":["post-7611","post","type-post","status-publish","format-standard","hentry","category-product","category-tutorials"],"yoast_head":"\n\n
Getting started<\/h1>\n
pip install comet_ml<\/span><\/pre>\npip install -q jax jaxlib flax<\/span><\/pre>\nimport comet_ml\nimport jax\nimport jax.numpy as jnp # JAX NumPy\n\nfrom flax import linen as nn # The Linen API\nfrom flax.training import train_state # Useful dataclass to keep train state\n\nimport numpy as np # Ordinary NumPy\nimport optax # Optimizers\nimport tensorflow_datasets as tfds # TFDS for MNIST<\/pre>\n
Log parameters<\/h1>\n
experiment = comet_ml.Experiment(\n api_key=\"YOUR_API_KEY\",\n project_name=\"JAX_Flax_CNN\", log_code=True)<\/pre>\n
# these will all get logged\nparams = {\n \"features_1\": 32,\n \"kernel_size\": 3,\n \"window_shape\": 2,\n \"categories\": 10,\n \"features_2\": 64,\n \"features_3\":256,\n \"strides\": 2,\n \"cross_entropy_loss\": \"softmax_cross_entropy\",\n \"dataset\": \"MNIST\"\n}\n\nexperiment.log_parameters(params)<\/pre>\nFlax network definition<\/h1>\n
setup<\/a><\/code> or inline<\/code> using nn.compact<\/a><\/code>.<\/p>\nclass CNN(nn.Module):\n \"\"\"A simple CNN model.\"\"\"\n\n @nn.compact\n def __call__(self, x):\n x = nn.Conv(features=params['features_1'], kernel_size=(params['kernel_size'], params['kernel_size']))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(params['window_shape'], params['window_shape']), strides=(params['strides'], params['strides']))\n x = nn.Conv(features=params['features_2'], kernel_size=(params['kernel_size'], params['kernel_size']))(x)\n x = nn.relu(x)\n x = nn.avg_pool(x, window_shape=(params['window_shape'], params['window_shape']), strides=(params['strides'], params['strides']))\n x = x.reshape((x.shape[0], -1)) # flatten\n x = nn.Dense(features=params['features_3'])(x)\n x = nn.relu(x)\n x = nn.Dense(features=params['categories'])(x)\n return x<\/pre>\n<\/div>\n<\/div>\n<\/div>\n<\/div>\n
Compute metrics<\/h1>\n
def cross_entropy_loss(*, logits, labels):\n labels_onehot = jax.nn.one_hot(labels, num_classes=params['categories'])\n return optax.softmax_cross_entropy(logits=logits, labels=labels_onehot).mean()<\/pre>\n
compute_metrics<\/code> function will calculate and return the loss and accuracy.<\/p>\ndef compute_metrics(*, logits, labels):\n loss = cross_entropy_loss(logits=logits, labels=labels)\n accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)\n metrics = {\n 'loss': loss,\n 'accuracy': accuracy,\n }\n return metrics<\/pre>\nLoading data in JAX<\/h1>\n
def get_datasets():\n \"\"\"Load MNIST train and test datasets into memory.\"\"\"\n ds_builder = tfds.builder('mnist')\n ds_builder.download_and_prepare()\n train_ds = tfds.as_numpy(ds_builder.as_dataset(split='train', batch_size=-1))\n test_ds = tfds.as_numpy(ds_builder.as_dataset(split='test', batch_size=-1))\n train_ds['image'] = jnp.float32(train_ds['image']) \/ 255.\n test_ds['image'] = jnp.float32(test_ds['image']) \/ 255.\n return train_ds, test_ds<\/pre>\nCreate Flax training state<\/h1>\n
train_state<\/a><\/code> from Flax. In the training state function:<\/p>\n\n
init<\/code> method to obtain network parameters by passing a sample data point.<\/li>\ndef create_train_state(rng, learning_rate, momentum):\n \"\"\"Creates initial `TrainState`.\"\"\"\n cnn = CNN()\n params = cnn.init(rng, jnp.ones([1, 28, 28, 1]))['params']\n tx = optax.sgd(learning_rate, momentum)\n return train_state.TrainState.create(\n apply_fn=cnn.apply, params=params, tx=tx)<\/pre>\n
Network training<\/h1>\n
\n
jax.jit<\/a><\/code> makes the function run faster.<\/p>\n@jax.jit\ndef train_step(state, batch):\n \"\"\"Train for a single step.\"\"\"\n def loss_fn(params):\n logits = CNN().apply({'params': params}, batch['image'])\n loss = cross_entropy_loss(logits=logits, labels=batch['label'])\n return loss, logits\n grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n (_, logits), grads = grad_fn(state.params)\n state = state.apply_gradients(grads=grads)\n metrics = compute_metrics(logits=logits, labels=batch['label'])\n return state, metrics<\/pre>\n\n
jax.device_get<\/code><\/li>\ndef train_epoch(state, train_ds, batch_size, epoch, rng):\n \"\"\"Train for a single epoch.\"\"\"\n train_ds_size = len(train_ds['image'])\n steps_per_epoch = train_ds_size \/\/ batch_size\n\n perms = jax.random.permutation(rng, train_ds_size)\n perms = perms[:steps_per_epoch * batch_size] # skip incomplete batch\n perms = perms.reshape((steps_per_epoch, batch_size))\n batch_metrics = []\n for perm in perms:\n batch = {k: v[perm, ...] for k, v in train_ds.items()}\n state, metrics = train_step(state, batch)\n batch_metrics.append(metrics)\n\n # compute mean of metrics across each batch in epoch.\n batch_metrics_np = jax.device_get(batch_metrics)\n epoch_metrics_np = {\n k: np.mean([metrics[k] for metrics in batch_metrics_np])\n for k in batch_metrics_np[0]}\n\n\n\n return state, epoch_metrics_np['loss'],epoch_metrics_np['accuracy'] * 100<\/pre>\nNetwork evaluation<\/h1>\n
@jax.jit\ndef eval_step(params, batch):\n logits = CNN().apply({'params': params}, batch['image'])\n return compute_metrics(logits=logits, labels=batch['label'])<\/pre>\ndef eval_model(params, test_ds):\n metrics = eval_step(params, test_ds)\n metrics = jax.device_get(metrics)\n summary = jax.tree_util.tree_map(lambda x: x.item(), metrics)\n return summary['loss'], summary['accuracy']<\/pre>\n<\/div>\n<\/div>\n<\/div>\n\n\n\n
Train the Flax network<\/h1>\n
Download data<\/h2>\n
train_ds, test_ds = get_datasets()<\/pre>\n
Set random seed<\/h2>\n
jax.random.PRNGKey<\/a><\/code>.<\/p>\nrng = jax.random.PRNGKey(0)\nrng, init_rng = jax.random.split(rng)<\/pre>\n
Initialize train state<\/h2>\n
learning_rate = 0.1\nmomentum = 0.9\n\nstate = create_train_state(init_rng, learning_rate, momentum)\ndel init_rng # Must not be used anymore.<\/pre>\n
Log model metrics<\/h2>\n
\n
num_epochs = 10\nbatch_size = 32\n\ntraining_loss = []\ntraining_accuracy = []\ntesting_loss = []\ntesting_accuracy = []\n\n\nfor epoch in range(1, num_epochs + 1):\n # Use a separate PRNG key to permute image data during shuffling\n rng, input_rng = jax.random.split(rng)\n # Run an optimization step over a training batch\n state, train_loss, train_accuracy = train_epoch(state, train_ds, batch_size, epoch, input_rng)\n training_loss.append(train_loss)\n training_accuracy.append(train_accuracy)\n # Evaluate on the test set after each training epoch\n test_loss, test_accuracy = eval_model(state.params, test_ds)\n testing_loss.append(test_loss)\n testing_accuracy.append(test_accuracy)\n\n experiment.log_metric(\"train_loss\", train_loss, step=None, epoch=epoch, include_context=True)\n experiment.log_metric(\"train_accuracy\", train_accuracy, step=None, epoch=epoch, include_context=True)\n experiment.log_metric(\"test_loss\", test_loss, step=None, epoch=epoch, include_context=True)\n experiment.log_metric(\"test_accuracy\", test_accuracy, step=None, epoch=epoch, include_context=True)\n print(f'Epoch {epoch} train loss {train_loss} train accuracy {train_accuracy}. Test_loss {test_loss} test accuracy {test_accuracy}')<\/pre>\n![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*F4atM7B3prIstZgCCZ5umA.png\")
<\/figure>Log model charts<\/h2>\n
fig = plt.figure(figsize=(8, 6))\n\nplt.plot(training_loss, label=\"Training\")\nplt.plot(testing_loss, label=\"Test\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Accuracy\")\nplt.legend()\nplt.show()\n\nexperiment.log_figure(figure_name=\"Loss visualization\", figure=fig)<\/pre>\n
experiment.end()<\/pre>\n
![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*NoAUQgwRbDFkCkSHMnexVQ.png\")
<\/figure>View the experiment on Comet<\/h2>\n
![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*jvfw4y3uvBQzHcRwepQ68A.png\")
<\/figure>![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*oGZGOjmTcivDeI0dQHxcwQ.png\")
<\/figure>![\"\"](\"https:\/\/miro.medium.com\/v2\/resize:fit:700\/1*WUkRiaT1EdKCeGoT4KnCHA.png\")
<\/figure>Final thoughts<\/h1>\n
\n
Resources<\/h2>\n
\n