| from __future__ import annotations |
|
|
| import copy |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import trackio |
| from data import generate_bindings |
| from model import FastWeightProgrammer, GRUControl, parameter_count |
| from safetensors.torch import save_file |
| from torch import nn |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "fast-weight-time-machine" |
| DATA_DIR = PROJECT_DIR / "data" |
|
|
|
|
| def seed_everything(seed: int) -> None: |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.set_num_threads(1) |
|
|
|
|
| def make_loader( |
| dataset: tuple[np.ndarray, ...], |
| *, |
| shuffle: bool, |
| seed: int, |
| ) -> DataLoader: |
| keys, values, writes, targets = dataset |
| return DataLoader( |
| TensorDataset( |
| torch.from_numpy(keys), |
| torch.from_numpy(values), |
| torch.from_numpy(writes), |
| torch.from_numpy(targets), |
| ), |
| batch_size=256, |
| shuffle=shuffle, |
| generator=torch.Generator().manual_seed(seed), |
| ) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model: nn.Module, loader: DataLoader) -> dict: |
| model.eval() |
| correct = 0 |
| total = 0 |
| losses = [] |
| for keys, values, writes, targets in loader: |
| logits = model(keys, values, writes) |
| losses.append(float(nn.functional.cross_entropy(logits, targets))) |
| correct += int((logits.argmax(1) == targets).sum()) |
| total += len(targets) |
| return {"accuracy": correct / total, "cross_entropy": float(np.mean(losses))} |
|
|
|
|
| def train_variant( |
| name: str, |
| model: nn.Module, |
| train_loader: DataLoader, |
| validation_loader: DataLoader, |
| ) -> tuple[nn.Module, list[dict]]: |
| optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-5) |
| best = copy.deepcopy(model.state_dict()) |
| best_accuracy = 0.0 |
| stale = 0 |
| history = [] |
| for epoch in range(1, 61): |
| model.train() |
| losses = [] |
| for keys, values, writes, targets in train_loader: |
| logits = model(keys, values, writes) |
| loss = nn.functional.cross_entropy(logits, targets) |
| optimizer.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| optimizer.step() |
| losses.append(float(loss.detach())) |
| validation = evaluate(model, validation_loader) |
| record = { |
| "variant": name, |
| "epoch": epoch, |
| "training_loss": float(np.mean(losses)), |
| "validation_accuracy": validation["accuracy"], |
| } |
| history.append(record) |
| if epoch % 5 == 0: |
| trackio.log(record) |
| if validation["accuracy"] > best_accuracy + 1e-4: |
| best_accuracy = validation["accuracy"] |
| best = copy.deepcopy(model.state_dict()) |
| stale = 0 |
| else: |
| stale += 1 |
| if stale >= 10 and epoch >= 20: |
| break |
| model.load_state_dict(best) |
| return model, history |
|
|
|
|
| def main() -> None: |
| seed_everything(2043) |
| train_data = generate_bindings(16_000, pairs=4, distractors=12, seed=2043) |
| validation_data = generate_bindings( |
| 2_000, pairs=4, distractors=12, seed=3043 |
| ) |
| evaluation_data = { |
| "four_pairs_12_distractors": generate_bindings( |
| 4_000, 4, 12, seed=4043 |
| ), |
| "eight_pairs_12_distractors": generate_bindings( |
| 4_000, 8, 12, seed=5043 |
| ), |
| "four_pairs_64_distractors": generate_bindings( |
| 4_000, 4, 64, seed=6043 |
| ), |
| "eight_pairs_64_distractors": generate_bindings( |
| 4_000, 8, 64, seed=7043 |
| ), |
| } |
| train_loader = make_loader(train_data, shuffle=True, seed=2043) |
| validation_loader = make_loader( |
| validation_data, shuffle=False, seed=3043 |
| ) |
| variants = { |
| "fast_weight": FastWeightProgrammer(), |
| "gru": GRUControl(), |
| } |
| trackio.init( |
| project="fast-weight-time-machine", |
| name="temporary-variable-binding-v1", |
| config={ |
| "training_examples": len(train_data[0]), |
| "training_pairs": 4, |
| "training_distractors": 12, |
| "parameters": { |
| name: parameter_count(model) for name, model in variants.items() |
| }, |
| }, |
| ) |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| results = {} |
| histories = {} |
| for name, model in variants.items(): |
| trained, history = train_variant( |
| name, model, train_loader, validation_loader |
| ) |
| histories[name] = history |
| results[name] = { |
| "parameters": parameter_count(trained), |
| "training_epochs": len(history), |
| **{ |
| condition: evaluate( |
| trained, make_loader(dataset, shuffle=False, seed=8043) |
| ) |
| for condition, dataset in evaluation_data.items() |
| }, |
| } |
| save_file(trained.state_dict(), ARTIFACT_DIR / f"{name}.safetensors") |
| report = { |
| "benchmark": "Temporary variable binding", |
| "training_examples": len(train_data[0]), |
| "results": results, |
| "training_history": histories, |
| } |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed( |
| DATA_DIR / "binding_test.npz", |
| keys=evaluation_data["four_pairs_12_distractors"][0], |
| values=evaluation_data["four_pairs_12_distractors"][1], |
| writes=evaluation_data["four_pairs_12_distractors"][2], |
| targets=evaluation_data["four_pairs_12_distractors"][3], |
| ) |
| trackio.log( |
| { |
| "fast_weight_accuracy": results["fast_weight"][ |
| "four_pairs_12_distractors" |
| ]["accuracy"], |
| "gru_accuracy": results["gru"]["four_pairs_12_distractors"][ |
| "accuracy" |
| ], |
| "fast_weight_long_accuracy": results["fast_weight"][ |
| "eight_pairs_64_distractors" |
| ]["accuracy"], |
| "gru_long_accuracy": results["gru"][ |
| "eight_pairs_64_distractors" |
| ]["accuracy"], |
| } |
| ) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
