Chapter 11.3 - Calculating the Batch Loss

Overview

When you take a test in school, your teacher gives you a grade based on how many mistakes you made. For an AI, this grade is called the "Loss". The lower the loss, the smarter the AI!

---

🎯 Why we do it

Rationale

The AI starts off completely clueless. It just guesses random words. To make it smarter, we have to measure exactly how wrong it is. If the AI was supposed to say "Apple" but said "Car", we give it a big penalty (high loss). If it said "Orange", we give it a smaller penalty.

🛠️ How we do it

Methodology

We use a math formula called Cross-Entropy Loss. It compares the percentages the AI guessed with the actual 100% correct answer. It calculates the difference for the whole "batch" (group) of words it just tried to guess.

import torch
import torch.nn.functional as F

# The AI's guesses (percentages)
ai_guesses = torch.tensor([[0.1, 0.8, 0.1]]) # It guessed 80% for word #1

# The actual correct answer (Word #1 is correct!)
correct_answers = torch.tensor([1]) 

# Calculate the penalty!
loss = F.cross_entropy(ai_guesses, correct_answers)
print("The AI's penalty score is:", loss.item())

---

💻 Code Implementation

Here is the exact PyTorch implementation for the concepts discussed above:

import torch
import torch.nn.functional as F


def calc_loss_batch(input_batch, target_batch, model, device):
    input_batch = input_batch.to(device)
    target_batch = target_batch.to(device)

    logits = model(input_batch)

    return F.cross_entropy(
        logits.flatten(0, 1),
        target_batch.flatten(),
    )


def calc_loss_loader(data_loader, model, device, num_batches=None):
    if len(data_loader) == 0:
        return float("nan")

    if num_batches is None:
        num_batches = len(data_loader)
    else:
        num_batches = min(num_batches, len(data_loader))

    total_loss = 0.0

    with torch.no_grad():
        for i, (input_batch, target_batch) in enumerate(data_loader):
            if i >= num_batches:
                break

            total_loss += calc_loss_batch(
                input_batch,
                target_batch,
                model,
                device,
            ).item()

    return total_loss / num_batches