损失函数是机器学习训练的核心。它们用于评估模型在数据集上的表现：性能较差的模型会导致较高的损失值，而表现良好的模型则具有较低的损失值。因此，在构建机器学习模型时，选择合适的损失函数至关重要。本文将介绍 JAX 中提供的损失函数以及如何使用它们。

损失函数

机器学习模型通过比较预测值与真实值来学习，并据此调整模型权重，目标是找到使损失函数最小化的权重组合。损失函数也被称为代价函数（cost function）。损失函数的选择取决于具体问题类型，最常见的两类问题是分类和回归，它们各自需要不同的损失函数。

在 JAX 中创建自定义损失函数

使用 JAX 训练神经网络时，通常需要在训练阶段获取 logits（即模型输出的原始得分）。这些 logits 将用于计算损失值。接着需要评估损失函数及其梯度，梯度用于更新模型参数。在此基础上，还可以进一步计算模型的训练指标。

注意：logits 是未归一化的对数概率。

1def compute_loss(params, images, labels):
2    logits = CNN().apply({'params': params}, images)
3    loss = cross_entropy_loss(logits=logits, labels=labels)
4    return loss, logits
5    
6@jax.jit
7def train_step(state, images, labels):
8    """执行单步训练"""
9    (_, logits), grads = jax.value_and_grad(compute_loss, has_aux=True)(state.params, images, labels)
10    state = state.apply_gradients(grads=grads) 
11    metrics = compute_metrics(logits=logits, labels=labels) 
12    return state, metrics

我们可以使用 JAX 提供的 log_sigmoid 和 log_softmax 等函数来构建自定义损失函数，也可以完全从零开始编写自己的损失函数。

以下是一个计算 Sigmoid 二元交叉熵损失的示例：

1def custom_sigmoid_binary_cross_entropy(logits, labels):
2    log_p = jax.nn.log_sigmoid(logits)
3    log_not_p = jax.nn.log_sigmoid(-logits)
4    return -labels * log_p - (1. - labels) * log_not_p
5
6print(custom_sigmoid_binary_cross_entropy(0.5, 0.0))