🎓 A1Learner
🔍

PyTorch Writing a Custom Loss

A custom loss is a plain function that scores predictions against targets. The function returns one number, and training drives that number down.

One mistake, scored two ways

The target is 10. A prediction of 7 misses low, and a prediction of 13 misses high.

import torch

def asymmetric_loss(pred, target):
    diff = pred - target
    squared = diff ** 2
    weight = 2.0 if diff < 0 else 1.0
    return weight * squared

print(asymmetric_loss(torch.tensor(7.0), torch.tensor(10.0)))
print(asymmetric_loss(torch.tensor(13.0), torch.tensor(10.0)))

A batch with torch.where

torch.where applies the double weight to every low prediction in one call.

import torch

pred = torch.tensor([8.0, 12.0])
target = torch.tensor([10.0, 10.0])
diff = pred - target
weights = torch.where(diff < 0, 2.0, 1.0)
print(weights)
print(weights * diff ** 2)

Gradients flow as usual

The loss is built from tensor operations. Calling backward fills the gradients.

import torch

model = torch.nn.Linear(1, 1)
pred = model(torch.tensor([[1.0]]))
target = torch.tensor([[3.0]])
diff = pred - target
loss = (torch.where(diff < 0, 2.0, 1.0) * diff ** 2).mean()
loss.backward()
print(model.weight.grad is not None)

Example

import torch

torch.manual_seed(0)

# Data follows the rule y = 2x + 1, with a little noise.
x = torch.linspace(0, 1, 40).reshape(-1, 1)
y = 2 * x + 1 + 0.1 * torch.randn(40, 1)


def asymmetric_loss(pred, target):
    diff = pred - target
    squared = diff ** 2
    weights = torch.where(diff < 0, 2.0, 1.0)
    return (weights * squared).mean()


model = torch.nn.Linear(1, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.3)

with torch.no_grad():
    before = model(x).mean().item()
print(f"Mean target: {y.mean().item():.4f}")
print(f"Mean prediction before training: {before:.4f}")

for step in range(300):
    loss = asymmetric_loss(model(x), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

with torch.no_grad():
    after = model(x).mean().item()
    final_loss = asymmetric_loss(model(x), y).item()
print(f"Mean prediction after training: {after:.4f}")
print(f"Final asymmetric loss: {final_loss:.4f}")

Key points

  • Loss function: a function that takes predictions and targets and returns one number
  • Use a custom loss when the built-in losses score mistakes the wrong way for the task
  • torch.where picks a value from one of two tensors based on a condition
  • An asymmetric loss scores one direction of mistake harder than the other
  • A custom loss plugs into the same training loop as any built-in loss
  • A custom loss is a scoring rule written as a function, and the loop treats it like any other loss

In this example, the data follows the rule y = 2x + 1 with a little noise. The mean target is 1.9919. The untrained model starts with a mean prediction of -0.6821, far below the targets. The asymmetric loss doubles the squared error whenever a prediction sits below its target. Training runs for 300 steps with that loss. The mean prediction climbs to 2.0188, just above the mean target. The double penalty on low misses pulls the predictions upward, exactly as the loss was written to do.

Note

The loss must return a tensor with gradients, so plain tensor maths is used instead of .item()

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A prediction is 4 and the target is 9. In the asymmetric loss from this lesson, what weight does this mistake get?