PyTorch Custom Loss Function
PyTorch's built-in loss functions cover most tasks. When they don't — because your task has a unique objective, class imbalance, or domain-specific requirements — you write your own. Any Python function that takes predictions and targets as tensors and returns a scalar tensor is a valid loss function in PyTorch.
Why Custom Loss Functions Are Needed
Built-in loss: CrossEntropyLoss
Works for: balanced multi-class classification
What it misses:
→ Class imbalance: 95% samples are class 0, 5% class 1
→ model learns to always predict class 0
→ Ordinal structure: predicting severity (1–5 stars)
→ predicting "5" when truth is "4" should cost less than predicting "1"
→ Task-specific metrics: optimizing directly for F1, IoU, SSIM
Rule: A Loss Function Must Return a Scalar
The loss must be a single number so that .backward() can compute gradients. If your computation produces a tensor, reduce it to a scalar using .mean() or .sum().
Example 1: Weighted Binary Cross-Entropy
When one class is rare, give the model a larger penalty for missing it by assigning it a higher weight.
import torch
import torch.nn.functional as F
def weighted_bce_loss(pred_logits, targets, pos_weight=5.0):
"""
pred_logits: raw model output (before sigmoid)
targets: binary labels (0 or 1)
pos_weight: how much more to penalize missing a positive sample
"""
bce = F.binary_cross_entropy_with_logits(
pred_logits, targets,
pos_weight=torch.tensor(pos_weight)
)
return bce
# Example
logits = torch.tensor([2.0, -1.0, 0.5, -2.0])
targets = torch.tensor([1.0, 1.0, 0.0, 0.0])
loss = weighted_bce_loss(logits, targets, pos_weight=5.0)
print(loss)Example 2: Focal Loss
Focal loss was designed for object detection with severe class imbalance. It reduces the loss contribution from easy examples (ones the model is already confident about) and focuses training on hard, misclassified examples.
import torch
import torch.nn.functional as F
def focal_loss(pred_logits, targets, alpha=0.25, gamma=2.0):
"""
alpha: weight for the positive class
gamma: focusing parameter (higher = more focus on hard examples)
"""
bce = F.binary_cross_entropy_with_logits(
pred_logits, targets, reduction='none')
probs = torch.sigmoid(pred_logits)
p_t = probs * targets + (1 - probs) * (1 - targets)
weight = alpha * targets + (1 - alpha) * (1 - targets)
focal = weight * ((1 - p_t) ** gamma) * bce
return focal.mean()
logits = torch.randn(16)
targets = torch.randint(0, 2, (16,)).float()
loss = focal_loss(logits, targets)
print(loss)Focal loss diagram: p_t (confidence) Standard BCE Focal weight 0.9 (easy) low × (1-0.9)² = 0.01 ← barely contributes 0.5 (medium) medium × (1-0.5)² = 0.25 0.1 (hard) high × (1-0.1)² = 0.81 ← drives most of learning
Example 3: Mean Absolute Error (L1 Loss)
def mae_loss(predictions, targets):
return torch.abs(predictions - targets).mean()
preds = torch.tensor([2.5, 3.0, 4.1])
targets = torch.tensor([3.0, 3.0, 4.0])
print(mae_loss(preds, targets)) # tensor(0.2000)Example 4: Dice Loss for Segmentation
Dice loss is popular in image segmentation tasks where most pixels belong to the background. It directly optimizes the overlap between predicted and true masks.
def dice_loss(pred, target, smooth=1e-6):
"""
pred: predicted probabilities after sigmoid, shape [B, H, W]
target: binary ground truth masks, shape [B, H, W]
"""
pred = pred.contiguous().view(-1)
target = target.contiguous().view(-1)
intersection = (pred * target).sum()
dice = (2.0 * intersection + smooth) / (pred.sum() + target.sum() + smooth)
return 1.0 - dice # loss = 1 - dice score
pred = torch.sigmoid(torch.randn(4, 64, 64))
target = torch.randint(0, 2, (4, 64, 64)).float()
print(dice_loss(pred, target))Dice score = 2 × |A ∩ B| / (|A| + |B|) Perfect overlap → Dice = 1 → Loss = 0 No overlap → Dice = 0 → Loss = 1
Using a Custom Loss in Training
model = MySegmentationModel()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
for images, masks in train_loader:
optimizer.zero_grad()
pred = torch.sigmoid(model(images)) # probabilities 0–1
loss = dice_loss(pred, masks)
loss.backward() # works because all operations use PyTorch tensors
optimizer.step()
print(f"Loss: {loss.item():.4f}")Custom Loss as an nn.Module
When your loss function has learnable parameters or you want to use it like a built-in module, implement it as a class inheriting from nn.Module.
class FocalLoss(nn.Module):
def __init__(self, alpha=0.25, gamma=2.0):
super().__init__()
self.alpha = alpha
self.gamma = gamma
def forward(self, pred_logits, targets):
bce = F.binary_cross_entropy_with_logits(
pred_logits, targets, reduction='none')
probs = torch.sigmoid(pred_logits)
p_t = probs * targets + (1 - probs) * (1 - targets)
wt = self.alpha * targets + (1 - self.alpha) * (1 - targets)
return (wt * (1 - p_t) ** self.gamma * bce).mean()
loss_fn = FocalLoss(alpha=0.25, gamma=2.0)
loss = loss_fn(logits, targets)
loss.backward() # gradients flow through normallyCombining Multiple Losses
Many tasks benefit from combining losses — for example, image reconstruction tasks often use pixel-level MSE plus a perceptual loss.
def combined_loss(pred, target, pred_features, target_features,
mse_weight=1.0, perceptual_weight=0.1):
mse = F.mse_loss(pred, target)
perceptual = F.mse_loss(pred_features, target_features)
return mse_weight * mse + perceptual_weight * perceptualSummary
Custom loss functions let you optimize for objectives that built-in losses don't cover. Write them as plain Python functions using PyTorch tensor operations — autograd handles gradient computation automatically. For class imbalance, use weighted BCE or Focal loss. For segmentation, use Dice loss. For regression with outliers, use Huber loss instead of MSE. Wrap custom losses in nn.Module when they have parameters or need to integrate cleanly with the rest of your model code. Always return a scalar (use .mean() or .sum()) so .backward() works correctly.