MLP

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

# ===== Define =====
class MLP(nn.Module):
    def __init__(self, in_features, hidden, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.fc3 = nn.Linear(hidden, num_classes)
        self.drop = nn.Dropout(0.2)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.drop(x)
        x = F.relu(self.fc2(x))
        return self.fc3(x)            # raw logits; CE loss applies softmax

# ===== Data =====
X = torch.randn(1000, 4)
y = torch.randint(0, 3, (1000,))
Xv = torch.randn(200, 4)
yv = torch.randint(0, 3, (200,))

train_ds = TensorDataset(X, y)
val_ds = TensorDataset(Xv, yv)
train_loader = DataLoader(train_ds, batch_size=32, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=64)

# ===== Setup =====
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = MLP(in_features=4, hidden=64, num_classes=3).to(device)
optim = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# ===== Training loop =====
def train_one_epoch():
    model.train()
    total_loss = 0
    for Xb, yb in train_loader:
        Xb = Xb.to(device); yb = yb.to(device)
        optim.zero_grad()
        out = model(Xb)
        loss = loss_fn(out, yb)
        loss.backward()
        optim.step()
        total_loss += loss.item() * yb.size(0)
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def evaluate():
    model.eval()
    correct = total = 0
    total_loss = 0
    for Xb, yb in val_loader:
        Xb = Xb.to(device); yb = yb.to(device)
        out = model(Xb)
        total_loss += loss_fn(out, yb).item() * yb.size(0)
        pred = out.argmax(1)
        correct += (pred == yb).sum().item()
        total += yb.size(0)
    return total_loss / total, correct / total

# ===== Run =====
best_acc = 0
for epoch in range(20):
    train_loss = train_one_epoch()
    val_loss, val_acc = evaluate()
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save(model.state_dict(), 'best.pt')
    print(f'epoch {epoch}: train_loss={train_loss:.3f} val_loss={val_loss:.3f} val_acc={val_acc:.3f}')

# ===== Inference =====
model.load_state_dict(torch.load('best.pt'))
model.eval()
with torch.no_grad():
    pred = model(Xv.to(device)).argmax(1)

# ===== Early stopping pattern =====
patience = 5
best_loss = float('inf')
bad_epochs = 0
for epoch in range(100):
    train_loss = train_one_epoch()
    val_loss, _ = evaluate()
    if val_loss < best_loss:
        best_loss = val_loss
        bad_epochs = 0
        torch.save(model.state_dict(), 'best.pt')
    else:
        bad_epochs += 1
        if bad_epochs >= patience:
            print('early stopping')
            break

# ===== Patterns to internalise =====
# - .to(device) on model + every batch
# - zero_grad before backward; step after
# - model.train() / model.eval() toggles
# - Save the best state_dict; restore for inference

# ===== Pitfalls =====
# - Calling loss.item() inside the loop without detaching gradients (use no_grad in eval)
# - Forgetting .to(device) on a batch -> device mismatch error
# - Comparing eval / train losses without remembering dropout + bn differences
# - Mutating learning rate without a scheduler (use ReduceLROnPlateau or CosineAnnealing)

import torch.nn as nn
mlp = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.3),
    nn.Linear(256, 10),
)