Evaluation

import torch
import torch.nn as nn

# ===== train vs eval =====
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Dropout(0.5), nn.Linear(16, 3))

model.train()    # default; enables Dropout + BatchNorm running stats updates
model.eval()     # disables Dropout (passes through); BatchNorm uses tracked stats

# These flags only affect modules with training-mode behaviour:
# Dropout, BatchNormNd, LayerNorm in some cases, GRU/LSTM with dropout, etc.

# ===== torch.no_grad() =====
# Disables autograd recording. Saves memory + speeds up.
with torch.no_grad():
    model.eval()
    out = model(X)

# Alternatively (Python 3.10+):
@torch.no_grad()
def predict(x):
    model.eval()
    return model(x)

# ===== Inference loop =====
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device).eval()

correct = total = 0
with torch.no_grad():
    for X, y in val_loader:
        X = X.to(device); y = y.to(device)
        out = model(X)
        pred = out.argmax(1)
        correct += (pred == y).sum().item()
        total += y.numel()
print(f'val_acc = {correct/total:.4f}')

# ===== Common bugs =====
# - Forgetting model.eval() -> Dropout still drops -> noisy / wrong inference
# - Forgetting torch.no_grad() -> memory blows up + slow + autograd graph kept
# - Calling loss.backward() inside no_grad() -> error
# - Forgetting to switch back to model.train() before next training epoch

# ===== A safe pattern =====
def evaluate(model, loader, device):
    was_training = model.training
    model.eval()
    total_loss = 0; n = 0
    with torch.no_grad():
        for X, y in loader:
            X = X.to(device); y = y.to(device)
            out = model(X)
            total_loss += loss_fn(out, y).item() * y.numel()
            n += y.numel()
    model.train(was_training)
    return total_loss / n

# ===== inference_mode (faster than no_grad on PyTorch 1.9+) =====
with torch.inference_mode():
    out = model(X)

# inference_mode is STRICTER: tensors created inside cannot be used in autograd later.
# Use for pure inference; use no_grad when results might re-enter graphs.

# ===== Loading a model for eval-only =====
state = torch.load('model.pt', map_location='cpu')
model.load_state_dict(state)
model.eval()

# ===== Mixed precision inference =====
from torch.amp import autocast
model.eval()
with torch.no_grad(), autocast('cuda', dtype=torch.float16):
    out = model(X.to('cuda'))

# ===== Calibration after training =====
# For models that benefit from temperature scaling:
class TemperatureScaler(nn.Module):
    def __init__(self):
        super().__init__()
        self.T = nn.Parameter(torch.ones(1) * 1.0)
    def forward(self, logits):
        return logits / self.T

# Fit T on a held-out set; use for calibrated probabilities at inference.

# ===== Patterns to internalise =====
# - model.eval() + torch.no_grad() for every inference path
# - inference_mode for performance-critical pure inference
# - autocast for fp16/bf16 inference on GPU
# - Switch back to model.train() before next epoch

# ===== Pitfalls =====
# - 'Eval works but training fails' -> often a forgotten model.train() after evaluate()
# - Dropout left enabled at inference -> non-deterministic predictions
# - BatchNorm with tiny batches at eval -> use running stats (default)
# - Comparing eval / train loss without remembering Dropout / BN differences

model.eval()
correct = total = 0
with torch.no_grad():
    for xb, yb in val_loader:
        out = model(xb.to(dev))
        correct += (out.argmax(1) == yb.to(dev)).sum().item()
        total += len(yb)
print('acc:', correct / total)