File size: 4,518 Bytes
ecc81b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """End-to-end: long-format rows -> lattice -> N-D model -> a trained model.
python examples/train_nd.py
torch-dimensions ships no trainer, on purpose — optimizers, schedules and loops
belong to you. What it does guarantee is that its models are ordinary
``nn.Module``s that train with ordinary PyTorch. This file is that guarantee
made runnable: every line below the model construction is plain torch, and
none of it is special.
The task is one-step-ahead forecasting over a (region x product) lattice where
one combination is never sold, so the lattice is genuinely sparse.
"""
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch_dimensions as td
REGIONS = ("north", "south", "east")
PRODUCTS = ("widget", "gadget")
NEVER_SOLD = ("east", "gadget")
def synthetic_rows(n_months=120, seed=0):
"""Long-format rows, the shape a database actually hands you."""
g = torch.Generator().manual_seed(seed)
coords, times, values = [], [], []
phase = {c: torch.rand(1, generator=g).item() * 6.28 for c in REGIONS + PRODUCTS}
for t in range(n_months):
for region in REGIONS:
for product in PRODUCTS:
if (region, product) == NEVER_SOLD:
continue # this combination does not exist
season = torch.sin(torch.tensor(t / 6.0 + phase[region])).item()
trend = t / n_months
noise = torch.randn(1, generator=g).item() * 0.05
coords.append((region, product))
times.append(t)
values.append([season + trend + noise])
return coords, times, values
def main() -> None:
torch.manual_seed(0)
# 1. Rows to a lattice. The grid, its axis vocabularies, and which cells
# exist are all inferred -- no hand-written coordinate mapping.
table = td.data.from_table(*synthetic_rows(), names=("region", "product"))
print(table)
print("lattice:", table.lattice)
# 2. Window the time axis, and split so no training window contains a
# timestep from after the cut.
windows = td.data.LatticeWindow(len(table), input_len=12, horizon=1)
train_w, test_w = windows.split(96)
print(f"windows: {len(train_w)} train, {len(test_w)} test")
source = td.data.TensorSource(table.series, table.lattice)
loader = DataLoader(
td.data.LatticeDataset(source, train_w),
batch_size=16,
shuffle=True,
collate_fn=td.data.collate_lattice,
)
# 3. An N-D model. Time stays causal; the categorical axes are swept both
# ways, since nothing about "region" is directional.
model = td.LSTM(
d_model=32,
n_layers=6,
lattice=table.lattice,
d_input=table.n_features,
bidirectional=("region", "product"),
)
head = nn.Linear(32, table.n_features)
print("plan:", model.plan)
print(
"params:", td.spec(model)["model"]["n_params"] + sum(p.numel() for p in head.parameters())
)
# 4. Your training loop. Nothing here comes from torch-dimensions.
opt = torch.optim.Adam([*model.parameters(), *head.parameters()], lr=3e-3)
mask = table.lattice.mask(torch.float32)
for epoch in range(15):
total, n = 0.0, 0
for batch in loader:
pred = head(model(batch.x))[:, -1:] # last step of the window
loss = ((pred - batch.y) * mask).pow(2).sum() / (mask.sum() * pred.shape[0])
opt.zero_grad()
loss.backward()
opt.step()
total += loss.item()
n += 1
if epoch % 3 == 0 or epoch == 14:
print(f"epoch {epoch:2d} train mse {total / n:.5f}")
# 5. Held-out evaluation on windows entirely after the split.
model.eval()
test_loader = DataLoader(
td.data.LatticeDataset(source, test_w),
batch_size=16,
collate_fn=td.data.collate_lattice,
)
with torch.no_grad():
errs = [
((head(model(b.x))[:, -1:] - b.y) * mask).pow(2).sum().item()
/ (mask.sum().item() * b.x.shape[0])
for b in test_loader
]
print(f"held-out mse {sum(errs) / len(errs):.5f}")
# Absent cells stay exactly zero all the way through.
with torch.no_grad():
out = head(model(next(iter(test_loader)).x)) * mask
absent = out.masked_select(~table.lattice.mask().expand_as(out))
print(f"absent-cell outputs: max |x| = {absent.abs().max().item():.1e}")
if __name__ == "__main__":
main()
|