Implementation:Kornia Kornia Lovasz Softmax Loss

Knowledge Sources	Kornia
Domains	Vision, Loss_Functions
Last Updated	2026-02-09 15:00 GMT

Overview

Lovasz Softmax Loss computes a surrogate multi-class intersection-over-union (IoU) loss using the Lovasz extension, enabling differentiable optimization of the IoU metric for multi-class segmentation.

Description

The Lovasz Softmax loss extends the Lovasz hinge loss to the multi-class setting. It provides a differentiable surrogate for directly optimizing the per-class IoU metric, based on the Lovasz extension of submodular set functions (Berman et al., 2018).

The IoU for each class is:

$\text{IoU}(X,\text{class})=\frac{|X\cap Y|}{|X\cup Y|}$

The loss is:

$\text{loss}(x,\text{class})=1-\text{IoU}(x,\text{class})$

The implementation computes softmax probabilities, sorts per-class errors, computes Lovasz gradients, and aggregates the loss across classes. Optional per-class weights allow emphasizing specific classes.

Note: This loss function only supports multi-class labels (C > 1). For binary labels, use the Lovasz-Hinge loss.

Usage

Import this loss for multi-class semantic segmentation tasks where you want to directly optimize the mean IoU metric. It is especially beneficial when class imbalance is present or when evaluation is based on IoU/mIoU.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/losses/lovasz_softmax.py
• Lines: 1-186

Signature

def lovasz_softmax_loss(
    pred: Tensor,
    target: Tensor,
    weight: Optional[Tensor] = None,
) -> Tensor: ...

class LovaszSoftmaxLoss(nn.Module):
    def __init__(self, weight: Optional[Tensor] = None) -> None: ...
    def forward(self, pred: Tensor, target: Tensor) -> Tensor: ...

Import

from kornia.losses import LovaszSoftmaxLoss
from kornia.losses import lovasz_softmax_loss

I/O Contract

Inputs

Name	Type	Required	Description
pred	torch.Tensor	Yes	Logits tensor with shape (N, C, H, W) where C = number of classes (C > 1)
target	torch.Tensor	Yes	Labels tensor with shape (N, H, W) where each value is in [0, C-1]
weight	Optional[torch.Tensor]	No	Per-class weights with shape (num_classes,)

Outputs

Name	Type	Description
loss	torch.Tensor	Scalar loss value representing the mean Lovasz softmax loss

Usage Examples

import torch
from kornia.losses import LovaszSoftmaxLoss

# Multi-class segmentation
N = 5  # num_classes
pred = torch.randn(1, N, 3, 5, requires_grad=True)
target = torch.empty(1, 3, 5, dtype=torch.long).random_(N)

# Using the module API
criterion = LovaszSoftmaxLoss()
output = criterion(pred, target)
output.backward()

# With per-class weights
weights = torch.tensor([1.0, 2.0, 1.5, 1.0, 0.5])
criterion_weighted = LovaszSoftmaxLoss(weight=weights)
output = criterion_weighted(pred, target)

Related Pages

• Principle:Kornia_Kornia_Segmentation_Losses