Deep Learning with PyTorch

import torch

x = torch.tensor([2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()
y.backward()
print(x.grad)          # tensor([4., 6.])  (dy/dx = 2x)

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

# toy data
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()
train_dl = DataLoader(TensorDataset(X[:800], y[:800]), batch_size=64, shuffle=True)
test_dl = DataLoader(TensorDataset(X[800:], y[800:]), batch_size=128)

model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, 2))
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    model.train()
    for xb, yb in train_dl:
        opt.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        opt.step()

    model.eval()
    correct = 0
    with torch.no_grad():
        for xb, yb in test_dl:
            correct += (model(xb).argmax(1) == yb).sum().item()
    print(f"epoch {epoch}: test acc {correct / len(test_dl.dataset):.3f}")

class SmallCNN(nn.Module):
    def __init__(self, classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 28 -> 14
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 14 -> 7
        )
        self.head = nn.Sequential(nn.Flatten(), nn.Linear(32 * 7 * 7, 128), nn.ReLU(), nn.Linear(128, classes))

    def forward(self, x):
        return self.head(self.features(x))