Träna ett neuronnät i PyTorch

import torch, torch.nn as nn

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(5):
    model.train()
    for X, y in train_loader:            # batchar av (N, 784), (N,)
        opt.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        correct = sum((model(X).argmax(1) == y).sum().item() for X, y in val_loader)
    print(epoch, loss.item(), correct / len(val_loader.dataset))

Fem rader som är hela deep learning: nollställ, forward, loss, backward, steg.

Läs kurvorna:

CrossEntropyLoss tar råa utdata (logits) och gör softmax internt — lägg inte till softmax själv. Validera alltid på data modellen inte tränats på.

En epok = ett varv över all träningsdata. Batchstorlek styr hur många exempel per gradientsteg. Adam (torch.optim.Adam) är en populär optimerare som anpassar steglängden per parameter och ofta fungerar utan finjustering av lr (börja med 1e-3).