KLDivLoss

class KLDivLoss(from_logits=True, axis=-1, weight=None, batch_axis=0, **kwargs)

The Kullback-Leibler divergence loss.

KL divergence measures the distance between contiguous distributions. It can be used to minimize information loss when approximating a distribution. If from_logits is True (default), loss is defined as:

(41)\[L = \sum_i {label}_i * \big[\log({label}_i) - {pred}_i\big]\]

If from_logits is False, loss is defined as:

(42)\[ \begin{align}\begin{aligned}\DeclareMathOperator{softmax}{softmax}\\prob = \softmax({pred})\\L = \sum_i {label}_i * \big[\log({label}_i) - \log({prob}_i)\big]\end{aligned}\end{align} \]

label and pred can have arbitrary shape as long as they have the same number of elements.

Parameters:

from_logits : bool, default is True

Whether the input is log probability (usually from log_softmax) instead of unnormalized numbers.

axis : int, default -1

The dimension along with to compute softmax. Only used when from_logits is False.

weight : float or None

Global scalar weight for loss.

batch_axis : int, default 0

The axis that represents mini-batch.

Inputs:

• pred: prediction tensor with arbitrary shape. If from_logits is True, pred should be log probabilities. Otherwise, it should be unnormalized predictions, i.e. from a dense layer.
• label: truth tensor with values in range (0, 1). Must have the same size as pred.
• sample_weight: element-wise weighting tensor. Must be broadcastable to the same shape as pred. For example, if pred has shape (64, 10) and you want to weigh each sample in the batch separately, sample_weight should have shape (64, 1).

Outputs:

• loss: loss tensor with shape (batch_size,). Dimenions other than batch_axis are averaged out.