diff --git a/benchmarks/adding_problem/adding_problem.py b/benchmarks/adding_problem/adding_problem.py
index 28570ea..624d6e9 100755
--- a/benchmarks/adding_problem/adding_problem.py
+++ b/benchmarks/adding_problem/adding_problem.py
@@ -1,267 +1,743 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import inspect
+import json
+import random
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Optional
+
import torch
import torch.nn as nn
import torch.optim as optim
-from torchrecurrent.benchmarks import adding_problem
from torch import Tensor
-import argparse
-import matplotlib.pyplot as plt
-import os
-import csv
+from torch.utils.data import DataLoader, TensorDataset
+
+import torchrecurrent
+from torchrecurrent.benchmarks import adding_problem
+
+try:
+ import matplotlib.pyplot as plt
+except ImportError: # pragma: no cover - benchmark convenience path
+ plt = None
+
+
+RECURRENT_LAYERS = [
+ torchrecurrent.AntisymmetricRNN,
+ torchrecurrent.ATR,
+ torchrecurrent.BR,
+ torchrecurrent.CFN,
+ torchrecurrent.coRNN,
+ torchrecurrent.DSGU,
+ torchrecurrent.FastGRNN,
+ torchrecurrent.FastRNN,
+ torchrecurrent.GatedAntisymmetricRNN,
+ torchrecurrent.IndRNN,
+ torchrecurrent.JANET,
+ torchrecurrent.LEM,
+ torchrecurrent.LightRU,
+ torchrecurrent.LiGRU,
+ torchrecurrent.MGU,
+ torchrecurrent.MiRU1,
+ torchrecurrent.MiRU2,
+ torchrecurrent.MultiplicativeLSTM,
+ torchrecurrent.MUT1,
+ torchrecurrent.MUT2,
+ torchrecurrent.MUT3,
+ torchrecurrent.NAS,
+ torchrecurrent.NBR,
+ torchrecurrent.OriginalLSTM,
+ torchrecurrent.PeepholeLSTM,
+ torchrecurrent.RAN,
+ torchrecurrent.ResLSTM,
+ torchrecurrent.SCRN,
+ torchrecurrent.SGRN,
+ torchrecurrent.SGU,
+ torchrecurrent.STAR,
+ torchrecurrent.tauGRU,
+ torchrecurrent.UGRNN,
+ torchrecurrent.UnICORNN,
+ torchrecurrent.WMCLSTM,
+]
+
+TORCH_BASELINES = [nn.RNN, nn.GRU, nn.LSTM]
+BASELINE_MSE = 1.0 / 6.0
+
+
+@dataclass
+class BenchmarkResult:
+ model: str
+ status: str
+ params: Optional[int] = None
+ best_test_mse: Optional[float] = None
+ best_epoch: Optional[int] = None
+ final_train_mse: Optional[float] = None
+ final_test_mse: Optional[float] = None
+ seconds: Optional[float] = None
+ error: Optional[str] = None
-class RecurrentModel(nn.Module):
- def __init__(self, cell, input_size: int, hidden_size: int, output_size: int, **kwargs):
+class RecurrentRegressor(nn.Module):
+ def __init__(
+ self,
+ recurrent_layer: type[nn.Module],
+ input_size: int,
+ hidden_size: int,
+ output_size: int,
+ **kwargs,
+ ):
super().__init__()
- self.rnn = cell(input_size, hidden_size, batch_first=True, **kwargs)
+ self.rnn = recurrent_layer(input_size, hidden_size, batch_first=True, **kwargs)
self.fc = nn.Linear(hidden_size, output_size)
- def forward(self, inp: Tensor):
+ def forward(self, inp: Tensor) -> Tensor:
output, _ = self.rnn(inp)
- last = output[:, -1, :]
- return self.fc(last)
+ return self.fc(output[:, -1, :])
+
+
+def model_name(layer: type[nn.Module]) -> str:
+ return layer.__name__
+
+
+def set_seed(seed: int) -> None:
+ random.seed(seed)
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+def resolve_device(args: argparse.Namespace) -> torch.device:
+ if args.device != "auto":
+ return torch.device(args.device)
+ if torch.cuda.is_available() and not args.no_cuda:
+ return torch.device("cuda")
+ if args.mps and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+ return torch.device("mps")
+ return torch.device("cpu")
+
+
+def layer_kwargs(layer: type[nn.Module], args: argparse.Namespace) -> dict[str, object]:
+ kwargs: dict[str, object] = {
+ "dropout": args.dropout,
+ "num_layers": args.num_layers,
+ }
+ signature = inspect.signature(layer)
+ if "nonlinearity" in signature.parameters:
+ kwargs["nonlinearity"] = "tanh"
+ return kwargs
+
+
+def make_loaders(
+ train_inputs: Tensor,
+ train_targets: Tensor,
+ test_inputs: Tensor,
+ test_targets: Tensor,
+ batch_size: int,
+ test_batch_size: int,
+ seed: int,
+) -> tuple[DataLoader, DataLoader]:
+ generator = torch.Generator()
+ generator.manual_seed(seed)
+ train_loader = DataLoader(
+ TensorDataset(train_inputs, train_targets),
+ batch_size=batch_size,
+ shuffle=True,
+ generator=generator,
+ )
+ test_loader = DataLoader(
+ TensorDataset(test_inputs, test_targets),
+ batch_size=test_batch_size,
+ shuffle=False,
+ )
+ return train_loader, test_loader
-def train(args, model, device, train_loader, optimizer, criterion, train_losses, epoch):
+def train_epoch(
+ model: nn.Module,
+ device: torch.device,
+ train_loader: DataLoader,
+ optimizer: optim.Optimizer,
+ criterion: nn.Module,
+ max_batches: Optional[int],
+) -> float:
model.train()
- total_loss = 0
+ total_loss = 0.0
+ n_batches = 0
for input_data, target_data in train_loader:
- input_data, target_data = input_data.to(device), target_data.to(device)
- optimizer.zero_grad()
- output = model(input_data)
- loss = criterion(output, target_data)
+ input_data = input_data.to(device)
+ target_data = target_data.to(device)
+
+ optimizer.zero_grad(set_to_none=True)
+ loss = criterion(model(input_data), target_data)
loss.backward()
optimizer.step()
- total_loss += loss.item()
- if args.dry_run:
- print("Dry run enabled, breaking after one batch.")
+ total_loss += float(loss.item())
+ n_batches += 1
+ if max_batches is not None and n_batches >= max_batches:
break
+ return total_loss / n_batches
- avg_loss = total_loss / len(train_loader)
- train_losses.append(avg_loss)
- print(f"Epoch {epoch}, Training Loss: {avg_loss:.6f}")
-
-def test(args, model, device, test_loader, criterion, test_losses, epoch):
+def evaluate(
+ model: nn.Module,
+ device: torch.device,
+ test_loader: DataLoader,
+ criterion: nn.Module,
+ max_batches: Optional[int],
+) -> float:
model.eval()
total_loss = 0.0
+ n_batches = 0
with torch.no_grad():
for input_data, target_data in test_loader:
- input_data, target_data = input_data.to(device), target_data.to(device)
- output = model(input_data)
- loss = criterion(output, target_data)
- total_loss += float(loss.item())
- if args.dry_run:
- print("Dry run enabled, breaking after one batch.")
+ input_data = input_data.to(device)
+ target_data = target_data.to(device)
+ total_loss += float(criterion(model(input_data), target_data).item())
+ n_batches += 1
+ if max_batches is not None and n_batches >= max_batches:
break
- avg_loss = total_loss / len(test_loader)
- test_losses.append(avg_loss)
- print(f"Epoch {epoch}, Test Loss: {avg_loss:.6f}")
+ return total_loss / n_batches
-def plot_learning_curves(train_losses, test_losses, out_png, title="", show=False):
- plt.figure()
- plt.plot(train_losses, label="Train MSE")
- plt.plot(test_losses, label="Test MSE")
- plt.axhline(1.0 / 6.0, linestyle="--", label="Baseline (1/6)")
- plt.xlabel("Epoch")
- plt.ylabel("MSE")
- plt.title(title or "Learning Curves")
- plt.legend()
- plt.grid(True, alpha=0.3)
- plt.savefig(out_png, bbox_inches="tight", dpi=150)
- if show:
- plt.show()
- plt.close()
+def run_one_model(
+ layer: type[nn.Module],
+ args: argparse.Namespace,
+ device: torch.device,
+ train_inputs: Tensor,
+ train_targets: Tensor,
+ test_inputs: Tensor,
+ test_targets: Tensor,
+ curves_dir: Path,
+) -> BenchmarkResult:
+ name = model_name(layer)
+ set_seed(args.seed + args.model_seed_offset)
+ train_loader, test_loader = make_loaders(
+ train_inputs,
+ train_targets,
+ test_inputs,
+ test_targets,
+ args.batch_size,
+ args.test_batch_size,
+ args.seed + args.shuffle_seed_offset,
+ )
+ model = RecurrentRegressor(
+ layer,
+ input_size=2,
+ hidden_size=args.hidden_size,
+ output_size=1,
+ **layer_kwargs(layer, args),
+ ).to(device)
+ params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ criterion = nn.MSELoss()
+ optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+ train_losses: list[float] = []
+ test_losses: list[float] = []
+ start = time.perf_counter()
+ for epoch in range(1, args.epochs + 1):
+ train_loss = train_epoch(
+ model,
+ device,
+ train_loader,
+ optimizer,
+ criterion,
+ args.max_batches,
+ )
+ test_loss = evaluate(
+ model,
+ device,
+ test_loader,
+ criterion,
+ args.max_test_batches,
+ )
+ train_losses.append(train_loss)
+ test_losses.append(test_loss)
+ if args.verbose:
+ print(
+ f"{name:>24s} epoch {epoch:03d}: train={train_loss:.6f} test={test_loss:.6f}"
+ )
-def save_losses_csv(train_losses, test_losses, out_csv):
- with open(out_csv, "w", newline="") as f:
- w = csv.writer(f)
- w.writerow(["epoch", "train_mse", "test_mse"])
- for i, (tr, te) in enumerate(zip(train_losses, test_losses), start=1):
- w.writerow([i, f"{tr:.8f}", f"{te:.8f}"])
+ seconds = time.perf_counter() - start
+ best_test = min(test_losses)
+ best_epoch = test_losses.index(best_test) + 1
+ write_curve_csv(curves_dir / f"{name}_losses.csv", train_losses, test_losses)
+ return BenchmarkResult(
+ model=name,
+ status="ok",
+ params=params,
+ best_test_mse=best_test,
+ best_epoch=best_epoch,
+ final_train_mse=train_losses[-1],
+ final_test_mse=test_losses[-1],
+ seconds=seconds,
+ )
-def main():
- parser = argparse.ArgumentParser(
- description="Addition problem benchmarks for recurrent layers"
- )
- parser.add_argument(
- "--batch-size",
- type=int,
- default=128,
- metavar="N",
- help="input batch size for training (default: 64)",
- )
- parser.add_argument(
- "--test-batch-size",
- type=int,
- default=1000,
- metavar="N",
- help="input batch size for testing (default: 1000)",
- )
- parser.add_argument(
- "--epochs",
- type=int,
- default=1000,
- metavar="N",
- help="number of epochs to train (default: 20)",
- )
- parser.add_argument(
- "--lr",
- type=float,
- default=0.001,
- metavar="LR",
- help="learning rate (default: 0.001)",
- )
- parser.add_argument(
- "--dropout",
- type=float,
- default=0.2,
- metavar="DO",
- help="dropout (default: 0.2)",
- )
- parser.add_argument(
- "--num_layers",
- type=int,
- default=2,
- metavar="NL",
- help="num_layers (default: 2)",
- )
- parser.add_argument(
- "--sequence-length",
- type=int,
- default=100,
- metavar="SL",
- help="length of the input sequences (default: 100)",
- )
- parser.add_argument(
- "--train-samples",
- type=int,
- default=5000,
- metavar="N",
- help="number of training samples (default: 5000)",
+def write_curve_csv(
+ path: Path, train_losses: list[float], test_losses: list[float]
+) -> None:
+ with path.open("w", newline="") as f:
+ writer = csv.writer(f)
+ writer.writerow(["epoch", "train_mse", "test_mse"])
+ for epoch, (train_loss, test_loss) in enumerate(
+ zip(train_losses, test_losses), start=1
+ ):
+ writer.writerow([epoch, f"{train_loss:.8f}", f"{test_loss:.8f}"])
+
+
+def write_summary_csv(path: Path, results: list[BenchmarkResult]) -> None:
+ fields = (
+ list(asdict(results[0]).keys())
+ if results
+ else list(BenchmarkResult.__annotations__)
)
- parser.add_argument(
- "--test-samples",
- type=int,
- default=1000,
- metavar="N",
- help="number of test samples (default: 1000)",
+ with path.open("w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=fields)
+ writer.writeheader()
+ for result in results:
+ writer.writerow(asdict(result))
+
+
+def read_curve(path: Path) -> tuple[list[int], list[float]]:
+ epochs: list[int] = []
+ losses: list[float] = []
+ with path.open(newline="") as f:
+ reader = csv.DictReader(f)
+ for row in reader:
+ epochs.append(int(row["epoch"]))
+ losses.append(float(row["test_mse"]))
+ return epochs, losses
+
+
+def svg_escape(text: str) -> str:
+ return (
+ text.replace("&", "&")
+ .replace("<", "<")
+ .replace(">", ">")
+ .replace('"', """)
)
- parser.add_argument(
- "--hidden-size",
- type=int,
- default=128,
- metavar="N",
- help="number of hidden units (default: 128)",
+
+
+def write_svg_bar(path: Path, results: list[BenchmarkResult]) -> None:
+ row_height = 22
+ margin_left = 190
+ margin_right = 40
+ margin_top = 42
+ width = 980
+ height = margin_top + row_height * len(results) + 48
+ values = [r.best_test_mse or 0.0 for r in results]
+ max_value = max(values + [BASELINE_MSE]) * 1.08
+ plot_width = width - margin_left - margin_right
+ baseline_x = margin_left + plot_width * BASELINE_MSE / max_value
+ lines = [
+ f'")
+ path.write_text("\n".join(lines))
+
+
+def write_svg_curves(
+ path: Path, results: list[BenchmarkResult], curves_dir: Path, top_k: int
+) -> None:
+ width = 1000
+ height = 650
+ margin_left = 70
+ margin_right = 180
+ margin_top = 48
+ margin_bottom = 55
+ plot_width = width - margin_left - margin_right
+ plot_height = height - margin_top - margin_bottom
+ curves = []
+ max_epoch = 1
+ max_loss = BASELINE_MSE
+ min_loss = 0.0
+ for result in results:
+ epochs, losses = read_curve(curves_dir / f"{result.model}_losses.csv")
+ curves.append((result.model, epochs, losses))
+ max_epoch = max(max_epoch, max(epochs))
+ max_loss = max(max_loss, max(losses))
+ max_loss *= 1.08
+
+ def point(epoch: int, loss: float) -> tuple[float, float]:
+ x = margin_left + plot_width * (epoch - 1) / max(1, max_epoch - 1)
+ y = margin_top + plot_height * (max_loss - loss) / max(1e-12, max_loss - min_loss)
+ return x, y
+
+ palette = [
+ "#4c78a8",
+ "#f58518",
+ "#54a24b",
+ "#e45756",
+ "#72b7b2",
+ "#b279a2",
+ "#ff9da6",
+ "#9d755d",
+ "#bab0ac",
+ "#8cd17d",
+ ]
+ top_names = {result.model for result in results[:top_k]}
+ top_order = [result.model for result in results[:top_k]]
+ lines = [
+ f'",
+ ]
)
- parser.add_argument(
- "--cuda", action="store_true", default=True, help="enables CUDA training"
+ path.write_text("\n".join(lines))
+
+
+def write_svg_scatter(path: Path, results: list[BenchmarkResult], top_k: int) -> None:
+ width = 900
+ height = 620
+ margin_left = 70
+ margin_right = 35
+ margin_top = 48
+ margin_bottom = 55
+ plot_width = width - margin_left - margin_right
+ plot_height = height - margin_top - margin_bottom
+ max_seconds = max((r.seconds or 0.0) for r in results) * 1.08
+ max_loss = max([r.best_test_mse or 0.0 for r in results] + [BASELINE_MSE]) * 1.08
+
+ def point(seconds: float, loss: float) -> tuple[float, float]:
+ x = margin_left + plot_width * seconds / max(1e-12, max_seconds)
+ y = margin_top + plot_height * (max_loss - loss) / max(1e-12, max_loss)
+ return x, y
+
+ baseline_y = point(0.0, BASELINE_MSE)[1]
+ lines = [
+ f'",
+ ]
)
- parser.add_argument(
- "--mps", action="store_true", default=False, help="enables MPS training"
+ path.write_text("\n".join(lines))
+
+
+def make_svg_plots(outdir: Path, results: list[BenchmarkResult], top_k: int) -> list[Path]:
+ ok_results = [r for r in results if r.status == "ok"]
+ ok_results.sort(
+ key=lambda r: r.best_test_mse if r.best_test_mse is not None else float("inf")
)
- parser.add_argument(
- "--dry-run",
- action="store_true",
- default=False,
- help="quickly check a single pass",
+ if not ok_results:
+ return []
+ curves_dir = outdir / "curves"
+ paths = [
+ outdir / "best_test_mse.svg",
+ outdir / "test_curves_all_models.svg",
+ outdir / "best_mse_vs_runtime.svg",
+ ]
+ write_svg_bar(paths[0], ok_results)
+ write_svg_curves(paths[1], ok_results, curves_dir, top_k)
+ write_svg_scatter(paths[2], ok_results, top_k)
+ return paths
+
+
+def make_plots(outdir: Path, results: list[BenchmarkResult], top_k: int) -> list[Path]:
+ if plt is None:
+ return make_svg_plots(outdir, results, top_k)
+
+ ok_results = [r for r in results if r.status == "ok"]
+ ok_results.sort(
+ key=lambda r: r.best_test_mse if r.best_test_mse is not None else float("inf")
)
- parser.add_argument(
- "--outdir", type=str, default="runs/adding_single", help="Where to save plots/logs."
+ if not ok_results:
+ return []
+
+ plot_paths: list[Path] = []
+ curves_dir = outdir / "curves"
+
+ bar_path = outdir / "best_test_mse.png"
+ fig_height = max(8.0, 0.28 * len(ok_results))
+ plt.figure(figsize=(11, fig_height))
+ labels = [r.model for r in ok_results]
+ values = [r.best_test_mse for r in ok_results]
+ plt.barh(labels, values, color="#4c78a8")
+ plt.axvline(BASELINE_MSE, color="#d62728", linestyle="--", label="baseline 1/6")
+ plt.gca().invert_yaxis()
+ plt.xlabel("Best test MSE")
+ plt.title("Adding Problem: best test MSE by recurrent layer")
+ plt.legend()
+ plt.grid(axis="x", alpha=0.25)
+ plt.tight_layout()
+ plt.savefig(bar_path, dpi=180)
+ plt.close()
+ plot_paths.append(bar_path)
+
+ curves_path = outdir / "test_curves_all_models.png"
+ top_names = {r.model for r in ok_results[:top_k]}
+ plt.figure(figsize=(13, 8))
+ for result in ok_results:
+ epochs, losses = read_curve(curves_dir / f"{result.model}_losses.csv")
+ if result.model in top_names:
+ plt.plot(epochs, losses, linewidth=2.2, label=result.model)
+ else:
+ plt.plot(epochs, losses, color="#9e9e9e", linewidth=0.8, alpha=0.45)
+ plt.axhline(BASELINE_MSE, color="#d62728", linestyle="--", label="baseline 1/6")
+ plt.xlabel("Epoch")
+ plt.ylabel("Test MSE")
+ plt.title(f"Adding Problem: test curves, top {min(top_k, len(ok_results))} highlighted")
+ plt.legend(ncol=2, fontsize=8)
+ plt.grid(alpha=0.25)
+ plt.tight_layout()
+ plt.savefig(curves_path, dpi=180)
+ plt.close()
+ plot_paths.append(curves_path)
+
+ scatter_path = outdir / "best_mse_vs_runtime.png"
+ plt.figure(figsize=(10, 7))
+ xs = [r.seconds for r in ok_results]
+ ys = [r.best_test_mse for r in ok_results]
+ sizes = [max(20, min(300, (r.params or 0) / 200)) for r in ok_results]
+ plt.scatter(xs, ys, s=sizes, color="#59a14f", alpha=0.75, edgecolors="#1f1f1f")
+ for result in ok_results[:top_k]:
+ plt.annotate(result.model, (result.seconds, result.best_test_mse), fontsize=8)
+ plt.axhline(BASELINE_MSE, color="#d62728", linestyle="--", label="baseline 1/6")
+ plt.xlabel("Training seconds")
+ plt.ylabel("Best test MSE")
+ plt.title("Adding Problem: quality/runtime tradeoff")
+ plt.legend()
+ plt.grid(alpha=0.25)
+ plt.tight_layout()
+ plt.savefig(scatter_path, dpi=180)
+ plt.close()
+ plot_paths.append(scatter_path)
+
+ return plot_paths
+
+
+def select_layers(args: argparse.Namespace) -> list[type[nn.Module]]:
+ layers = list(RECURRENT_LAYERS)
+ if args.include_torch_baselines:
+ layers.extend(TORCH_BASELINES)
+ if args.models:
+ requested = [name.strip() for name in args.models.split(",") if name.strip()]
+ by_name = {model_name(layer): layer for layer in layers}
+ unknown = sorted(set(requested).difference(by_name))
+ if unknown:
+ raise ValueError(f"Unknown model(s): {', '.join(unknown)}")
+ layers = [by_name[name] for name in requested]
+ if args.skip_models:
+ skipped = {name.strip() for name in args.skip_models.split(",") if name.strip()}
+ layers = [layer for layer in layers if model_name(layer) not in skipped]
+ return layers
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Unified adding-problem benchmark for torchrecurrent layers."
)
+ parser.add_argument("--batch-size", type=int, default=128)
+ parser.add_argument("--test-batch-size", type=int, default=512)
+ parser.add_argument("--epochs", type=int, default=100)
+ parser.add_argument("--lr", type=float, default=1e-3)
+ parser.add_argument("--weight-decay", type=float, default=0.0)
+ parser.add_argument("--dropout", type=float, default=0.0)
+ parser.add_argument("--num-layers", type=int, default=1)
+ parser.add_argument("--sequence-length", type=int, default=100)
+ parser.add_argument("--train-samples", type=int, default=5000)
+ parser.add_argument("--test-samples", type=int, default=1000)
+ parser.add_argument("--hidden-size", type=int, default=128)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--model-seed-offset", type=int, default=10_000)
+ parser.add_argument("--shuffle-seed-offset", type=int, default=20_000)
+ parser.add_argument("--device", default="auto", help="auto, cpu, cuda, cuda:0, or mps")
+ parser.add_argument("--no-cuda", action="store_true")
+ parser.add_argument("--mps", action="store_true", help="Allow MPS when device=auto.")
+ parser.add_argument("--models", default="", help="Comma-separated model names to run.")
parser.add_argument(
- "--show",
- action="store_true",
- default=False,
- help="Call plt.show() after saving the plot.",
+ "--skip-models", default="", help="Comma-separated model names to skip."
)
+ parser.add_argument("--include-torch-baselines", action="store_true")
+ parser.add_argument("--continue-on-error", action="store_true", default=True)
+ parser.add_argument("--fail-fast", dest="continue_on_error", action="store_false")
+ parser.add_argument("--max-batches", type=int, default=None)
+ parser.add_argument("--max-test-batches", type=int, default=None)
+ parser.add_argument("--top-k", type=int, default=10)
+ parser.add_argument("--verbose", action="store_true")
parser.add_argument(
- "--save-csv", action="store_true", default=False, help="Also save losses as CSV."
+ "--outdir", type=Path, default=Path("/tmp/torchrecurrent_adding_problem")
)
- args = parser.parse_args()
+ return parser.parse_args()
- if args.cuda and torch.cuda.is_available():
- device = torch.device("cuda")
- elif args.mps and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
- device = torch.device("mps")
- else:
- device = torch.device("cpu")
- torch.manual_seed(args.seed)
- if device.type == "cuda":
- torch.cuda.manual_seed(args.seed)
+def main() -> None:
+ args = parse_args()
+ outdir = args.outdir.expanduser().resolve()
+ curves_dir = outdir / "curves"
+ outdir.mkdir(parents=True, exist_ok=True)
+ curves_dir.mkdir(parents=True, exist_ok=True)
- print(f"Using device: {device}")
-
- train_loader = adding_problem(
+ device = resolve_device(args)
+ layers = select_layers(args)
+ set_seed(args.seed)
+ train_inputs, train_targets = adding_problem(
sequence_length=args.sequence_length,
n_samples=args.train_samples,
- batch_size=args.batch_size,
- shuffle=True,
+ return_dataloader=False,
)
-
- test_loader = adding_problem(
+ test_inputs, test_targets = adding_problem(
sequence_length=args.sequence_length,
n_samples=args.test_samples,
- batch_size=args.test_batch_size,
- shuffle=False,
+ return_dataloader=False,
)
- input_size = 2
- output_size = 1
- model = RecurrentModel(
- nn.GRU,
- input_size,
- args.hidden_size,
- output_size,
- dropout=args.dropout,
- num_layers=args.num_layers,
- ).to(device)
- # model = torch.jit.script(model)
- criterion = nn.MSELoss()
- optimizer = optim.Adam(model.parameters(), lr=args.lr) # Adam
-
- train_losses = []
- test_losses = []
+ config = vars(args).copy()
+ config["outdir"] = str(outdir)
+ config["device"] = str(device)
+ config["models"] = [model_name(layer) for layer in layers]
+ with (outdir / "config.json").open("w") as f:
+ json.dump(config, f, indent=2, sort_keys=True)
- # train and validate
- for epoch in range(1, args.epochs + 1):
- train(args, model, device, train_loader, optimizer, criterion, train_losses, epoch)
- test(args, model, device, test_loader, criterion, test_losses, epoch)
+ print(f"Output directory: {outdir}")
+ print(f"Device: {device}")
+ print(f"Models: {len(layers)}")
- if args.dry_run:
- print("Dry run enabled, stopping training.")
- break
+ results: list[BenchmarkResult] = []
+ for idx, layer in enumerate(layers, start=1):
+ name = model_name(layer)
+ print(f"[{idx:02d}/{len(layers):02d}] {name}")
+ try:
+ result = run_one_model(
+ layer,
+ args,
+ device,
+ train_inputs,
+ train_targets,
+ test_inputs,
+ test_targets,
+ curves_dir,
+ )
+ print(
+ f" best_test_mse={result.best_test_mse:.6f} "
+ f"epoch={result.best_epoch} seconds={result.seconds:.1f}"
+ )
+ except Exception as exc:
+ if not args.continue_on_error:
+ raise
+ result = BenchmarkResult(model=name, status="error", error=repr(exc))
+ print(f" error: {result.error}")
+ results.append(result)
+ write_summary_csv(outdir / "summary.csv", results)
+ with (outdir / "summary.json").open("w") as f:
+ json.dump([asdict(result) for result in results], f, indent=2)
- # ---- visualize & save ----
- os.makedirs(args.outdir, exist_ok=True)
- model_name = type(model.rnn).__name__ if hasattr(model, "rnn") else type(model).__name__
- tag = f"{model_name}_T{args.sequence_length}_H{args.hidden_size}"
- png_path = os.path.join(args.outdir, f"{tag}_learning_curves.png")
- plot_learning_curves(
- train_losses,
- test_losses,
- png_path,
- title=f"{model_name} on Adding Problem (T={args.sequence_length})",
- show=args.show,
+ plot_paths = make_plots(outdir, results, args.top_k)
+ ok_results = [result for result in results if result.status == "ok"]
+ ok_results.sort(
+ key=lambda r: r.best_test_mse if r.best_test_mse is not None else float("inf")
)
- print(f"Saved plot: {png_path}")
-
- if args.save_csv:
- csv_path = os.path.join(args.outdir, f"{tag}_losses.csv")
- save_losses_csv(train_losses, test_losses, csv_path)
- print(f"Saved CSV: {csv_path}")
-
- best_test = min(test_losses)
- best_epoch = 1 + int(test_losses.index(best_test))
- print(f"Best Test MSE: {best_test:.6f} @ epoch {best_epoch}")
- print("Note: trivial baseline MSE ≈ 1/6 ≈ 0.1667")
+ if ok_results:
+ print("\nTop results:")
+ for rank, result in enumerate(ok_results[: args.top_k], start=1):
+ print(
+ f"{rank:2d}. {result.model:24s} "
+ f"best_test_mse={result.best_test_mse:.6f} "
+ f"epoch={result.best_epoch} seconds={result.seconds:.1f}"
+ )
+ print("\nSaved:")
+ print(f" {outdir / 'summary.csv'}")
+ print(f" {outdir / 'summary.json'}")
+ for path in plot_paths:
+ print(f" {path}")
+ if plt is None and not plot_paths:
+ print(" plots skipped because matplotlib is not installed")
if __name__ == "__main__":
diff --git a/tests/test_cells.py b/tests/test_cells.py
index 17a0987..85f06ff 100755
--- a/tests/test_cells.py
+++ b/tests/test_cells.py
@@ -143,6 +143,74 @@ def test_taugru_cell_parameter_shapes():
assert cell.bias_hh.shape == (36,)
+def test_cornn_cell_defaults_are_damped():
+ cell = coRNNCell(4, 9)
+
+ assert cell.dt == pytest.approx(0.1)
+ assert cell.gamma == pytest.approx(1.0)
+ assert cell.epsilon == pytest.approx(1.0)
+
+
+def test_cornn_cell_matches_official_explicit_update():
+ cell = coRNNCell(
+ 2,
+ 2,
+ bias=False,
+ recurrent_bias=False,
+ cell_bias=False,
+ dt=0.5,
+ gamma=2.0,
+ epsilon=3.0,
+ )
+ with torch.no_grad():
+ cell.weight_ih.copy_(torch.tensor([[0.1, 0.2], [0.3, 0.4]]))
+ cell.weight_hh.copy_(torch.tensor([[0.5, 0.6], [0.7, 0.8]]))
+ cell.weight_ch.copy_(torch.tensor([[0.9, 1.0], [1.1, 1.2]]))
+
+ x = torch.tensor([[0.2, -0.3]])
+ h = torch.tensor([[0.4, -0.5]])
+ z = torch.tensor([[0.6, -0.7]])
+
+ new_h, new_z = cell(x, (h, z))
+ act = torch.tanh(
+ x @ cell.weight_ih.t() + h @ cell.weight_hh.t() + z @ cell.weight_ch.t()
+ )
+ expected_z = z + 0.5 * (act - 2.0 * h - 3.0 * z)
+ expected_h = h + 0.5 * expected_z
+
+ assert torch.allclose(new_z, expected_z)
+ assert torch.allclose(new_h, expected_h)
+
+
+def test_unicornn_cell_parameter_shapes_match_independent_recurrence():
+ cell = UnICORNNCell(4, 9)
+
+ assert cell.weight_ih.shape == (9, 4)
+ assert cell.weight_hh.shape == (9,)
+ assert cell.weight_ch.shape == (9,)
+
+
+def test_unicornn_cell_matches_official_independent_update():
+ cell = UnICORNNCell(2, 2, bias=False, recurrent_bias=False, dt=0.25, alpha=0.5)
+ with torch.no_grad():
+ cell.weight_ih.copy_(torch.tensor([[0.1, 0.2], [0.3, 0.4]]))
+ cell.weight_hh.copy_(torch.tensor([0.5, 0.6]))
+ cell.weight_ch.copy_(torch.tensor([0.7, -0.8]))
+
+ x = torch.tensor([[0.2, -0.3]])
+ h = torch.tensor([[0.4, -0.5]])
+ z = torch.tensor([[0.6, -0.7]])
+
+ new_h, new_z = cell(x, (h, z))
+ step = 0.25 * torch.sigmoid(cell.weight_ch)
+ candidate = torch.tanh(x @ cell.weight_ih.t() + h * cell.weight_hh)
+ expected_z = z - step * (candidate + 0.5 * h)
+ expected_h = h + step * expected_z
+
+ assert torch.allclose(new_z, expected_z)
+ assert torch.allclose(new_h, expected_h)
+
+
def test_taugru_cell_uses_delayed_state():
cell = tauGRUCell(1, 1)
with torch.no_grad():
diff --git a/torchrecurrent/base.py b/torchrecurrent/base.py
index 85eaecd..d27b322 100755
--- a/torchrecurrent/base.py
+++ b/torchrecurrent/base.py
@@ -100,6 +100,7 @@ def resolve_init_name(init: Any, default: str) -> str:
"kaiming_uniform",
"kaiming_normal",
"orthogonal",
+ "uniform_centered",
}
if name in known:
return name
@@ -119,6 +120,8 @@ def apply_init_(t: Tensor, name: str) -> None:
nn.init.normal_(t)
elif key == "uniform":
nn.init.uniform_(t)
+ elif key == "uniform_centered":
+ nn.init.uniform_(t, a=-0.1, b=0.1)
elif key == "xavier_uniform":
nn.init.xavier_uniform_(t)
elif key == "xavier_normal":
diff --git a/torchrecurrent/cells/cornn_cell.py b/torchrecurrent/cells/cornn_cell.py
index 3ffaadf..162de18 100755
--- a/torchrecurrent/cells/cornn_cell.py
+++ b/torchrecurrent/cells/cornn_cell.py
@@ -2,7 +2,12 @@
import torch.nn as nn
from torch import Tensor
from typing import Optional, Tuple
-from ..base import DoubleStateRecurrentLayerBase, DoubleStateCellBase, resolve_init_name, apply_init_
+from ..base import (
+ DoubleStateRecurrentLayerBase,
+ DoubleStateCellBase,
+ resolve_init_name,
+ apply_init_,
+)
class coRNN(DoubleStateRecurrentLayerBase):
@@ -52,9 +57,9 @@ class coRNN(DoubleStateRecurrentLayerBase):
`b_hh`. Default: True
cell_bias: If ``False``, then the layer does not use cell bias `b_ch`.
Default: True
- dt: Integration step size :math:`\Delta t`. Default: 1.0
- gamma: Damping coefficient on the hidden-state term. Default: 0.0
- epsilon: Damping coefficient on the cell-state term. Default: 0.0
+ dt: Integration step size :math:`\Delta t`. Default: 0.1
+ gamma: Damping coefficient on the hidden-state term. Default: 1.0
+ epsilon: Damping coefficient on the cell-state term. Default: 1.0
kernel_init: Initializer for `W_ih`. Default:
:func:`torch.nn.init.xavier_uniform_`
recurrent_kernel_init: Initializer for `W_{hh}`. Default:
@@ -190,9 +195,9 @@ class coRNNCell(DoubleStateCellBase):
``b_hh``. Default: ``True``.
cell_bias: If ``False``, the layer does not use cell bias ``b_ch``.
Default: ``True``.
- dt: Integration step size :math:`\Delta t`. Default: ``1.0``.
- gamma: Damping on hidden-state term. Default: ``0.0``.
- epsilon: Damping on cell-state term. Default: ``0.0``.
+ dt: Integration step size :math:`\Delta t`. Default: ``0.1``.
+ gamma: Damping on hidden-state term. Default: ``1.0``.
+ epsilon: Damping on cell-state term. Default: ``1.0``.
kernel_init: Initializer for ``W_{ih}``.
Default: :func:`torch.nn.init.xavier_uniform_`.
recurrent_kernel_init: Initializer for ``W_{hh}``.
@@ -279,9 +284,9 @@ def __init__(
bias: bool = True,
recurrent_bias: bool = True,
cell_bias: bool = True,
- dt: float = 1.0,
- gamma: float = 0.0,
- epsilon: float = 0.0,
+ dt: float = 0.1,
+ gamma: float = 1.0,
+ epsilon: float = 1.0,
kernel_init=nn.init.xavier_uniform_,
recurrent_kernel_init=nn.init.xavier_uniform_,
cell_kernel_init=nn.init.xavier_uniform_,
@@ -350,8 +355,16 @@ def forward(
b_c = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
else:
h, c = state
- b_h = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) if h is None else (h.unsqueeze(0) if (not is_batched and h.dim() == 1) else h)
- b_c = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) if c is None else (c.unsqueeze(0) if (not is_batched and c.dim() == 1) else c)
+ b_h = (
+ self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
+ if h is None
+ else (h.unsqueeze(0) if (not is_batched and h.dim() == 1) else h)
+ )
+ b_c = (
+ self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
+ if c is None
+ else (c.unsqueeze(0) if (not is_batched and c.dim() == 1) else c)
+ )
pre_act = (
b_inp @ self.weight_ih.t()
@@ -363,10 +376,7 @@ def forward(
)
act = torch.tanh(pre_act)
new_c = (
- b_c
- + self.dt * act
- - self.dt * self.gamma * b_h
- - self.dt * self.epsilon * b_c
+ b_c + self.dt * act - self.dt * self.gamma * b_h - self.dt * self.epsilon * b_c
)
new_h = b_h + self.dt * new_c
diff --git a/torchrecurrent/cells/unicornn_cell.py b/torchrecurrent/cells/unicornn_cell.py
index ad0af4a..d42bcaf 100755
--- a/torchrecurrent/cells/unicornn_cell.py
+++ b/torchrecurrent/cells/unicornn_cell.py
@@ -2,7 +2,12 @@
from torch import Tensor
import torch.nn as nn
from typing import Optional, Tuple
-from ..base import DoubleStateRecurrentLayerBase, DoubleStateCellBase, resolve_init_name, apply_init_
+from ..base import (
+ DoubleStateRecurrentLayerBase,
+ DoubleStateCellBase,
+ resolve_init_name,
+ apply_init_,
+)
class UnICORNN(DoubleStateRecurrentLayerBase):
@@ -18,7 +23,7 @@ class UnICORNN(DoubleStateRecurrentLayerBase):
\begin{aligned}
h(t) &= h(t-1) + \Delta t \, \hat{\sigma}(w_{ch}) \circ z(t), \\
z(t) &= z(t-1) - \Delta t \, \hat{\sigma}(w_{ch}) \circ
- \Bigl[\sigma(W_{hh} h(t-1) + W_{ih} x(t) + b_{ih})
+ \Bigl[\sigma(w_{hh} \circ h(t-1) + W_{ih} x(t) + b_{ih})
+ \alpha h(t-1)\Bigr],
\end{aligned}
@@ -36,18 +41,19 @@ class UnICORNN(DoubleStateRecurrentLayerBase):
`(batch, seq, feature)` format instead of `(seq, batch, feature)`.
Default: False
bias: If ``False``, disables input bias `b_{ih}`. Default: True
- recurrent_bias: If ``False``, disables hidden bias `b_{hh}`. Default: True
+ recurrent_bias: If ``True``, enables an extra hidden bias `b_{hh}`.
+ Default: False
kernel_init: Initializer for `W_{ih}`.
Default: :func:`torch.nn.init.xavier_uniform_`
- recurrent_kernel_init: Initializer for `W_{hh}`.
- Default: :func:`torch.nn.init.xavier_uniform_`
+ recurrent_kernel_init: Initializer for `w_{hh}`.
+ Default: :func:`torch.nn.init.uniform_`
control_kernel_init: Initializer for `w_{ch}`.
- Default: :func:`torch.nn.init.normal_`
+ Default: uniform on ``[-0.1, 0.1]``
bias_init: Initializer for `b_{ih}` when ``bias=True``.
Default: :func:`torch.nn.init.zeros_`
recurrent_bias_init: Initializer for `b_{hh}` when ``recurrent_bias=True``.
Default: :func:`torch.nn.init.zeros_`
- dt: Integration step :math:`\Delta t`. Default: 1.0
+ dt: Integration step :math:`\Delta t`. Default: 0.1
alpha: Leakage coefficient :math:`\alpha`. Default: 0.0
device: Desired device of parameters.
dtype: Desired floating point type of parameters.
@@ -85,13 +91,13 @@ class UnICORNN(DoubleStateRecurrentLayerBase):
cells.{k}.weight_ih : input–hidden weights of the :math:`k`-th layer,
shape `(hidden_size, input_size)` for `k=0`,
otherwise `(hidden_size, hidden_size)`.
- cells.{k}.weight_hh : hidden–hidden weights of the :math:`k`-th layer,
- shape `(hidden_size, hidden_size)`.
+ cells.{k}.weight_hh : elementwise hidden weights of the :math:`k`-th layer,
+ shape `(hidden_size,)`.
cells.{k}.weight_ch : control weights of the :math:`k`-th layer,
shape `(hidden_size,)`.
cells.{k}.bias_ih : input bias of the :math:`k`-th layer,
shape `(hidden_size,)` if ``bias=True``.
- cells.{k}.bias_hh : hidden bias of the :math:`k`-th layer,
+ cells.{k}.bias_hh : optional hidden bias of the :math:`k`-th layer,
shape `(hidden_size,)` if ``recurrent_bias=True``.
.. seealso::
@@ -140,7 +146,7 @@ class UnICORNNCell(DoubleStateCellBase):
- \Delta t\,\hat{\sigma}(\mathbf{w}_{ch}) \circ
\Bigl[
\sigma\bigl(
- \mathbf{W}_{hh}\,\mathbf{h}(t-1)
+ \mathbf{w}_{hh} \circ \mathbf{h}(t-1)
+ \mathbf{W}_{ih}\,\mathbf{x}(t)
+ \mathbf{b}_{ih}
\bigr)
@@ -156,19 +162,19 @@ class UnICORNNCell(DoubleStateCellBase):
hidden_size: The number of features in the hidden state ``h``.
bias: If ``False``, the layer does not use the input bias
``b_{ih}``. Default: ``True``.
- recurrent_bias: If ``False``, the layer does not use the hidden
- bias ``b_{hh}``. Default: ``True``.
+ recurrent_bias: If ``True``, the layer uses an extra hidden
+ bias ``b_{hh}``. Default: ``False``.
kernel_init: Initializer for ``W_{ih}``.
Default: :func:`torch.nn.init.xavier_uniform_`.
- recurrent_kernel_init: Initializer for ``W_{hh}``.
- Default: :func:`torch.nn.init.xavier_uniform_`.
+ recurrent_kernel_init: Initializer for ``w_{hh}``.
+ Default: :func:`torch.nn.init.uniform_`.
control_kernel_init: Initializer for ``w_{ch}``.
- Default: :func:`torch.nn.init.normal_`.
+ Default: uniform on ``[-0.1, 0.1]``.
bias_init: Initializer for ``b_{ih}``.
Default: :func:`torch.nn.init.zeros_`.
recurrent_bias_init: Initializer for ``b_{hh}``.
Default: :func:`torch.nn.init.zeros_`.
- dt: Time step :math:`\Delta t` between updates. Default: ``1.0``.
+ dt: Time step :math:`\Delta t` between updates. Default: ``0.1``.
alpha: Leakage coefficient in the control update. Default: ``0.0``.
device: The desired device of parameters.
dtype: The desired floating point type of parameters.
@@ -192,14 +198,14 @@ class UnICORNNCell(DoubleStateCellBase):
Variables:
weight_ih: The learnable input–hidden weights,
of shape ``(hidden_size, input_size)``.
- weight_hh: The learnable hidden–hidden weights,
- of shape ``(hidden_size, hidden_size)``.
+ weight_hh: The learnable elementwise hidden weights,
+ of shape ``(hidden_size,)``.
weight_ch: The learnable control weights,
of shape ``(hidden_size,)``.
bias_ih: The learnable input bias,
of shape ``(hidden_size,)``.
- bias_hh: The learnable hidden bias,
- of shape ``(hidden_size,)``.
+ bias_hh: The optional learnable hidden bias,
+ of shape ``(hidden_size,)`` when ``recurrent_bias=True``.
Examples::
@@ -230,13 +236,13 @@ def __init__(
input_size: int,
hidden_size: int,
bias: bool = True,
- recurrent_bias: bool = True,
+ recurrent_bias: bool = False,
kernel_init=nn.init.xavier_uniform_,
- recurrent_kernel_init=nn.init.xavier_uniform_,
- control_kernel_init=nn.init.normal_,
+ recurrent_kernel_init=nn.init.uniform_,
+ control_kernel_init="uniform_centered",
bias_init=nn.init.zeros_,
recurrent_bias_init=nn.init.zeros_,
- dt: float = 1.0,
+ dt: float = 0.1,
alpha: float = 0.0,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
@@ -248,9 +254,11 @@ def __init__(
self.alpha = alpha
self.init_cfg["kernel"] = resolve_init_name(kernel_init, self.init_cfg["kernel"])
self.init_cfg["recurrent_kernel"] = resolve_init_name(
- recurrent_kernel_init, self.init_cfg["recurrent_kernel"]
+ recurrent_kernel_init, "uniform"
+ )
+ self.init_cfg["control_kernel"] = resolve_init_name(
+ control_kernel_init, "uniform_centered"
)
- self.init_cfg["control_kernel"] = resolve_init_name(control_kernel_init, "normal")
self.init_cfg["bias"] = resolve_init_name(bias_init, self.init_cfg["bias"])
self.init_cfg["recurrent_bias"] = resolve_init_name(
recurrent_bias_init, self.init_cfg["recurrent_bias"]
@@ -259,7 +267,7 @@ def __init__(
self._register_tensors(
{
"weight_ih": ((hidden_size, input_size), True),
- "weight_hh": ((hidden_size, hidden_size), True),
+ "weight_hh": ((hidden_size,), True),
"weight_ch": ((hidden_size,), True),
"bias_ih": ((hidden_size,), bias),
"bias_hh": ((hidden_size,), recurrent_bias),
@@ -288,14 +296,19 @@ def forward(
b_c = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
else:
h, c = state
- b_h = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) if h is None else (h.unsqueeze(0) if (not is_batched and h.dim() == 1) else h)
- b_c = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) if c is None else (c.unsqueeze(0) if (not is_batched and c.dim() == 1) else c)
+ b_h = (
+ self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
+ if h is None
+ else (h.unsqueeze(0) if (not is_batched and h.dim() == 1) else h)
+ )
+ b_c = (
+ self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype)
+ if c is None
+ else (c.unsqueeze(0) if (not is_batched and c.dim() == 1) else c)
+ )
candidate_state = torch.tanh(
- b_inp @ self.weight_ih.t()
- + self.bias_ih
- + b_h @ self.weight_hh.t()
- + self.bias_hh
+ b_inp @ self.weight_ih.t() + self.bias_ih + b_h * self.weight_hh + self.bias_hh
)
new_c = b_c - self.dt * torch.sigmoid(self.weight_ch) * (
candidate_state + self.alpha * b_h