INNER CODE UNIT · Python

cross_entropy_loss

nv-tlabs/ATISS · scene_synthesis/losses/__init__.py:16

def cross_entropy_loss(pred, target):
    """Cross entropy loss."""
    B, L, C = target.shape
    loss = torch.nn.functional.cross_entropy(
        pred.reshape(-1, C),
        target.reshape(-1, C).argmax(-1),
        reduction="none"
    ).reshape(B, L)

    return loss


def log_sum_exp(x):
    """Numerically stable log_sum_exp implementation that prevents
    overflow.
    """
    # TF ordering
    axis = len(x.size()) - 1

View source record →