Upload folder using huggingface_hub

ecc81b3 verified 3 months ago

4.52 kB

"""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()

	"""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()