How to create a custom loss function in PyTorch

Loss functions turn prediction errors into the scalar value that a model minimizes. When under-prediction has a higher cost than over-prediction, a symmetric built-in criterion cannot represent that domain rule by itself.

A loss composed from PyTorch tensor operations can inherit from nn.Module and rely on autograd for its backward graph. A separate torch.autograd.Function is unnecessary unless the calculation uses operations that autograd cannot record or needs a custom gradient rule.

An asymmetric mean-squared error makes the behavior visible with two equal-sized errors: the under-prediction receives a multiplier while the over-prediction does not. Numerical gradient checking and a backward pass through nn.Linear then confirm that the criterion supplies finite model-parameter gradients.

1. Create asymmetric_mse_loss.py with the imports and loss-module constructor.

   asymmetric_mse_loss.py

   import torch
   from torch import nn
   from torch.autograd import gradcheck
    
    
   class AsymmetricMSELoss(nn.Module):
       def __init__(self, under_prediction_weight=2.0):
           super().__init__()
           if under_prediction_weight < 1:
               raise ValueError("under_prediction_weight must be at least 1")
           self.under_prediction_weight = under_prediction_weight

   The constructor keeps the penalty factor as module state and rejects values that would make under-prediction cheaper than ordinary mean-squared error.

2. Add the asymmetric mean-squared-error calculation inside AsymmetricMSELoss.

   asymmetric_mse_loss.py

       def forward(self, prediction, target):
           squared_error = (prediction - target).square()
           weights = torch.where(
               prediction < target,
               self.under_prediction_weight,
               1.0,
           )
           return (weights * squared_error).mean()

   torch.where() selects the multiplier element by element, and mean() reduces the weighted errors to the scalar loss required by backward().

3. Append paired under-prediction and over-prediction cases below the class.

   asymmetric_mse_loss.py

   if __name__ == "__main__":
       loss_fn = AsymmetricMSELoss(under_prediction_weight=3.0)
       under_prediction_loss = loss_fn(torch.tensor([0.0]), torch.tensor([1.0]))
       over_prediction_loss = loss_fn(torch.tensor([2.0]), torch.tensor([1.0]))

   Both predictions are one unit from the target. The configured multiplier should therefore make under_prediction_loss three times over_prediction_loss.

4. Append numerical and model-gradient checks after the paired cases.

   asymmetric_mse_loss.py

       check_prediction = torch.tensor(
           [0.25, 1.75],
           dtype=torch.float64,
           requires_grad=True,
       )
       check_target = torch.tensor([1.0, 1.0], dtype=torch.float64)
       gradient_check = gradcheck(loss_fn, (check_prediction, check_target))
    
       model = nn.Linear(1, 1, bias=False, dtype=torch.float64)
       features = torch.tensor([[1.0], [2.0]], dtype=torch.float64)
       targets = torch.tensor([[1.0], [2.0]], dtype=torch.float64)
       training_loss = loss_fn(model(features), targets)
       training_loss.backward()
       parameter_gradient = model.weight.grad
    
       assert under_prediction_loss > over_prediction_loss
       assert gradient_check
       assert parameter_gradient is not None
       assert torch.isfinite(parameter_gradient).all()
    
       print(f"under_prediction_loss={under_prediction_loss.item():.6f}")
       print(f"over_prediction_loss={over_prediction_loss.item():.6f}")
       print(f"gradient_check={gradient_check}")
       print(f"parameter_gradient_finite={bool(torch.isfinite(parameter_gradient).all())}")

   gradcheck() is designed for double-precision inputs and compares finite-difference gradients with the analytical gradients from autograd. The chosen values stay away from the branch boundary where prediction equals target.

5. Run the completed custom loss check.

   $ python3 asymmetric_mse_loss.py
   under_prediction_loss=3.000000
   over_prediction_loss=1.000000
   gradient_check=True
   parameter_gradient_finite=True

   The two loss values confirm the asymmetric penalty. The final two lines confirm that the loss passes numerical gradient checking and propagates a finite gradient to a model parameter.