Search⌘ K
AI Features
Home
Courses
Deep Learning with JAX and Flax

Solution: Loss and Activation function

Explore how to implement multiple loss and activation functions with JAX for neural networks. Understand the coding of Softmax cross-entropy, Huber loss, ReLU, and ELU functions. Learn to apply these functions in convolutional layers and evaluate model accuracy to compare their effectiveness, gaining practical experience with neural network activation techniques.

We'll cover the following...

Solution 1: Implementation of loss functions

Here is the implementation of the following functions:

  1. Softmax cross-entropy
  2. Cosine similarity
  3. Huber loss
  4. CELU
  5. Softplus
Python 3.8
import jax.numpy as jnp
import optax
logits = jnp.array([0.50,0.60,0.70,0.30,0.25])
labels = jnp.array([0.20,0.30,0.10,0.20,0.2])
print("Softmax function: ", optax.softmax_cross_entropy(logits,labels))
print("Cosine similarity: ", optax.cosine_similarity(logits,labels,epsilon=0.5))
print("Huber Loss: ", optax.huber_loss(logits,labels))
print("L2 loss: ", optax.l2_loss(logits,labels))
print("CELU: ", jax.nn.celu(logits))
print("Softplus: ", jax.nn.softplus(logits))

Let’s review the code line-by-line:

  • Lines 1–3: We import all the necessary libraries. JAX or, more specifically, jax.nn, provides the necessary activation functions, while optax contains all the optimizer and loss functions.

  • Lines 4–5: We define the JAX arrays: logits and labels.

  • Lines 8–13: We call the ...