PyTorch

uv add torch torchvision   # or: pip install torch ...
uv pip install torch --torch-backend=auto  # picks CUDA/CPU
uv run python -c "import torch; print(torch.__version__)"

[tool.uv.sources]
torch = [{ index = "pytorch-cu130" }]
 
[[tool.uv.index]]
name = "pytorch-cu130"      # or "pytorch-cpu" + /whl/cpu
url = "https://download.pytorch.org/whl/cu130"
explicit = true             # only torch comes from here

import numpy as np
import torch
 
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])  # copies data
z = torch.zeros(2, 3)                        # float32
g = torch.Generator().manual_seed(0)
r = torch.randn(4, 3, generator=g)           # seeded
i = torch.arange(0, 10, 2)                   # int64
o = torch.ones_like(a)                       # same shape
n = torch.from_numpy(np.ones(3))             # float64!
print(a.shape, a.dtype, a.device)  # (2, 2) float32 cpu
 
a @ a            # matrix multiply
a * a            # elementwise
a.sum(dim=0)     # reduce over rows -> shape (2,)
a.mean(dim=1, keepdim=True)  # shape (2, 1)
a[a > 2]         # boolean mask -> tensor([3., 4.])
a[0, 0].item()   # Python float 1.0

x = torch.randn(32, 10)       # batch of 32 rows
mu = x.mean(dim=0)            # (10,)
std = x.std(dim=0)            # (10,)
xn = (x - mu) / std           # (32, 10) op (10,) -> (32, 10)
col = torch.randn(32, 1)
col + mu                      # (32, 1) + (10,) -> (32, 10)

w = torch.tensor([2.0, -1.0], requires_grad=True)
x = torch.tensor([3.0, 4.0])
loss = ((w * x).sum() - 1.0) ** 2   # (6 - 4 - 1)^2 = 1
loss.backward()                     # fills w.grad
print(w.grad)                       # tensor([6., 8.])
 
with torch.no_grad():               # manual SGD step
    w -= 0.01 * w.grad
w.grad = None                       # grads accumulate!

import torch
from torch import Tensor, nn
 
 
class MLP(nn.Module):
    def __init__(
        self, d_in: int, d_hidden: int, n_classes: int
    ) -> None:
        super().__init__()          # required first
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(d_hidden, n_classes),  # logits
        )
 
    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)
 
 
model = MLP(20, 64, 3)
logits = model(torch.randn(8, 20))   # (8, 3)
n_params = sum(p.numel() for p in model.parameters())

opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
warmup = torch.optim.lr_scheduler.LinearLR(
    opt, start_factor=0.1, total_iters=5
)
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
    opt, T_max=45
)
sched = torch.optim.lr_scheduler.SequentialLR(
    opt, [warmup, cosine], milestones=[5]
)

from torch.utils.data import DataLoader
 
Batch = tuple[Tensor, Tensor]
 
 
def train_epoch(
    model: nn.Module,
    loader: DataLoader[Batch],
    loss_fn: nn.Module,
    opt: torch.optim.Optimizer,
    device: torch.device,
) -> float:
    model.train()                     # dropout/BN on
    total, seen = 0.0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad()               # 1. clear old grads
        loss = loss_fn(model(xb), yb) # 2. forward
        loss.backward()               # 3. backprop
        opt.step()                    # 4. update weights
        total += loss.item() * len(xb)
        seen += len(xb)
    return total / seen

@torch.inference_mode()               # no graph, faster
def evaluate(
    model: nn.Module,
    loader: DataLoader[Batch],
    loss_fn: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()                      # dropout/BN off
    loss_sum, correct, seen = 0.0, 0, 0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)
        loss_sum += loss_fn(logits, yb).item() * len(xb)
        correct += (logits.argmax(1) == yb).sum().item()
        seen += len(xb)
    return loss_sum / seen, correct / seen

from torch.utils.data import TensorDataset, random_split
 
g = torch.Generator().manual_seed(0)
X = torch.randn(1_000, 20, generator=g)
y = (X[:, 0] + X[:, 1] > 0).long() + (X[:, 2] > 1).long()
full = TensorDataset(X, y)            # rows of (x, y)
train_ds, val_ds = random_split(
    full, [0.8, 0.2], generator=g
)
 
train_loader = DataLoader(
    train_ds, batch_size=64, shuffle=True, generator=g
)
val_loader = DataLoader(val_ds, batch_size=256)
xb, yb = next(iter(train_loader))     # (64, 20), (64,)

def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():   # Apple GPU
        return torch.device("mps")
    return torch.device("cpu")
 
# device-agnostic API; None when no accelerator
acc = torch.accelerator.current_accelerator(
    check_available=True
)
device = acc if acc is not None else torch.device("cpu")
model = model.to(device)            # before the optimizer

use_fp16 = device.type == "cuda"
amp_dtype = torch.float16 if use_fp16 else torch.bfloat16
scaler = torch.amp.GradScaler(device.type, enabled=use_fp16)
loss_fn = nn.CrossEntropyLoss()
 
for xb, yb in train_loader:
    xb, yb = xb.to(device), yb.to(device)
    opt.zero_grad()
    with torch.autocast(device.type, dtype=amp_dtype):
        loss = loss_fn(model(xb), yb)   # forward in 16-bit
    scaler.scale(loss).backward()       # no-op if disabled
    scaler.step(opt)                    # skips inf/NaN steps
    scaler.update()

compiled = torch.compile(model)       # returns a wrapper
model.compile()                       # or compile in place
 
@torch.compile(fullgraph=True)        # error on graph breaks
def gelu_mlp(x: Tensor, w: Tensor) -> Tensor:
    return torch.nn.functional.gelu(x @ w)

torch.save(model.state_dict(), "model.pt")  # weights only
 
restored = MLP(20, 64, 3)                   # rebuild class
state = torch.load(
    "model.pt", map_location="cpu", weights_only=True
)
missing, unexpected = restored.load_state_dict(state)
restored.eval()

import os
import random
 
import numpy as np
 
 
def seed_everything(seed: int = 0) -> torch.Generator:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)          # CPU, CUDA and MPS
    return torch.Generator().manual_seed(seed)
 
 
def strict_determinism() -> None:
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    torch.use_deterministic_algorithms(True)  # raise if not
    torch.backends.cudnn.benchmark = False

torch.manual_seed(0)
device = pick_device()
model = MLP(20, 64, 3).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
loss_fn = nn.CrossEntropyLoss()
 
for epoch in range(20):
    tr = train_epoch(
        model, train_loader, loss_fn, opt, device
    )
    va, acc = evaluate(model, val_loader, loss_fn, device)
    if epoch % 5 == 4:
        print(f"{epoch:2} train {tr:.3f} val {va:.3f} "
              f"acc {acc:.2%}")

import pandas as pd
from torch.utils.data import Dataset
 
 
class CsvDataset(Dataset[tuple[Tensor, Tensor]]):
    def __init__(self, path: str, target: str) -> None:
        df = pd.read_csv(path)
        feats = df.drop(columns=[target])
        self.x = torch.tensor(
            feats.to_numpy(), dtype=torch.float32
        )
        self.y = torch.tensor(
            df[target].to_numpy(), dtype=torch.long
        )
 
    def __len__(self) -> int:
        return len(self.y)
 
    def __getitem__(self, i: int) -> tuple[Tensor, Tensor]:
        return self.x[i], self.y[i]
 
 
toy = {"a": [0.1, 0.9], "b": [1.0, 2.0], "label": [0, 1]}
pd.DataFrame(toy).to_csv("toy.csv", index=False)
ds = CsvDataset("toy.csv", target="label")
loader = DataLoader(ds, batch_size=2, shuffle=True)

import copy
 
 
class EarlyStopping:
    def __init__(self, patience: int = 5) -> None:
        self.patience, self.bad = patience, 0
        self.best = float("inf")
        self.state: dict[str, Tensor] | None = None
 
    def step(self, loss: float, model: nn.Module) -> bool:
        """Return True when training should stop."""
        if loss < self.best:
            self.best, self.bad = loss, 0
            self.state = copy.deepcopy(model.state_dict())
        else:
            self.bad += 1
        return self.bad >= self.patience
 
 
stopper = EarlyStopping(patience=3)
for epoch in range(100):
    train_epoch(model, train_loader, loss_fn, opt, device)
    val, _ = evaluate(model, val_loader, loss_fn, device)
    if stopper.step(val, model):
        break
assert stopper.state is not None
model.load_state_dict(stopper.state)  # best weights

from pathlib import Path
 
ckpt = Path("ckpt.pt")
sched = torch.optim.lr_scheduler.StepLR(opt, step_size=10)
 
 
def save(epoch: int) -> None:
    torch.save({
        "epoch": epoch,
        "model": model.state_dict(),
        "opt": opt.state_dict(),
        "sched": sched.state_dict(),
        "rng": torch.get_rng_state(),
    }, ckpt)
 
 
start = 0
if ckpt.exists():
    c = torch.load(ckpt, map_location=device)  # weights_only
    model.load_state_dict(c["model"])
    opt.load_state_dict(c["opt"])
    sched.load_state_dict(c["sched"])
    torch.set_rng_state(c["rng"].cpu())
    start = c["epoch"] + 1
for epoch in range(start, start + 3):
    train_epoch(model, train_loader, loss_fn, opt, device)
    sched.step()
    save(epoch)

from torchvision.models import ResNet18_Weights, resnet18
 
weights = ResNet18_Weights.DEFAULT      # downloads once
net = resnet18(weights=weights)
preprocess = weights.transforms()       # resize + normalize
 
for p in net.parameters():
    p.requires_grad_(False)             # freeze everything
net.fc = nn.Linear(net.fc.in_features, 5)  # new 5-class head
 
net = net.to(device)
head_opt = torch.optim.AdamW(
    (p for p in net.parameters() if p.requires_grad), lr=1e-3
)
imgs = torch.rand(4, 3, 256, 256)       # stand-in batch
logits = net(preprocess(imgs).to(device))  # (4, 5)
# later: unfreeze net.layer4 with a 10x lower LR group

accum = 4                            # effective batch x4
model.train()
opt.zero_grad()
for step, (xb, yb) in enumerate(train_loader, start=1):
    xb, yb = xb.to(device), yb.to(device)
    loss = loss_fn(model(xb), yb) / accum  # mean over steps
    loss.backward()                        # grads add up
    if step % accum == 0:
        nn.utils.clip_grad_norm_(      # returns the norm
            model.parameters(), max_norm=1.0
        )
        opt.step()
        opt.zero_grad()

@torch.inference_mode()
def predict(
    model: nn.Module, loader: DataLoader[Batch]
) -> tuple[Tensor, Tensor]:
    model.eval()
    preds, targets = [], []
    for xb, yb in loader:
        logits = model(xb.to(device))
        preds.append(logits.argmax(dim=1).cpu())
        targets.append(yb)
    return torch.cat(preds), torch.cat(targets)
 
 
pred, true = predict(model, val_loader)
acc = (pred == true).float().mean().item()
print(f"accuracy {acc:.2%}")