mindspore.nn.KLDivLoss

class mindspore.nn.KLDivLoss(reduction='mean')[source]

Computes the Kullback-Leibler divergence between the logits and the labels.

For tensors of the same shape \(x\) and \(target\), the updating formulas of KLDivLoss algorithm are as follows,

\[L(x, target) = target \cdot (\log target - x)\]

Then,

\[\begin{split}\ell(x, target) = \begin{cases} L(x, target), & \text{if reduction} = \text{'none';}\\ \operatorname{mean}(L(x, target)), & \text{if reduction} = \text{'mean';}\\ \operatorname{sum}(L(x, target)) / x.\operatorname{shape}[0], & \text{if reduction} = \text{'batchmean';}\\ \operatorname{sum}(L(x, target)), & \text{if reduction} = \text{'sum'.} \end{cases}\end{split}\]

where \(x\) represents logits, \(target\) represents labels, and \(\ell(x, target)\) represents output.

Note

  • Currently it does not support float64 input on Ascend.

  • The output aligns with the mathematical definition of Kullback-Leibler divergence only when reduction is set to 'batchmean'.

Parameters

reduction (str) –

Specifies the reduction to be applied to the output. Default: 'mean' .

  • On Ascend, the value of reduction must be one of 'batchmean' , 'none' or 'sum' .

  • On GPU, the value of reduction must be one of 'mean' , 'none' or 'sum' .

  • On CPU, the value of reduction must be one of 'mean' , 'batchmean' , 'none' or 'sum' .

Inputs:

  • logits (Tensor) - The input Tensor. The data type must be float16, float32 or float64.

  • labels (Tensor) - The label Tensor which has the same shape and data type as logits.

Outputs:

Tensor or Scalar, if reduction is 'none', then output is a tensor and has the same shape as logits. Otherwise, it is a scalar.

Raises

  • TypeError – If reduction is not a str.

  • TypeError – If neither logits nor labels is a Tensor.

  • TypeError – If dtype of logits or labels is not currently supported.

  • ValueError – If shape of logits is not the same as labels.

  • RuntimeError – If logits or labels is a scalar when reduction is 'batchmean'.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore as ms
>>> import mindspore.nn as nn
>>> import numpy as np
>>> logits = ms.Tensor(np.array([0.2, 0.7, 0.1]), ms.float32)
>>> labels = ms.Tensor(np.array([0., 1., 0.]), ms.float32)
>>> loss = nn.KLDivLoss(reduction='mean')
>>> output = loss(logits, labels)
>>> print(output)
-0.23333333