Function ctc_loss

imperative/python/megengine/functional/loss.py:334–446

r"""The Connectionist Temporal Classification loss. Args: pred: The probabilities of the output, shape is (T, N, C) , where T=input length, N=batch size, and C=number of classes (including blank). pred_lengths: number of time steps for each sequence in ``pred``, sha

(
    pred: Tensor,
    pred_lengths: Tensor,
    label: Tensor,
    label_lengths: Tensor,
    blank: int = 0,
    reduction: str = "mean",
)

def ctc_loss(
    pred: Tensor,
    pred_lengths: Tensor,
    label: Tensor,
    label_lengths: Tensor,
    blank: int = 0,
    reduction: str = "mean",
) -> Tensor:
    r"""The Connectionist Temporal Classification loss.


    Args:
        pred: The probabilities of the output, shape is (T, N, C) ,
            where T=input length, N=batch size, and C=number of classes (including blank).
        pred_lengths: number of time steps for each sequence in ``pred``, shape is (N, )
        label: groundtruth labels, containing the indices of groundtruth
            symbols for each sequence at each output time step, and the blank
            symbol should not be included. shape is (N, S) or (sum(label_lengths)).
        label_lengths: number of time steps for each sequence in the groundtruth, shape is (N, )
        blank: the blank symbol number, default 0
        reduction: the reduction to apply to the output: 'none' | 'mean' | 'sum'. Default: 'mean'

    Returns:
        loss value.

    Examples:

        >>> pred = Tensor([[[0.0614, 0.9386],[0.8812, 0.1188]],[[0.699, 0.301 ],[0.2572, 0.7428]]])
        >>> pred_lengths = Tensor([2, 2])
        >>> label = Tensor([1, 1])
        >>> label_lengths = Tensor([1, 1])
        >>> F.nn.ctc_loss(pred, pred_lengths, label, label_lengths)
        Tensor(0.1504417, device=xpux:0)

    """
    T, N, C = pred.shape

    assert (
        pred_lengths.size == N
    ), "pred_lengths must be equal to batch_size {}, but got {}".format(
        N, pred_lengths.size
    )
    assert (
        label_lengths.size == N
    ), "label_lengths must be euqal to batch_size {}, but got {}".format(
        N, label_lengths.size
    )
    assert (
        blank >= 0 and blank < C
    ), "blank must be in label range [0, {}), but got {}".format(C, blank)
    assert (
        pred_lengths.min() > 0 and pred_lengths.max() <= T
    ), "pred_lengths must be in range ({}, {}], bug got min {}, max {}".format(
        0, T, pred_lengths.min(), pred_lengths.max()
    )

    if label.ndim == 1:  # concatenated label
        assert label_lengths.min() > 0, "label lengths muse be positive"