diff --git a/README.md b/README.md index 7af2d07..af8ea13 100644 --- a/README.md +++ b/README.md @@ -45,15 +45,15 @@ We provide [datasets](memax/datasets) to test our recurrent models. > > **Sequence Lengths:** `[784]` -### MNIST Math [[HuggingFace]](https://huggingface.co/datasets?sort=trending&search=bolt-lab%2Fmnist-math) [[Code]](memax/datasets/sequential_mnist.py) +### MNIST Math [[HuggingFace]](https://huggingface.co/datasets/bolt-lab/mnist-math-100) [[Code]](memax/datasets/mnist_math.py) > The recurrent model receives a sequence of MNIST images and operators, pixel by pixel, and must predict the percentile of the operators applied to the MNIST image classes. > -> **Sequence Lengths:** `[784 * 5, 784 * 100, 784 * 1_000, 784 * 10_000, 784 * 1_000_000]` +> **Sequence Lengths:** `[784 * 100, …]` on Hub as `bolt-lab/mnist-math-{length}` (default loader: `100`; `100_000` is under `smorad/mnist-math-100000`). -### Continuous Localization [[HuggingFace]](https://huggingface.co/datasets?sort=trending&search=bolt-lab%2Fcontinuous-localization) [[Code]](memax/datasets/sequential_mnist.py) +### Continuous Localization [[HuggingFace]](https://huggingface.co/datasets/bolt-lab/continuous-localization-20) [[Code]](memax/datasets/continuous_localization.py) > The recurrent model receives a sequence of translation and rotation vectors **in the local coordinate frame**, and must predict the corresponding position and orientation **in the global coordinate frame**. > -> **Sequence Lengths:** `[20, 100, 1_000]` +> **Sequence Lengths:** `[20, 100]` on Hub as `bolt-lab/continuous-localization-{length}`. # Getting Started Install `memax` using pip for your specific framework: @@ -105,10 +105,23 @@ hs, ys = filter_jit(filter_vmap(model))(hs_0, (xs, starts)) ``` ## Running Baselines -You can compare various recurrent models on our datasets with a single command +You can compare various recurrent models on our datasets with a single command: ```bash -python run_equinox_experiments.py # equinox framework -python run_linen_experiments.py # flax linen framework +python run_equinox_experiments.py --dataset-name sequential_mnist +python run_linen_experiments.py --dataset-name sequential_mnist +``` + +Dataset names: see `memax.experiments.datasets.DATASET_NAMES` (e.g. `sequential_mnist`, `mnist_math`, `mnist-math-100`, `continuous-localization-20`). Legacy alias: `sequential_rotation`. + +For a quick CPU smoke run (tiny model, few optimizer steps): +```bash +python run_equinox_experiments.py --dataset-name sequential_mnist --smoke --models GRU +python run_linen_experiments.py --dataset-name sequential_mnist --smoke --models GRU +``` + +CI runs synthetic smoke tests only (`pytest`; no Hugging Face downloads). To exercise real datasets locally: +```bash +pytest -m integration tests/test_experiment_smoke.py ``` diff --git a/memax/datasets/continuous_localization.py b/memax/datasets/continuous_localization.py index efbbd0e..f21a82a 100644 --- a/memax/datasets/continuous_localization.py +++ b/memax/datasets/continuous_localization.py @@ -1,17 +1,26 @@ """This module generates the continuous localization dataset and uploads it to Hugging Face. The dataset consists of sequences of 3D rotations and translations, along with the absolute positions and -orientations. Each sequence is generated by applying a series of random small rotations and translations to an initial position and orientation. The goal is to predict the absolute position and orientation from the sequence of relative movements.""" +orientations. Each sequence is generated by applying a series of random small rotations and translations to an initial position and orientation. The goal is to predict the absolute position and orientation from the sequence of relative movements. +""" -import jax.numpy as jnp import jax +import jax.numpy as jnp +from datasets import Array2D, Dataset, Features, load_dataset from jax.scipy.spatial.transform import Rotation -from datasets import Dataset, Features, Array2D, load_dataset -SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000] +from memax.datasets.hub import ( + CONTINUOUS_LOCALIZATION_ORG, + CONTINUOUS_LOCALIZATION_SEQ_LENS, + continuous_localization_hub_id, +) + +# Lengths used when generating/uploading new Hub revisions. +UPLOAD_SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000] + def step(carry, inputs): - (x, rot) = carry - (dx, drot) = inputs + x, rot = carry + dx, drot = inputs transform_x = jax.vmap(lambda rot_t, x_t, dx: rot_t.apply(dx) + x_t) transform_r = jax.vmap(lambda rot_t, drot: rot_t * drot) x = transform_x(rot, x, dx) @@ -19,19 +28,19 @@ def step(carry, inputs): return ((x, rot), (x, rot)) -def generate_dataset( - key, - batch_size, - num_steps -): +def generate_dataset(key, batch_size, num_steps): spatial_dim = 3 keys = jax.random.split(key, 4) x = jnp.zeros((batch_size, spatial_dim)) rot = jax.vmap(Rotation.identity, axis_size=batch_size)() dx_mag = jax.random.exponential(keys[0], (num_steps, batch_size, 1)) * 2 - dx_dir = jax.random.uniform(keys[1], (num_steps, batch_size, 3), minval=-1., maxval=1.) + dx_dir = jax.random.uniform( + keys[1], (num_steps, batch_size, 3), minval=-1.0, maxval=1.0 + ) dx = dx_dir / jnp.linalg.norm(dx_dir, axis=-1, keepdims=True) * dx_mag - drot_dir = jax.random.uniform(keys[2], (num_steps, batch_size, spatial_dim), minval=-1, maxval=1) + drot_dir = jax.random.uniform( + keys[2], (num_steps, batch_size, spatial_dim), minval=-1, maxval=1 + ) drot_mag = jax.random.uniform(keys[3], (num_steps, batch_size, 1), maxval=jnp.pi) drot_vec = drot_dir / jnp.linalg.norm(drot_dir, axis=-1, keepdims=True) * drot_mag drot = Rotation.from_rotvec(drot_vec) @@ -45,56 +54,62 @@ def generate_dataset( outputs = jnp.concatenate([rot_abs_vec, x_abs], axis=-1) return { - "inputs": jnp.permute_dims(inputs, (1,0,2)), - "outputs": jnp.permute_dims(outputs, (1,0,2)), - "delta_rotation": jnp.permute_dims(drot_vec, (1,0,2)), - "delta_position": jnp.permute_dims(dx, (1,0,2)), - "absolute_rotation": jnp.permute_dims(rot_abs_vec, (1,0,2)), - "absolute_position": jnp.permute_dims(x_abs, (1,0,2)), + "inputs": jnp.permute_dims(inputs, (1, 0, 2)), + "outputs": jnp.permute_dims(outputs, (1, 0, 2)), + "delta_rotation": jnp.permute_dims(drot_vec, (1, 0, 2)), + "delta_position": jnp.permute_dims(dx, (1, 0, 2)), + "absolute_rotation": jnp.permute_dims(rot_abs_vec, (1, 0, 2)), + "absolute_position": jnp.permute_dims(x_abs, (1, 0, 2)), } def upload_hf_datasets( - batch_size_train = 1e6, - batch_size_test = 2e5, - sequence_length = 20, + batch_size_train=1e6, + batch_size_test=2e5, + sequence_length=20, ): - name = f"bolt-lab/continuous-localization-{sequence_length}", + name = f"{CONTINUOUS_LOCALIZATION_ORG}/continuous-localization-{sequence_length}" batch_size_train = int(batch_size_train) batch_size_test = int(batch_size_test) - FEATURES = Features({ - "inputs": Array2D(dtype='float32', shape=(sequence_length, 6)), - "outputs": Array2D(dtype='float32', shape=(sequence_length, 6)), - "delta_rotation": Array2D(dtype='float32', shape=(sequence_length, 3)), - "delta_position": Array2D(dtype='float32', shape=(sequence_length, 3)), - "absolute_rotation": Array2D(dtype='float32', shape=(sequence_length, 3)), - "absolute_position": Array2D(dtype='float32', shape=(sequence_length, 3)), - }) + FEATURES = Features( + { + "inputs": Array2D(dtype="float32", shape=(sequence_length, 6)), + "outputs": Array2D(dtype="float32", shape=(sequence_length, 6)), + "delta_rotation": Array2D(dtype="float32", shape=(sequence_length, 3)), + "delta_position": Array2D(dtype="float32", shape=(sequence_length, 3)), + "absolute_rotation": Array2D(dtype="float32", shape=(sequence_length, 3)), + "absolute_position": Array2D(dtype="float32", shape=(sequence_length, 3)), + } + ) key = jax.random.key(0) keys = jax.random.split(key, 2) - train_dict = generate_dataset(keys[0], batch_size=batch_size_train, num_steps=sequence_length) - test_dict = generate_dataset(keys[1], batch_size=batch_size_test, num_steps=sequence_length) - train_dataset = Dataset.from_dict(train_dict).cast(features=FEATURES, batch_size=batch_size_train) - test_dataset = Dataset.from_dict(test_dict).cast(features=FEATURES, batch_size=batch_size_test) + train_dict = generate_dataset( + keys[0], batch_size=batch_size_train, num_steps=sequence_length + ) + test_dict = generate_dataset( + keys[1], batch_size=batch_size_test, num_steps=sequence_length + ) + train_dataset = Dataset.from_dict(train_dict).cast( + features=FEATURES, batch_size=batch_size_train + ) + test_dataset = Dataset.from_dict(test_dict).cast( + features=FEATURES, batch_size=batch_size_test + ) train_dataset.push_to_hub(name, split="train") test_dataset.push_to_hub(name, split="test") + def get_rot_dataset(sequence_length=20): - seq_lens = [20] - assert sequence_length in seq_lens, f"Invalid sequenec length, must be one of {seq_lens}" - num_labels = 10 - dataset = load_dataset(f"bolt-lab/continuous-localization-{sequence_length}") + assert ( + sequence_length in CONTINUOUS_LOCALIZATION_SEQ_LENS + ), f"Invalid sequence length, must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}" + dataset = load_dataset(continuous_localization_hub_id(sequence_length)) x = jnp.array(dataset["train"]["delta_rotation"]) - y = jnp.array(dataset["train"]["absolute_rotation"][-1]) - #y = jnp.array(dataset["train"]["absolute_rotation_percentile"]) - # TODO: One-hot based on abs of rotation angle? - #y = jax.nn.one_hot(y, num_labels) + y = jnp.array(dataset["train"]["absolute_rotation"])[:, -1, :] test_x = jnp.array(dataset["test"]["delta_rotation"]) - #test_y = jnp.array(dataset["test"]["absolute_rotation_percentile"]) - test_y = jnp.array(dataset["test"]["absolute_rotation"][-1]) - #test_y = jax.nn.one_hot(test_y, num_labels) + test_y = jnp.array(dataset["test"]["absolute_rotation"])[:, -1, :] return { "x_train": x, @@ -102,25 +117,21 @@ def get_rot_dataset(sequence_length=20): "x_test": test_x, "y_test": test_y, "size": x.shape[0], - "num_labels": 10, + "num_labels": y.shape[-1], } + def get_trans_dataset(sequence_length=20): - seq_lens = [20, 1024] - assert sequence_length in seq_lens, f"Invalid sequenec length, must be one of {seq_lens}" - dataset = load_dataset(f"bolt-lab/continuous-localization-{sequence_length}") - - x = jnp.concatenate([ - jnp.array(dataset["train"]["delta_rotation"]), - jnp.array(dataset["train"]["delta_translation"]) - ], axis=-1) - y = jnp.array(dataset["train"]["absolute_translation_percentile"]) - - test_x = jnp.concatenate([ - jnp.array(dataset["test"]["delta_rotation"]), - jnp.array(dataset["test"]["delta_translation"]) - ], axis=-1) - test_y = jnp.array(dataset["test"]["absolute_translation_percentile"]) + assert ( + sequence_length in CONTINUOUS_LOCALIZATION_SEQ_LENS + ), f"Invalid sequence length, must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}" + dataset = load_dataset(continuous_localization_hub_id(sequence_length)) + + x = jnp.array(dataset["train"]["inputs"]) + y = jnp.array(dataset["train"]["outputs"])[:, -1, :] + + test_x = jnp.array(dataset["test"]["inputs"]) + test_y = jnp.array(dataset["test"]["outputs"])[:, -1, :] return { "x_train": x, @@ -128,10 +139,12 @@ def get_trans_dataset(sequence_length=20): "x_test": test_x, "y_test": test_y, "size": x.shape[0], - "num_labels": 10, + "num_labels": y.shape[-1], } if __name__ == "__main__": - for seq_len in SEQ_LENS: - upload_hf_datasets(batch_size_train=60_000, batch_size_test=10_000, sequence_length=seq_len) \ No newline at end of file + for seq_len in UPLOAD_SEQ_LENS: + upload_hf_datasets( + batch_size_train=60_000, batch_size_test=10_000, sequence_length=seq_len + ) diff --git a/memax/datasets/hub.py b/memax/datasets/hub.py new file mode 100644 index 0000000..1971322 --- /dev/null +++ b/memax/datasets/hub.py @@ -0,0 +1,47 @@ +"""Hugging Face Hub dataset identifiers used by memax loaders. + +Canonical dataset pages: +- https://huggingface.co/datasets/ylecun/mnist +- https://huggingface.co/datasets/bolt-lab/mnist-math-{seq_len} +- https://huggingface.co/datasets/bolt-lab/continuous-localization-{seq_len} +""" + +# Raw MNIST images (README: ylecun/mnist). The legacy id "mnist" redirects but logs hub warnings. +MNIST = "ylecun/mnist" + +# Published under bolt-lab; 100_000 lives only under smorad. +MNIST_MATH_ORG = "bolt-lab" +MNIST_MATH_ORG_ALT = "smorad" +MNIST_MATH_SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000] + +CONTINUOUS_LOCALIZATION_ORG = "bolt-lab" +# Datasets available on the Hub as of upload (longer seq_lens may be added separately). +CONTINUOUS_LOCALIZATION_SEQ_LENS = [20, 100] + + +def mnist_math_hub_id(seq_len: int) -> str: + """Return the Hub repo id for a MNIST Math sequence length.""" + if seq_len not in MNIST_MATH_SEQ_LENS: + raise ValueError( + f"Invalid seq_len {seq_len}, must be one of {MNIST_MATH_SEQ_LENS}" + ) + if seq_len == 100_000: + return f"{MNIST_MATH_ORG_ALT}/mnist-math-{seq_len}" + return f"{MNIST_MATH_ORG}/mnist-math-{seq_len}" + + +def mnist_math_hub_url(seq_len: int) -> str: + return f"https://huggingface.co/datasets/{mnist_math_hub_id(seq_len)}" + + +def continuous_localization_hub_id(sequence_length: int) -> str: + if sequence_length not in CONTINUOUS_LOCALIZATION_SEQ_LENS: + raise ValueError( + f"Invalid sequence_length {sequence_length}, " + f"must be one of {CONTINUOUS_LOCALIZATION_SEQ_LENS}" + ) + return f"{CONTINUOUS_LOCALIZATION_ORG}/continuous-localization-{sequence_length}" + + +def continuous_localization_hub_url(sequence_length: int) -> str: + return f"https://huggingface.co/datasets/{continuous_localization_hub_id(sequence_length)}" diff --git a/memax/datasets/mnist_listops.py b/memax/datasets/mnist_listops.py index bcea732..a09f4ce 100644 --- a/memax/datasets/mnist_listops.py +++ b/memax/datasets/mnist_listops.py @@ -4,9 +4,9 @@ import jax.numpy as jnp from datasets import load_dataset # huggingface datasets +from memax.datasets.hub import MNIST from memax.train_utils import get_residual_memory_models - NUM_EPOCHS = 100 BATCH_SIZE = 32 SEQ_LEN = 784 @@ -105,7 +105,7 @@ def normalize_and_flatten(x): def make_dataset(dataset_size=3, num_terms=5, key=jax.random.key(0), batch_size=32): - dataset = load_dataset("mnist") + dataset = load_dataset(MNIST) x, y = jnp.array(dataset["train"]["image"]), jnp.array(dataset["train"]["label"]) keys = jax.random.split(key, 3) diff --git a/memax/datasets/mnist_math.py b/memax/datasets/mnist_math.py index adec845..ddf2c33 100644 --- a/memax/datasets/mnist_math.py +++ b/memax/datasets/mnist_math.py @@ -1,6 +1,6 @@ """ This file contains the code for creating the MNIST Math dataset, -as well as the code that preprocesses the dataset before training. +as well as the code that preprocesses the dataset before training. The MNIST Math dataset consists of sequences of MNIST digits interspersed with plus and minus operators. The task is to compute the cumulative result of the arithmetic expression formed by the digits @@ -16,8 +16,10 @@ import tqdm from datasets import Dataset, Features, Image, Sequence, Value, load_dataset +from memax.datasets.hub import MNIST, MNIST_MATH_SEQ_LENS, mnist_math_hub_id + NUM_LABELS = 10 -SEQ_LENS = [100, 1_000, 10_000, 100_000, 1_000_000] +SEQ_LENS = MNIST_MATH_SEQ_LENS FEATURES = Features( { @@ -173,7 +175,7 @@ def generate_dataset(key, mnist, dataset_size=10, num_terms=5): def make_hf_datasets(): - mnist = load_dataset("mnist") + mnist = load_dataset(MNIST) train_dset = generate_dataset(jax.random.key(0), mnist["train"], 60_000) test_dset = generate_dataset(jax.random.key(1), mnist["test"], 10_000) train_dset = Dataset.from_dict(train_dset).cast(FEATURES) @@ -183,13 +185,14 @@ def make_hf_datasets(): def upload_hf_datasets(length): train, test = make_hf_datasets() - train.push_to_hub(f"smorad/mnist-math-{length}", split="train") - test.push_to_hub(f"smorad/mnist-math-{length}", split="test") + repo_id = mnist_math_hub_id(length) + train.push_to_hub(repo_id, split="train") + test.push_to_hub(repo_id, split="test") -def get_dataset(seq_len=5): +def get_dataset(seq_len=100): assert seq_len in SEQ_LENS, f"Invalid length, must be in {SEQ_LENS}" - dataset = load_dataset("smorad/mnist-math").with_format("np") + dataset = load_dataset(mnist_math_hub_id(seq_len)).with_format("np") num_labels = 10 x = dataset["train"]["image_flat"] diff --git a/memax/datasets/sequential_mnist.py b/memax/datasets/sequential_mnist.py index ce33422..8589999 100644 --- a/memax/datasets/sequential_mnist.py +++ b/memax/datasets/sequential_mnist.py @@ -7,6 +7,8 @@ import jax.numpy as jnp from datasets import load_dataset +from memax.datasets.hub import MNIST + def normalize_and_flatten(x): # batch, time, feature @@ -15,7 +17,7 @@ def normalize_and_flatten(x): def get_dataset(): - dataset = load_dataset("mnist") + dataset = load_dataset(MNIST) num_labels = 10 x = jnp.array(dataset["train"]["image"]) @@ -35,4 +37,4 @@ def get_dataset(): "y_test": test_y, "num_labels": num_labels, "size": x.shape[0], - } \ No newline at end of file + } diff --git a/memax/experiments/__init__.py b/memax/experiments/__init__.py new file mode 100644 index 0000000..0a1b272 --- /dev/null +++ b/memax/experiments/__init__.py @@ -0,0 +1,34 @@ +"""Shared utilities for memax training experiment scripts.""" + +from memax.experiments.config import ( + ExperimentConfig, + apply_smoke_defaults, + fill_dataset_defaults, + prepare_config, +) +from memax.experiments.datasets import ( + DATASET_NAMES, + get_default_hyperparameters, + get_default_loss_name, + load_dataset, + make_synthetic_dataset, + slice_dataset, +) + +# Re-export sorted registry keys for CLI help / docs +from memax.experiments.runner import run_equinox_training, run_linen_training + +__all__ = [ + "DATASET_NAMES", + "ExperimentConfig", + "apply_smoke_defaults", + "fill_dataset_defaults", + "prepare_config", + "get_default_hyperparameters", + "get_default_loss_name", + "load_dataset", + "make_synthetic_dataset", + "run_equinox_training", + "run_linen_training", + "slice_dataset", +] diff --git a/memax/experiments/cli.py b/memax/experiments/cli.py new file mode 100644 index 0000000..3e3f168 --- /dev/null +++ b/memax/experiments/cli.py @@ -0,0 +1,71 @@ +"""Shared argparse flags for experiment scripts.""" + +import argparse + + +def add_experiment_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--seed", type=int, default=0, help="Random seed") + parser.add_argument("--num-epochs", type=int, default=None, help="Training epochs") + parser.add_argument("--batch-size", type=int, default=None, help="Batch size") + parser.add_argument( + "--recurrent-size", + "--recurrent_size", + dest="recurrent_size", + type=int, + default=None, + help="Recurrent hidden size", + ) + parser.add_argument("--num-layers", type=int, default=None, help="Residual layers") + parser.add_argument("--lr", type=float, default=None, help="Learning rate") + parser.add_argument( + "--use-wandb", + action="store_true", + default=False, + help="Log to Weights & Biases", + ) + parser.add_argument( + "--project-name", + type=str, + default="memax-debug", + help="W&B project name", + ) + parser.add_argument( + "--dataset-name", + "--dataset_name", + dest="dataset_name", + type=str, + default="sequential_mnist", + help=("Dataset name (see memax.experiments.datasets.DATASET_NAMES)"), + ) + parser.add_argument( + "--loss-function", + "--loss_function", + dest="loss_function", + type=str, + default=None, + help="loss_classify_terminal_output or loss_regress_terminal_output", + ) + parser.add_argument( + "--models", + type=str, + nargs="+", + default=None, + help="Model names (default: all)", + ) + parser.add_argument( + "--max-train-samples", + type=int, + default=None, + help="Cap training set size after loading", + ) + parser.add_argument( + "--max-updates", + type=int, + default=None, + help="Max optimizer steps per epoch", + ) + parser.add_argument( + "--smoke", + action="store_true", + help="CPU-friendly defaults (tiny model, few steps)", + ) diff --git a/memax/experiments/config.py b/memax/experiments/config.py new file mode 100644 index 0000000..1b30e28 --- /dev/null +++ b/memax/experiments/config.py @@ -0,0 +1,76 @@ +"""Experiment configuration shared by equinox and linen training scripts.""" + +from dataclasses import dataclass, fields +from typing import List, Optional, Union + +from memax.experiments.datasets import get_default_hyperparameters + + +@dataclass +class ExperimentConfig: + seed: int = 0 + num_epochs: Optional[int] = None + batch_size: Optional[int] = None + recurrent_size: Optional[int] = None + num_layers: Optional[int] = None + lr: Optional[float] = None + use_wandb: bool = False + project_name: str = "memax-debug" + dataset_name: str = "sequential_mnist" + loss_function: Optional[str] = None + models: Union[str, List[str]] = "all" + max_train_samples: Optional[int] = None + max_updates: Optional[int] = None + smoke: bool = False + + @classmethod + def from_namespace(cls, args) -> "ExperimentConfig": + models = getattr(args, "models", None) + if models is None: + models = "all" + return cls( + seed=args.seed, + num_epochs=getattr(args, "num_epochs", None), + batch_size=getattr(args, "batch_size", None), + recurrent_size=getattr(args, "recurrent_size", None), + num_layers=getattr(args, "num_layers", None), + lr=getattr(args, "lr", None), + use_wandb=args.use_wandb, + project_name=args.project_name, + dataset_name=getattr(args, "dataset_name", "sequential_mnist"), + loss_function=getattr(args, "loss_function", None), + models=models, + max_train_samples=getattr(args, "max_train_samples", None), + max_updates=getattr(args, "max_updates", None), + smoke=getattr(args, "smoke", False), + ) + + +def apply_smoke_defaults(config: ExperimentConfig) -> ExperimentConfig: + if not config.smoke: + return config + config.num_epochs = 1 + config.batch_size = 2 + config.recurrent_size = 16 + config.num_layers = 1 + config.lr = 1e-3 + config.max_train_samples = config.max_train_samples or 8 + config.max_updates = config.max_updates or 2 + if config.models == "all": + config.models = ["GRU"] + return config + + +def fill_dataset_defaults(config: ExperimentConfig) -> ExperimentConfig: + defaults = get_default_hyperparameters(config.dataset_name) + for f in fields(config): + if f.name in defaults and getattr(config, f.name) is None: + setattr(config, f.name, defaults[f.name]) + return config + + +def prepare_config(args) -> ExperimentConfig: + config = ExperimentConfig.from_namespace(args) + config = apply_smoke_defaults(config) + config = fill_dataset_defaults(config) + return config diff --git a/memax/experiments/datasets.py b/memax/experiments/datasets.py new file mode 100644 index 0000000..749e9fb --- /dev/null +++ b/memax/experiments/datasets.py @@ -0,0 +1,187 @@ +"""Dataset loading for memax experiment scripts.""" + +from typing import Any, Dict, Optional, Tuple + +import jax +import jax.numpy as jnp + +from memax.datasets.continuous_localization import get_rot_dataset, get_trans_dataset +from memax.datasets.hub import ( + CONTINUOUS_LOCALIZATION_SEQ_LENS, + MNIST_MATH_SEQ_LENS, + continuous_localization_hub_id, + mnist_math_hub_id, +) +from memax.datasets.mnist_math import get_dataset as get_mnist_math +from memax.datasets.sequential_mnist import get_dataset as get_sequential_mnist + +# canonical_name -> defaults (loss + hyperparameters use canonical keys only) +_CANONICAL_LOSS = { + "sequential_mnist": "loss_classify_terminal_output", + "mnist_math": "loss_classify_terminal_output", + "continuous_localization_rotation": "loss_regress_terminal_output", + "continuous_localization": "loss_regress_terminal_output", +} + +_CANONICAL_HYPERPARAMETERS = { + "sequential_mnist": { + "num_epochs": 5, + "batch_size": 16, + "recurrent_size": 256, + "num_layers": 2, + "lr": 1e-4, + }, + "mnist_math": { + "num_epochs": 5, + "batch_size": 16, + "recurrent_size": 256, + "num_layers": 2, + "lr": 1e-4, + }, + "continuous_localization_rotation": { + "num_epochs": 5, + "batch_size": 16, + "recurrent_size": 256, + "num_layers": 2, + "lr": 1e-4, + }, + "continuous_localization": { + "num_epochs": 5, + "batch_size": 16, + "recurrent_size": 256, + "num_layers": 2, + "lr": 1e-4, + }, +} + +# CLI / Hub name -> (canonical task, kwargs for the loader) +_DATASET_REGISTRY: Dict[str, Tuple[str, Dict[str, Any]]] = { + "sequential_mnist": ("sequential_mnist", {}), + # MNIST Math (default length 100) + "mnist_math": ("mnist_math", {"seq_len": 100}), + # Continuous localization (default length 20) + "continuous_localization": ( + "continuous_localization", + {"sequence_length": 20}, + ), + "continuous_localization_rotation": ( + "continuous_localization_rotation", + {"sequence_length": 20}, + ), + "sequential_rotation": ( + "continuous_localization_rotation", + {"sequence_length": 20}, + ), # legacy alias +} + +# Register every Hub MNIST Math repo id and friendly CLI names. +for _seq_len in MNIST_MATH_SEQ_LENS: + _entry = ("mnist_math", {"seq_len": _seq_len}) + _DATASET_REGISTRY[f"mnist-math-{_seq_len}"] = _entry + _DATASET_REGISTRY[mnist_math_hub_id(_seq_len)] = _entry + +# Register every Hub continuous-localization repo id and friendly CLI names. +for _length in CONTINUOUS_LOCALIZATION_SEQ_LENS: + _full = ("continuous_localization", {"sequence_length": _length}) + _rot = ("continuous_localization_rotation", {"sequence_length": _length}) + _DATASET_REGISTRY[f"continuous-localization-{_length}"] = _full + _DATASET_REGISTRY[continuous_localization_hub_id(_length)] = _full + _DATASET_REGISTRY[f"continuous-localization-rotation-{_length}"] = _rot + +DATASET_NAMES = tuple(sorted(_DATASET_REGISTRY.keys())) + + +def get_default_loss_name(dataset_name: str) -> str: + canonical, _ = _resolve_name(dataset_name) + if canonical not in _CANONICAL_LOSS: + raise ValueError(f"No default loss for dataset: {dataset_name}") + return _CANONICAL_LOSS[canonical] + + +def get_default_hyperparameters(dataset_name: str) -> Dict[str, Any]: + canonical, _ = _resolve_name(dataset_name) + if canonical not in _CANONICAL_HYPERPARAMETERS: + raise ValueError(f"No default hyperparameters for dataset: {dataset_name}") + return dict(_CANONICAL_HYPERPARAMETERS[canonical]) + + +def _resolve_name(dataset_name: str) -> Tuple[str, Dict[str, Any]]: + if dataset_name not in _DATASET_REGISTRY: + raise ValueError( + f"Unknown dataset: {dataset_name}. " + f"Choose from: {', '.join(DATASET_NAMES)}" + ) + canonical, extra = _DATASET_REGISTRY[dataset_name] + return canonical, dict(extra) + + +def load_dataset(dataset_name: str, **kwargs) -> Dict[str, Any]: + canonical, extra = _resolve_name(dataset_name) + opts = {**extra, **kwargs} + if canonical == "sequential_mnist": + return get_sequential_mnist() + if canonical == "mnist_math": + return get_mnist_math(seq_len=opts.get("seq_len", 100)) + if canonical == "continuous_localization_rotation": + return get_rot_dataset( + sequence_length=opts.get("sequence_length", 20), + ) + if canonical == "continuous_localization": + return get_trans_dataset( + sequence_length=opts.get("sequence_length", 20), + ) + raise ValueError(f"No loader for canonical dataset: {canonical}") + + +def slice_dataset( + dataset: Dict[str, Any], max_train_samples: Optional[int] +) -> Dict[str, Any]: + if max_train_samples is None: + return dataset + n = min(max_train_samples, dataset["size"]) + out = dict(dataset) + out["x_train"] = dataset["x_train"][:n] + out["y_train"] = dataset["y_train"][:n] + out["size"] = n + return out + + +def make_synthetic_dataset( + task: str, + *, + num_train: int = 8, + time: int = 16, + key: Optional[jax.Array] = None, +) -> Dict[str, Any]: + """Small in-memory dataset for smoke tests (no Hugging Face).""" + if key is None: + key = jax.random.PRNGKey(0) + k1, k2, k3, k4 = jax.random.split(key, 4) + if task == "classify": + feature_in, feature_out, num_labels = 1, 10, 10 + x = jax.random.uniform(k1, (num_train, time, feature_in)) + labels = jax.random.randint(k2, (num_train,), 0, num_labels) + y = jax.nn.one_hot(labels, num_labels) + elif task == "regress": + feature_in, feature_out = 3, 6 + num_labels = feature_out + x = jax.random.normal(k1, (num_train, time, feature_in)) * 0.1 + y = jax.random.normal(k2, (num_train, feature_out)) * 0.1 + else: + raise ValueError(f"Unknown synthetic task: {task}") + + x_test = jax.random.uniform(k3, (4, time, feature_in)) + if task == "classify": + test_labels = jax.random.randint(k4, (4,), 0, num_labels) + y_test = jax.nn.one_hot(test_labels, num_labels) + else: + y_test = jax.random.normal(k4, (4, feature_out)) * 0.1 + + return { + "x_train": x, + "y_train": y, + "x_test": x_test, + "y_test": y_test, + "num_labels": num_labels, + "size": num_train, + } diff --git a/memax/experiments/losses.py b/memax/experiments/losses.py new file mode 100644 index 0000000..765599b --- /dev/null +++ b/memax/experiments/losses.py @@ -0,0 +1,21 @@ +"""Loss function resolution for experiment scripts.""" + +from typing import Callable, Optional + +from memax.equinox.train_utils import loss_classify_terminal_output as eqx_loss_classify +from memax.equinox.train_utils import loss_regress_terminal_output as eqx_loss_regress +from memax.experiments.datasets import get_default_loss_name + + +def resolve_loss_name(dataset_name: str, loss_function: Optional[str]) -> str: + if loss_function is not None: + return loss_function + return get_default_loss_name(dataset_name) + + +def get_equinox_loss(loss_name: str) -> Callable: + if loss_name == "loss_classify_terminal_output": + return eqx_loss_classify + if loss_name == "loss_regress_terminal_output": + return eqx_loss_regress + raise ValueError(f"Unknown equinox loss: {loss_name}") diff --git a/memax/experiments/runner.py b/memax/experiments/runner.py new file mode 100644 index 0000000..7a0ef2a --- /dev/null +++ b/memax/experiments/runner.py @@ -0,0 +1,134 @@ +"""Training loops for equinox and linen experiment scripts.""" + +from typing import Any, Callable, Dict, Optional + +import equinox as eqx +import jax +import jax.numpy as jnp +import optax +import tqdm + +from memax.experiments.config import ExperimentConfig +from memax.linen.train_utils import update_model as linen_update_model + + +def _num_batches(dataset_size: int, batch_size: int) -> int: + return dataset_size // batch_size + + +def _update_indices( + dataset_size: int, batch_size: int, max_updates: Optional[int] +) -> range: + n = _num_batches(dataset_size, batch_size) + if max_updates is not None: + n = min(n, max_updates) + return range(n) + + +def run_equinox_training( + config: ExperimentConfig, + name: str, + model: eqx.Module, + dataset: Dict[str, Any], + loss_fn: Callable, + wandb_module=None, +) -> Dict[str, float]: + if config.use_wandb and wandb_module is not None: + wandb_module.init(project=config.project_name, name=name) + + lr_schedule = optax.constant_schedule(config.lr) + opt = optax.chain(optax.zero_nans(), optax.adamw(lr_schedule)) + opt_state = opt.init(eqx.filter(model, eqx.is_inexact_array)) + key = jax.random.PRNGKey(config.seed) + last_metrics: Dict[str, float] = {} + + for epoch in range(config.num_epochs): + key, shuffle_key = jax.random.split(key) + shuffle_idx = jax.random.permutation(shuffle_key, dataset["size"]) + x = dataset["x_train"][shuffle_idx] + y = dataset["y_train"][shuffle_idx] + updates = _update_indices( + dataset["size"], config.batch_size, config.max_updates + ) + pbar = tqdm.tqdm(updates, desc=f"{name} epoch {epoch}", leave=False) + + for update in pbar: + key, subkey = jax.random.split(key) + start = update * config.batch_size + end = start + config.batch_size + x_batch = x[start:end] + y_batch = y[start:end] + + from memax.equinox.train_utils import update_model + + model, opt_state, metrics = eqx.filter_jit(update_model)( + model=model, + loss_fn=loss_fn, + opt=opt, + opt_state=opt_state, + x=x_batch, + y=y_batch, + key=subkey, + ) + last_metrics = {k: float(jnp.mean(v)) for k, v in metrics.items()} + pbar.set_postfix({k: f"{v:.4f}" for k, v in last_metrics.items()}) + if config.use_wandb and wandb_module is not None: + wandb_module.log({**last_metrics, "epoch": epoch}) + + if config.use_wandb and wandb_module is not None: + wandb_module.finish() + return last_metrics + + +def run_linen_training( + config: ExperimentConfig, + name: str, + model, + dataset: Dict[str, Any], + loss_fn: Callable, + wandb_module=None, +) -> Dict[str, float]: + if config.use_wandb and wandb_module is not None: + wandb_module.init(project=config.project_name, name=name) + + lr_schedule = optax.constant_schedule(config.lr) + opt = optax.chain(optax.zero_nans(), optax.adamw(lr_schedule)) + key = jax.random.PRNGKey(config.seed) + + dummy_x = dataset["x_train"][0] + dummy_starts = jnp.zeros(dummy_x.shape[0], dtype=bool) + dummy_h = model.zero_carry() + params = model.init(key, dummy_h, (dummy_x, dummy_starts)) + opt_state = opt.init(params) + last_metrics: Dict[str, float] = {} + + jitted_update = jax.jit(linen_update_model, static_argnames=("loss_fn", "opt")) + + for epoch in range(config.num_epochs): + key, shuffle_key = jax.random.split(key) + shuffle_idx = jax.random.permutation(shuffle_key, dataset["size"]) + x = dataset["x_train"][shuffle_idx] + y = dataset["y_train"][shuffle_idx] + updates = _update_indices( + dataset["size"], config.batch_size, config.max_updates + ) + pbar = tqdm.tqdm(updates, desc=f"{name} epoch {epoch}", leave=False) + + for update in pbar: + key, subkey = jax.random.split(key) + start = update * config.batch_size + end = start + config.batch_size + x_batch = x[start:end] + y_batch = y[start:end] + + params, opt_state, metrics = jitted_update( + params, loss_fn, opt, opt_state, x_batch, y_batch, key=subkey + ) + last_metrics = {k: float(jnp.mean(v)) for k, v in metrics.items()} + pbar.set_postfix({k: f"{v:.4f}" for k, v in last_metrics.items()}) + if config.use_wandb and wandb_module is not None: + wandb_module.log({**last_metrics, "epoch": epoch}) + + if config.use_wandb and wandb_module is not None: + wandb_module.finish() + return last_metrics diff --git a/memax/linen/train_utils.py b/memax/linen/train_utils.py index 74d595e..ca3df61 100644 --- a/memax/linen/train_utils.py +++ b/memax/linen/train_utils.py @@ -2,8 +2,6 @@ It includes loss functions, accuracy metrics, and training loops. It also provides a straightforward way to construct multi-layer recurrent models.""" -from typing import Any - import jax import jax.numpy as jnp import optax @@ -67,7 +65,6 @@ def loss_classify_terminal_output( seq_len = x.shape[1] starts = jnp.zeros((batch_size, seq_len), dtype=bool) h0 = init_carry_fn(params) - # h0 = jax.tree_map(partial(add_batch_dim, batch_size=batch_size), h0) h0 = add_batch_dim(h0, batch_size) _, y_preds = jax.vmap(model_apply_fn, in_axes=[None, 0, 0])(params, h0, (x, starts)) @@ -78,6 +75,60 @@ def loss_classify_terminal_output( return loss, {"loss": loss, "accuracy": acc} +def mse( + y_hat: Shaped[Array, "Batch ... Feature"], y: Shaped[Array, "Batch ... Feature"] +) -> Shaped[Array, "1"]: + return jnp.mean(jnp.linalg.norm(y - y_hat, axis=-1, ord=2)) + + +def l1_error( + y_hat: Shaped[Array, "Batch ... Feature"], y: Shaped[Array, "Batch ... Feature"] +) -> Shaped[Array, "1"]: + return jnp.mean(jnp.linalg.norm(y - y_hat, axis=-1, ord=1)) + + +def loss_regress_terminal_output( + params: FrozenDict, + x: Shaped[Array, "Batch Time Feature"], + y: Shaped[Array, "Batch Feature"], + init_carry_fn, + model_apply_fn, +) -> Tuple[Shaped[Array, "1"], Dict[str, Array]]: + batch_size = x.shape[0] + seq_len = x.shape[1] + starts = jnp.zeros((batch_size, seq_len), dtype=bool) + h0 = init_carry_fn(params) + h0 = add_batch_dim(h0, batch_size) + + _, y_preds = jax.vmap(model_apply_fn, in_axes=[None, 0, 0])(params, h0, (x, starts)) + y_pred = y_preds[:, -1] + + loss = mse(y_pred, y) + l1 = l1_error(y_pred, y) + return loss, {"loss": loss, "l1_error": l1} + + +def make_linen_loss_fn(model, loss_name: str): + """Bind a terminal loss to a linen ResidualModel.""" + from functools import partial + + init_carry_fn = partial(model.apply, method="initialize_carry") + model_apply_fn = model.apply + if loss_name == "loss_classify_terminal_output": + return partial( + loss_classify_terminal_output, + init_carry_fn=init_carry_fn, + model_apply_fn=model_apply_fn, + ) + if loss_name == "loss_regress_terminal_output": + return partial( + loss_regress_terminal_output, + init_carry_fn=init_carry_fn, + model_apply_fn=model_apply_fn, + ) + raise ValueError(f"Unknown loss: {loss_name}") + + def get_semigroups( recurrent_size: int, semigroup_kwargs: Optional[Dict[str, Any]] = None, diff --git a/pyproject.toml b/pyproject.toml index 25d0eea..15cb766 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,12 @@ all = [ ] test = ["pytest"] +[tool.pytest.ini_options] +markers = [ + "integration: downloads Hugging Face datasets (not run in default CI)", +] +addopts = "-m 'not integration'" + [build-system] requires = ["setuptools>=61", "wheel"] build-backend = "setuptools.build_meta" diff --git a/run_equinox_experiments.py b/run_equinox_experiments.py index 3d43d5c..81b9ab0 100644 --- a/run_equinox_experiments.py +++ b/run_equinox_experiments.py @@ -1,208 +1,50 @@ -"""This script runs experiments training various recurrent memory models -on different datasets using Equinox. It serves as a reference implementation -for training and evaluating memax modules.""" +"""Train recurrent memory models (Equinox) on memax datasets.""" import argparse -import equinox as eqx import jax -import jax.numpy as jnp -import optax -import tqdm - import wandb -from memax.datasets.mnist_math import get_dataset as get_mnist_math -from memax.datasets.sequential_mnist import get_dataset as get_sequential_mnist -from memax.datasets.continuous_localization import get_rot_dataset, get_trans_dataset -from memax.equinox.train_utils import ( - get_residual_memory_models, - loss_classify_terminal_output, - loss_regress_terminal_output, - update_model, -) - - -def parse_args(): - parser = argparse.ArgumentParser(description="Train recurrent memory models.") - parser.add_argument("--seed", type=int, default=0, help="Random seed") - parser.add_argument("--num-epochs", type=int, help="Number of training epochs") - parser.add_argument("--batch-size", type=int, help="Batch size") - parser.add_argument( - "--recurrent_size", type=int, help="Recurrent size of the model" - ) - parser.add_argument("--num_layers", type=int, help="Number of layers in the model") - parser.add_argument("--lr", type=float, help="Learning rate") - parser.add_argument( - "--use-wandb", - action="store_true", - default=False, - help="Use Weights & Biases for logging", - ) - parser.add_argument( - "--project-name", - type=str, - default="memax-debug", - help="Weights & Biases project name", - ) - parser.add_argument( - "--dataset-name", - type=str, - default="sequential_mnist", - help="Dataset name (e.g., mnist_math, other_dataset)", - ) - parser.add_argument( - "--loss-function", - type=str, - default=None, - help="Loss function to use (e.g., loss_classify_terminal_output, other_loss_fn)", - ) - parser.add_argument("--models", type=str, nargs="+", default="all") - return parser.parse_args() - - -def get_default_loss(dataset_name): - defaults = { - "sequential_mnist": "loss_classify_terminal_output", - "mnist_math_5": "loss_classify_terminal_output", - "sequential_rotation": "loss_regress_terminal_output" - } - if dataset_name in defaults: - return defaults[dataset_name] - else: - raise ValueError( - f"No default loss function defined for dataset: {dataset_name}" - ) - -def get_default_hyperparameters(dataset_name): - defaults = { - "mnist_math_5": { - "num_epochs": 5, - "batch_size": 16, - "recurrent_size": 256, - "num_layers": 2, - "lr": 0.0001, - }, - "sequential_mnist": { - "num_epochs": 5, - "batch_size": 16, - "recurrent_size": 256, - "num_layers": 2, - "lr": 0.0001, - }, - "sequential_rotation": { - "num_epochs": 5, - "batch_size": 16, - "recurrent_size": 256, - "num_layers": 2, - "lr": 0.0001, - }, - # Add more datasets and their default hyperparameters here - } - if dataset_name in defaults: - return defaults[dataset_name] - else: - raise ValueError( - f"No default hyperparameters defined for dataset: {dataset_name}" - ) - -def update_config_with_defaults(args): - defaults = get_default_hyperparameters(args.dataset_name) - for key, value in defaults.items(): - if getattr(args, key) is None: - setattr(args, key, value) - - -def run_test(config, name, model, dataset, loss_fn): - if config.use_wandb: - wandb.init(project=config.project_name, name=name) - - lr_schedule = optax.constant_schedule(config.lr) - opt = optax.chain( - optax.zero_nans(), - optax.adamw(lr_schedule), - ) - opt_state = opt.init(eqx.filter(model, eqx.is_inexact_array)) - key = jax.random.PRNGKey(config.seed) - - for epoch in range(config.num_epochs): - key, shuffle_key = jax.random.split(key) - shuffle_idx = jax.random.permutation(shuffle_key, dataset["size"]) - x = dataset["x_train"][shuffle_idx] - y = dataset["y_train"][shuffle_idx] - pbar = tqdm.tqdm(range(x.shape[0] // config.batch_size)) - - for update in pbar: - key, subkey = jax.random.split(key) - x_batch = x[update * config.batch_size : (update + 1) * config.batch_size] - y_batch = y[update * config.batch_size : (update + 1) * config.batch_size] - - model, opt_state, metrics = eqx.filter_jit(update_model)( - model=model, - loss_fn=loss_fn, - opt=opt, - opt_state=opt_state, - x=x_batch, - y=y_batch, - key=subkey, - ) - - mean_metrics = {k: jnp.mean(v).item() for k, v in metrics.items()} - pbar.set_description( - f"{name} epoch: {epoch}, " - + ", ".join(f"{k}: {v:.4f}" for k, v in mean_metrics.items()) - ) - if config.use_wandb: - wandb.log({**mean_metrics, "epoch": epoch}) - - if config.use_wandb: - wandb.finish() +from memax.equinox.train_utils import get_residual_memory_models +from memax.experiments.cli import add_experiment_args +from memax.experiments.config import prepare_config +from memax.experiments.datasets import load_dataset, slice_dataset +from memax.experiments.losses import get_equinox_loss, resolve_loss_name +from memax.experiments.runner import run_equinox_training def main(): - args = parse_args() - update_config_with_defaults(args) + parser = argparse.ArgumentParser( + description="Train recurrent memory models (Equinox)." + ) + add_experiment_args(parser) + args = parser.parse_args() + config = prepare_config(args) - # Dynamically load dataset - if args.dataset_name == "mnist_math_5": - dataset = get_mnist_math(5) - elif args.dataset_name == "sequential_mnist": - dataset = get_sequential_mnist() - elif args.dataset_name == "sequential_rotation": - dataset = get_rot_dataset() - else: - raise ValueError(f"Unknown dataset: {args.dataset_name}") + dataset = load_dataset(config.dataset_name) + dataset = slice_dataset(dataset, config.max_train_samples) feature_in = dataset["x_test"].shape[-1] - feature_out = dataset["y_test"].shape[-1] + feature_out = dataset["y_test"].shape[-1] - # Select loss function - if args.loss_function is None: - loss_fn_name = get_default_loss(args.dataset_name) - else: - loss_fn_name = args.loss_fn + loss_name = resolve_loss_name(config.dataset_name, config.loss_function) + loss_fn = get_equinox_loss(loss_name) - if loss_fn_name =="loss_classify_terminal_output": - loss_fn = loss_classify_terminal_output - elif loss_fn_name == "loss_regress_terminal_output": - loss_fn = loss_regress_terminal_output - else: - raise ValueError(f"Unknown loss function: {args.loss_function}") - - # Create model - key = jax.random.PRNGKey(args.seed) + key = jax.random.PRNGKey(config.seed) models = get_residual_memory_models( input=feature_in, - hidden=args.recurrent_size, + hidden=config.recurrent_size, output=feature_out, - num_layers=args.num_layers, - models=args.models, + num_layers=config.num_layers, + models=config.models, key=key, ) - # Run experiments + wandb_module = wandb if config.use_wandb else None for name, model in models.items(): - run_test(args, name, model, dataset, loss_fn) + run_equinox_training( + config, name, model, dataset, loss_fn, wandb_module=wandb_module + ) if __name__ == "__main__": diff --git a/run_linen_experiments.py b/run_linen_experiments.py index edf8ab6..8b77840 100644 --- a/run_linen_experiments.py +++ b/run_linen_experiments.py @@ -1,174 +1,47 @@ -"""This script runs experiments training various recurrent memory models -on different datasets using Flax Linen. It serves as a reference implementation -for training and evaluating memax modules.""" +"""Train recurrent memory models (Flax Linen) on memax datasets.""" + import argparse -from functools import partial -import equinox as eqx import jax -import jax.numpy as jnp -import optax -import tqdm - import wandb -from memax.datasets.mnist_math import get_dataset as get_mnist_math -from memax.datasets.sequential_mnist import get_dataset as get_sequential_mnist -from memax.linen.train_utils import ( - get_residual_memory_models, - loss_classify_terminal_output, - update_model, -) - - -def parse_args(): - parser = argparse.ArgumentParser(description="Train recurrent memory models.") - parser.add_argument("--seed", type=int, default=0, help="Random seed") - parser.add_argument("--num-epochs", type=int, help="Number of training epochs") - parser.add_argument("--batch-size", type=int, help="Batch size") - parser.add_argument( - "--recurrent-size", type=int, help="Recurrent size of the model" - ) - parser.add_argument("--num-layers", type=int, help="Number of layers in the model") - parser.add_argument("--lr", type=float, help="Learning rate") - parser.add_argument( - "--use-wandb", - action="store_true", - default=False, - help="Use Weights & Biases for logging", - ) - parser.add_argument( - "--project-name", - type=str, - default="memax-debug", - help="Weights & Biases project name", - ) - parser.add_argument( - "--dataset_name", - type=str, - default="sequential_mnist", - help="Dataset name (e.g., mnist_math, other_dataset)", - ) - parser.add_argument( - "--loss-function", - type=str, - default="loss_classify_terminal_output", - help="Loss function to use (e.g., loss_classify_terminal_output, other_loss_fn)", - ) - parser.add_argument("--models", type=str, nargs="+") - return parser.parse_args() - - -def get_default_hyperparameters(dataset_name): - defaults = { - "mnist_math_5": { - "num_epochs": 5, - "batch_size": 16, - "recurrent_size": 256, - "num_layers": 2, - "lr": 0.0001, - }, - "sequential_mnist": { - "num_epochs": 5, - "batch_size": 16, - "recurrent_size": 256, - "num_layers": 2, - "lr": 0.0001, - }, - # Add more datasets and their default hyperparameters here - } - if dataset_name in defaults: - return defaults[dataset_name] - else: - raise ValueError( - f"No default hyperparameters defined for dataset: {dataset_name}" - ) - -def update_config_with_defaults(args): - defaults = get_default_hyperparameters(args.dataset_name) - for key, value in defaults.items(): - if getattr(args, key) is None: - setattr(args, key, value) - - -def run_test(config, name, model, dataset, loss_fn): - if config.use_wandb: - wandb.init(project=config.project_name, name=name) - - lr_schedule = optax.constant_schedule(config.lr) - opt = optax.chain( - optax.zero_nans(), - optax.adamw(lr_schedule), - ) - key = jax.random.PRNGKey(config.seed) - - dummy_x = dataset["x_train"][0] - dummy_starts = jnp.zeros(dummy_x.shape[0], dtype=bool) - dummy_h = model.zero_carry() - params = model.init(key, dummy_h, (dummy_x, dummy_starts)) - opt_state = opt.init(params) - initialise_carry_fn = partial(model.apply, method="initialize_carry") - model_apply_fn = model.apply - loss_fn = partial( - loss_classify_terminal_output, - init_carry_fn=initialise_carry_fn, - model_apply_fn=model_apply_fn, - ) - - for epoch in range(config.num_epochs): - key, shuffle_key = jax.random.split(key) - shuffle_idx = jax.random.permutation(shuffle_key, dataset["size"]) - x = dataset["x_train"][shuffle_idx] - y = dataset["y_train"][shuffle_idx] - pbar = tqdm.tqdm(range(x.shape[0] // config.batch_size)) - - for update in pbar: - key, subkey = jax.random.split(key) - x_batch = x[update * config.batch_size : (update + 1) * config.batch_size] - y_batch = y[update * config.batch_size : (update + 1) * config.batch_size] - - params, opt_state, metrics = jax.jit( - update_model, static_argnames=("loss_fn", "opt") - )(params, loss_fn, opt, opt_state, x_batch, y_batch, key=subkey) - - mean_metrics = {k: jnp.mean(v).item() for k, v in metrics.items()} - pbar.set_description( - f"{name} epoch: {epoch}, " - + ", ".join(f"{k}: {v:.4f}" for k, v in mean_metrics.items()) - ) - if config.use_wandb: - wandb.log({**mean_metrics, "epoch": epoch}) - - if config.use_wandb: - wandb.finish() +from memax.experiments.cli import add_experiment_args +from memax.experiments.config import prepare_config +from memax.experiments.datasets import load_dataset, slice_dataset +from memax.experiments.losses import resolve_loss_name +from memax.experiments.runner import run_linen_training +from memax.linen.train_utils import get_residual_memory_models, make_linen_loss_fn def main(): - args = parse_args() - update_config_with_defaults(args) + parser = argparse.ArgumentParser( + description="Train recurrent memory models (Linen)." + ) + add_experiment_args(parser) + args = parser.parse_args() + config = prepare_config(args) - # Dynamically load dataset - if args.dataset_name == "mnist_math_5": - dataset = get_mnist_math(5) - elif args.dataset_name == "sequential_mnist": - dataset = get_sequential_mnist() - else: - raise ValueError(f"Unknown dataset: {args.dataset_name}") + dataset = load_dataset(config.dataset_name) + dataset = slice_dataset(dataset, config.max_train_samples) - # Dynamically select loss function - if args.loss_function == "loss_classify_terminal_output": - loss_fn = loss_classify_terminal_output - else: - raise ValueError(f"Unknown loss function: {args.loss_function}") + feature_in = dataset["x_test"].shape[-1] + feature_out = dataset["y_test"].shape[-1] + loss_name = resolve_loss_name(config.dataset_name, config.loss_function) models = get_residual_memory_models( - hidden=args.recurrent_size, - output=dataset["num_labels"], - num_layers=args.num_layers, + hidden=config.recurrent_size, + output=feature_out, + num_layers=config.num_layers, + models=config.models, + model_kwargs={"input": feature_in}, ) + wandb_module = wandb if config.use_wandb else None for name, model in models.items(): - run_test(args, name, model, dataset, loss_fn) + loss_fn = make_linen_loss_fn(model, loss_name) + run_linen_training( + config, name, model, dataset, loss_fn, wandb_module=wandb_module + ) if __name__ == "__main__":