diff --git a/docs/api/benchmarks.rst b/docs/api/benchmarks.rst index 2e3b034..96469da 100644 --- a/docs/api/benchmarks.rst +++ b/docs/api/benchmarks.rst @@ -1,7 +1,21 @@ Benchmarks Reference ========================= -This page documents all the benchmarks included in the `torchrecurrent.benchmarks` module. +Synthetic benchmarks for evaluating recurrent architectures. The generators use +batch-first tensors, support reproducible :class:`torch.Generator` instances, and +can return either raw tensors or :class:`torch.utils.data.DataLoader` objects. + +``adding_problem`` implements the two-half marker sampling from the canonical +adding task. ``copy_memory`` implements the categorical ``T + 20`` protocol by +default and can one-hot encode inputs for direct use with recurrent layers. +``sequential_mnist`` adapts standard MNIST tensors to the sequential and fixed +permutation variants without introducing a dataset-download dependency. +``sequential_cifar10`` adapts CIFAR-10 tensors the same way, flattening each +image into a 1024-step, 3-channel pixel sequence. ``penn_treebank`` prepares +licensed, preprocessed PTB split files for canonical word-level language +modeling with contiguous truncated-BPTT batches. ``timit`` batches aligned +120-dimensional log-Mel, delta, and acceleration features for the canonical +180-state frame-classification protocol. .. autosummary:: :toctree: ../generated/ @@ -9,3 +23,7 @@ This page documents all the benchmarks included in the `torchrecurrent.benchmark torchrecurrent.benchmarks.adding_problem torchrecurrent.benchmarks.copy_memory + torchrecurrent.benchmarks.penn_treebank + torchrecurrent.benchmarks.sequential_cifar10 + torchrecurrent.benchmarks.sequential_mnist + torchrecurrent.benchmarks.timit diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py new file mode 100644 index 0000000..c817fca --- /dev/null +++ b/tests/test_benchmarks.py @@ -0,0 +1,293 @@ +import pytest +import torch + +from torchrecurrent.benchmarks import ( + adding_problem, + copy_memory, + penn_treebank, + sequential_cifar10, + sequential_mnist, + timit, +) + + +def test_adding_problem_matches_definition(): + inputs, targets = adding_problem( + 9, + 32, + return_dataloader=False, + generator=torch.Generator().manual_seed(0), + dtype=torch.float64, + ) + + assert inputs.shape == (32, 9, 2) + assert targets.shape == (32, 1) + assert inputs.dtype == targets.dtype == torch.float64 + assert torch.all(inputs[:, :4, 1].sum(dim=1) == 1) + assert torch.all(inputs[:, 4:, 1].sum(dim=1) == 1) + torch.testing.assert_close(targets[:, 0], (inputs[:, :, 0] * inputs[:, :, 1]).sum(1)) + + +def test_adding_problem_dataloader(): + loader = adding_problem(8, 7, batch_size=4, shuffle=False) + inputs, targets = next(iter(loader)) + + assert inputs.shape == (4, 8, 2) + assert targets.shape == (4, 1) + + +@pytest.mark.parametrize("one_hot", [False, True]) +def test_copy_memory_matches_definition(one_hot): + inputs, targets = copy_memory( + 5, + 4, + memory_length=3, + return_dataloader=False, + one_hot=one_hot, + generator=torch.Generator().manual_seed(0), + ) + token_inputs = inputs.argmax(-1) if one_hot else inputs + + assert token_inputs.shape == (4, 11) + assert targets.shape == (4, 11) + assert torch.all(token_inputs[:, :3] < 8) + assert torch.all(token_inputs[:, 3:7] == 8) + assert torch.all(token_inputs[:, 7] == 9) + assert torch.all(token_inputs[:, 8:] == 8) + assert torch.all(targets[:, :8] == 8) + assert torch.equal(targets[:, 8:], token_inputs[:, :3]) + if one_hot: + assert inputs.shape == (4, 11, 10) + assert inputs.is_floating_point() + + +@pytest.mark.parametrize( + "function,args", + [ + (adding_problem, (1, 2)), + (adding_problem, (2, 0)), + (copy_memory, (0, 2)), + (copy_memory, (2, 0)), + ], +) +def test_benchmarks_validate_sizes(function, args): + with pytest.raises(ValueError): + function(*args) + + +def test_sequential_mnist_normalizes_and_flattens_images(): + images = (torch.arange(2 * 28 * 28) % 256).to(torch.uint8).reshape(2, 1, 28, 28) + targets = torch.tensor([3, 7], dtype=torch.int32) + + sequences, result_targets = sequential_mnist( + images, targets, return_dataloader=False, dtype=torch.float64 + ) + + assert sequences.shape == (2, 784, 1) + assert sequences.dtype == torch.float64 + assert result_targets.dtype == torch.long + torch.testing.assert_close(sequences[0, :, 0], images[0].flatten().double() / 255) + assert torch.equal(result_targets, targets.long()) + + +def test_sequential_mnist_applies_one_fixed_permutation(): + images = torch.arange(2 * 28 * 28).reshape(2, 28, 28).float() + targets = torch.tensor([0, 1]) + permutation = torch.arange(783, -1, -1) + + sequences, _ = sequential_mnist( + images, + targets, + permutation=permutation, + return_dataloader=False, + normalize=False, + ) + + assert torch.equal(sequences[:, :, 0], images.flatten(1)[:, permutation]) + + +def test_sequential_mnist_dataloader_is_layer_ready(): + loader = sequential_mnist( + torch.zeros(5, 28, 28, dtype=torch.uint8), + torch.arange(5), + batch_size=3, + shuffle=False, + ) + sequences, targets = next(iter(loader)) + + assert sequences.shape == (3, 784, 1) + assert targets.shape == (3,) + + +@pytest.mark.parametrize( + "images,targets,permutation", + [ + (torch.zeros(2, 27, 28), torch.zeros(2), None), + (torch.zeros(2, 3, 28, 28), torch.zeros(2), None), + (torch.zeros(2, 28, 28), torch.zeros(3), None), + (torch.zeros(2, 28, 28), torch.zeros(2), torch.zeros(783)), + (torch.zeros(2, 28, 28), torch.zeros(2), torch.zeros(784)), + ], +) +def test_sequential_mnist_validates_inputs(images, targets, permutation): + with pytest.raises(ValueError): + sequential_mnist(images, targets, permutation=permutation) + + +def test_sequential_cifar10_normalizes_and_flattens_images(): + images = (torch.arange(2 * 32 * 32 * 3) % 256).to(torch.uint8).reshape(2, 32, 32, 3) + targets = torch.tensor([3, 7], dtype=torch.int32) + + sequences, result_targets = sequential_cifar10( + images, targets, return_dataloader=False, dtype=torch.float64 + ) + + assert sequences.shape == (2, 1024, 3) + assert sequences.dtype == torch.float64 + assert result_targets.dtype == torch.long + torch.testing.assert_close(sequences[0], images[0].reshape(1024, 3).double() / 255) + assert torch.equal(result_targets, targets.long()) + + +def test_sequential_cifar10_accepts_channels_first_images(): + images = torch.arange(2 * 3 * 32 * 32).reshape(2, 3, 32, 32).float() + targets = torch.tensor([0, 1]) + + sequences, _ = sequential_cifar10( + images, targets, return_dataloader=False, normalize=False + ) + + expected = images.permute(0, 2, 3, 1).reshape(2, 1024, 3) + torch.testing.assert_close(sequences, expected) + + +def test_sequential_cifar10_applies_one_fixed_permutation(): + images = torch.arange(2 * 32 * 32 * 3).reshape(2, 32, 32, 3).float() + targets = torch.tensor([0, 1]) + permutation = torch.arange(1023, -1, -1) + + sequences, _ = sequential_cifar10( + images, + targets, + permutation=permutation, + return_dataloader=False, + normalize=False, + ) + + expected = images.reshape(2, 1024, 3)[:, permutation] + torch.testing.assert_close(sequences, expected) + + +def test_sequential_cifar10_dataloader_is_layer_ready(): + loader = sequential_cifar10( + torch.zeros(5, 32, 32, 3, dtype=torch.uint8), + torch.arange(5), + batch_size=3, + shuffle=False, + ) + sequences, targets = next(iter(loader)) + + assert sequences.shape == (3, 1024, 3) + assert targets.shape == (3,) + + +@pytest.mark.parametrize( + "images,targets,permutation", + [ + (torch.zeros(2, 31, 32, 3), torch.zeros(2), None), + (torch.zeros(2, 32, 32, 4), torch.zeros(2), None), + (torch.zeros(2, 32, 32, 3), torch.zeros(3), None), + (torch.zeros(2, 32, 32, 3), torch.zeros(2), torch.zeros(1023)), + (torch.zeros(2, 32, 32, 3), torch.zeros(2), torch.zeros(1024)), + ], +) +def test_sequential_cifar10_validates_inputs(images, targets, permutation): + with pytest.raises(ValueError): + sequential_cifar10(images, targets, permutation=permutation) + + +def test_penn_treebank_builds_vocabulary_and_shifted_streams(tmp_path): + train = tmp_path / "ptb.train.txt" + validation = tmp_path / "ptb.valid.txt" + test = tmp_path / "ptb.test.txt" + train.write_text("the cat sat\nthe dog sat\n", encoding="utf-8") + validation.write_text("the fox sat\n", encoding="utf-8") + test.write_text("the cat ran\n", encoding="utf-8") + + corpus = penn_treebank(train, validation, test, batch_size=2, sequence_length=2) + + assert corpus.vocabulary[""] == 0 + assert "" in corpus.vocabulary + inputs, targets = next(iter(corpus.train)) + assert inputs.shape == targets.shape == (2, 2) + stream = corpus.train.dataset.tokens + assert torch.equal(inputs, stream[:, :2]) + assert torch.equal(targets, stream[:, 1:3]) + + assert corpus.vocabulary[""] in corpus.validation.dataset.tokens + + +def test_penn_treebank_caps_vocabulary_deterministically(tmp_path): + files = [] + for name in ("train", "valid", "test"): + path = tmp_path / name + path.write_text("z b a b a z c\n", encoding="utf-8") + files.append(path) + + corpus = penn_treebank(*files, batch_size=1, max_vocab_size=4) + + assert list(corpus.vocabulary) == ["", "", "a", "b"] + + +def test_penn_treebank_rejects_too_short_splits(tmp_path): + files = [] + for name in ("train", "valid", "test"): + path = tmp_path / name + path.write_text("token\n", encoding="utf-8") + files.append(path) + + with pytest.raises(ValueError, match="more tokens"): + penn_treebank(*files, batch_size=2) + + +def test_timit_pads_variable_length_utterances(): + features = [torch.ones(3, 120), torch.full((5, 120), 2.0)] + targets = [torch.tensor([1, 2, 3]), torch.tensor([4, 5, 6, 7, 8])] + + batch = next(iter(timit(features, targets, batch_size=2, shuffle=False))) + + assert batch.features.shape == (2, 5, 120) + assert batch.targets.shape == (2, 5) + assert torch.equal(batch.lengths, torch.tensor([3, 5])) + assert torch.all(batch.features[0, 3:] == 0) + assert torch.all(batch.targets[0, 3:] == -100) + assert torch.equal(batch.targets[1], targets[1]) + + +def test_timit_supports_alternative_feature_and_class_counts(): + loader = timit( + [torch.ones(2, 4)], + [torch.tensor([0, 2])], + feature_size=4, + num_classes=3, + batch_size=1, + ) + + batch = next(iter(loader)) + assert batch.features.shape == (1, 2, 4) + + +@pytest.mark.parametrize( + "features,targets,error", + [ + ([], [], ValueError), + ([torch.ones(2, 119)], [torch.zeros(2)], ValueError), + ([torch.ones(2, 120)], [torch.zeros(3)], ValueError), + ([torch.ones(0, 120)], [torch.zeros(0)], ValueError), + ([torch.ones(2, 120)], [torch.tensor([0, 180])], ValueError), + ([torch.ones(2, 120, dtype=torch.long)], [torch.zeros(2)], TypeError), + ], +) +def test_timit_validates_aligned_sequences(features, targets, error): + with pytest.raises(error): + timit(features, targets) diff --git a/torchrecurrent/benchmarks/__init__.py b/torchrecurrent/benchmarks/__init__.py index 788c44c..445cde8 100644 --- a/torchrecurrent/benchmarks/__init__.py +++ b/torchrecurrent/benchmarks/__init__.py @@ -1,4 +1,18 @@ from .adding import adding_problem from .copymemory import copy_memory +from .penn_treebank import PennTreebankCorpus, penn_treebank +from .sequential_cifar10 import sequential_cifar10 +from .sequential_mnist import sequential_mnist +from .timit import TIMITBatch, TIMITDataset, timit -__all__ = ["adding_problem", "copy_memory"] +__all__ = [ + "PennTreebankCorpus", + "TIMITBatch", + "TIMITDataset", + "adding_problem", + "copy_memory", + "penn_treebank", + "sequential_cifar10", + "sequential_mnist", + "timit", +] diff --git a/torchrecurrent/benchmarks/adding.py b/torchrecurrent/benchmarks/adding.py index 54a026e..2d14dd0 100644 --- a/torchrecurrent/benchmarks/adding.py +++ b/torchrecurrent/benchmarks/adding.py @@ -1,5 +1,5 @@ import torch -from torch.utils.data import TensorDataset, DataLoader +from torch.utils.data import DataLoader, TensorDataset def adding_problem( @@ -7,56 +7,91 @@ def adding_problem( n_samples: int, return_dataloader: bool = True, batch_size: int = 64, - shuffle=True, + shuffle: bool = True, + *, + generator: torch.Generator = None, + dtype: torch.dtype = None, + device: torch.device = None, + **dataloader_kwargs, ): - """Generate data for the adding problem benchmark. - - The adding problem is a synthetic task where each input sequence - consists of two features per time step: - - 1. A random number sampled uniformly from [0, 1]. - 2. A binary mask indicating which two positions in the sequence - should be summed. - - The target is the sum of the two masked numbers. - - Parameters - ---------- - sequence_length : int - Length of each input sequence. - n_samples : int - Number of samples to generate. - return_dataloader : bool, default=True - If True, return a DataLoader wrapping the dataset. - If False, return raw tensors instead. - batch_size : int, default=64 - Batch size used when returning a DataLoader. - shuffle : bool, default=True - Whether to shuffle the dataset when returning a DataLoader. - - Returns - ------- - torch.utils.data.DataLoader or tuple of (torch.Tensor, torch.Tensor) - - If ``return_dataloader`` is True: a DataLoader yielding batches of - (inputs, targets). - - If ``return_dataloader`` is False: - * inputs: torch.Tensor of shape (n_samples, sequence_length, 2) - * targets: torch.Tensor of shape (n_samples, 1) + """Generate the adding problem introduced for long-term memory tests. + + Each input contains a uniform random sequence and a binary indicator + sequence. One indicator is sampled from each half of the sequence, and the + regression target is the sum of the two indicated values. This is the + formulation used by Arjovsky et al. (2016), Section 5.2 + (https://proceedings.mlr.press/v48/arjovsky16.html). + + Args: + sequence_length: Number of time steps. Must be at least 2. + n_samples: Number of independent sequences to generate. + return_dataloader: Return a data loader when true, otherwise tensors. + batch_size: Batch size of the returned data loader. + shuffle: Whether the returned data loader shuffles samples. + generator: Optional random number generator used for reproducibility. + Only forwarded to the returned data loader's shuffling when it is a + CPU generator; a non-CPU generator is still used for tensor + generation but the loader falls back to its own seeding. + dtype: Floating-point dtype of inputs and targets. + device: Device on which to create the tensors. + **dataloader_kwargs: Additional arguments passed to + :class:`torch.utils.data.DataLoader`. + + Returns: + A data loader, or ``(inputs, targets)`` when ``return_dataloader=False``. + Inputs have shape ``(n_samples, sequence_length, 2)`` and targets have + shape ``(n_samples, 1)``. """ - random_sequence = torch.rand(n_samples, sequence_length, 1) - mask_sequence = torch.zeros(n_samples, sequence_length, 1) - targets = torch.zeros(n_samples, 1) - - for i in range(n_samples): - idx = torch.randperm(sequence_length)[:2] - mask_sequence[i, idx, 0] = 1 - targets[i] = random_sequence[i, idx, 0].sum() - - inputs = torch.cat((random_sequence, mask_sequence), dim=2) - if return_dataloader: - dataset = TensorDataset(inputs, targets) - data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle) - - return data_loader - else: + if sequence_length < 2: + raise ValueError("sequence_length must be at least 2") + if n_samples < 1: + raise ValueError("n_samples must be positive") + if dtype is None: + dtype = torch.get_default_dtype() + if not dtype.is_floating_point: + raise TypeError("dtype must be a floating-point dtype") + + random_sequence = torch.rand( + n_samples, + sequence_length, + 1, + generator=generator, + dtype=dtype, + device=device, + ) + first_half = sequence_length // 2 + first_indices = torch.randint( + first_half, (n_samples, 1), generator=generator, device=device + ) + second_indices = torch.randint( + first_half, + sequence_length, + (n_samples, 1), + generator=generator, + device=device, + ) + indices = torch.cat((first_indices, second_indices), dim=1) + + mask_sequence = torch.zeros_like(random_sequence) + mask_sequence.scatter_(1, indices.unsqueeze(-1), 1) + targets = random_sequence.squeeze(-1).gather(1, indices).sum(dim=1, keepdim=True) + inputs = torch.cat((random_sequence, mask_sequence), dim=-1) + + if not return_dataloader: return inputs, targets + + # torch.utils.data.DataLoader shuffling always runs its generator on CPU + # (torch.randperm has no device argument), so a generator created for a + # non-CPU generation device cannot also drive the loader's shuffling. + loader_generator = generator + if generator is not None and generator.device.type != "cpu": + loader_generator = None + + dataset = TensorDataset(inputs, targets) + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + generator=loader_generator, + **dataloader_kwargs, + ) diff --git a/torchrecurrent/benchmarks/copymemory.py b/torchrecurrent/benchmarks/copymemory.py index ea7c557..e790a5d 100644 --- a/torchrecurrent/benchmarks/copymemory.py +++ b/torchrecurrent/benchmarks/copymemory.py @@ -1,50 +1,103 @@ import torch -from torch.utils.data import TensorDataset, DataLoader +from torch.nn import functional as F +from torch.utils.data import DataLoader, TensorDataset -def copy_memory(seq_len: int, n_samples: int, num_classes: int = 10, **kwargs): - """Generate data for the copy memory benchmark. +def copy_memory( + seq_len: int, + n_samples: int, + num_classes: int = 10, + *, + memory_length: int = 10, + return_dataloader: bool = True, + one_hot: bool = False, + generator: torch.Generator = None, + device: torch.device = None, + **dataloader_kwargs, +): + """Generate the canonical copy-memory benchmark. - The copy memory task is a synthetic sequence learning problem where a - model must memorize and reproduce an input sequence after a long delay. - Each sample consists of: - - - A random sequence of integers (the content to be memorized). - - A delimiter symbol marking the end of the input. - - A sequence of zeros acting as distractors. - - The target sequence requires the model to output padding until the - delimiter, then reproduce the original random sequence. + The first ``memory_length`` tokens are sampled from the content classes. + They are followed by ``seq_len - 1`` blank tokens, a delimiter, and another + ``memory_length`` blanks. Targets are blank until the final segment, where + they reproduce the initial tokens. This follows Arjovsky et al. (2016), + Section 5.1 (https://proceedings.mlr.press/v48/arjovsky16.html). Args: - seq_len (int): Length of the random sequence to memorize. - n_samples (int): Number of samples to generate. - num_classes (int, optional): Number of distinct classes used for the - random sequence. Defaults to 10. The delimiter token uses the - value ``num_classes``. - **kwargs: Additional keyword arguments passed to - :class:`torch.utils.data.DataLoader` (e.g. ``batch_size``, - ``shuffle``). + seq_len: Time lag ``T`` in the paper. Must be positive. The complete + sequence length is ``seq_len + 2 * memory_length``. + n_samples: Number of independent sequences to generate. + num_classes: Alphabet size. The last two classes are reserved for the + blank and delimiter tokens, respectively. + memory_length: Number of content tokens to remember. + return_dataloader: Return a data loader when true, otherwise tensors. + one_hot: Convert inputs to floating-point one-hot vectors so they can be + passed directly to recurrent layers. Targets remain integer class + indices suitable for :class:`torch.nn.CrossEntropyLoss`. + generator: Optional random number generator used for reproducibility. + Only forwarded to the returned data loader's shuffling when it is a + CPU generator; a non-CPU generator is still used for tensor + generation but the loader falls back to its own seeding. + device: Device on which to create the tensors. + **dataloader_kwargs: Arguments passed to + :class:`torch.utils.data.DataLoader`. Returns: - torch.utils.data.DataLoader: A DataLoader yielding batches of - ``(input_seq, target_seq)`` where: - - - ``input_seq`` has shape ``(n_samples, 2 * seq_len + 1)`` and contains - the random sequence, followed by a delimiter token, followed by - distractor zeros. - - ``target_seq`` has shape ``(n_samples, 2 * seq_len + 1)`` and - contains padding + delimiter, followed by the original random - sequence. + A data loader, or ``(inputs, targets)`` when ``return_dataloader=False``. + Integer inputs and targets have shape ``(n_samples, total_length)``. + With ``one_hot=True``, inputs have an additional final dimension of size + ``num_classes``. """ - random_seq = torch.randint(0, num_classes, (n_samples, seq_len)) - delimiter = torch.full((n_samples, 1), num_classes) - distractor_seq = torch.zeros((n_samples, seq_len), dtype=torch.long) - input_seq = torch.cat([random_seq, delimiter, distractor_seq], dim=1) - target_seq = torch.cat( - [torch.full((n_samples, seq_len + 1), num_classes), random_seq], dim=1 + if seq_len < 1: + raise ValueError("seq_len must be positive") + if n_samples < 1: + raise ValueError("n_samples must be positive") + if num_classes < 3: + raise ValueError("num_classes must be at least 3") + if memory_length < 1: + raise ValueError("memory_length must be positive") + + blank_token = num_classes - 2 + delimiter_token = num_classes - 1 + content = torch.randint( + blank_token, + (n_samples, memory_length), + generator=generator, + device=device, + ) + leading_blanks = torch.full( + (n_samples, seq_len - 1), blank_token, dtype=torch.long, device=device + ) + delimiter = torch.full((n_samples, 1), delimiter_token, dtype=torch.long, device=device) + trailing_blanks = torch.full( + (n_samples, memory_length), blank_token, dtype=torch.long, device=device ) + inputs = torch.cat((content, leading_blanks, delimiter, trailing_blanks), dim=1) + targets = torch.cat( + ( + torch.full( + (n_samples, seq_len + memory_length), + blank_token, + dtype=torch.long, + device=device, + ), + content, + ), + dim=1, + ) + + if one_hot: + inputs = F.one_hot(inputs, num_classes=num_classes).to(torch.get_default_dtype()) + if not return_dataloader: + return inputs, targets - dataset = TensorDataset(input_seq, target_seq) - dataloader = DataLoader(dataset, **kwargs) + # torch.utils.data.DataLoader shuffling always runs its generator on CPU + # (torch.randperm has no device argument), so a generator created for a + # non-CPU generation device cannot also drive the loader's shuffling. + loader_generator = generator + if generator is not None and generator.device.type != "cpu": + loader_generator = None - return dataloader + return DataLoader( + TensorDataset(inputs, targets), generator=loader_generator, **dataloader_kwargs + ) diff --git a/torchrecurrent/benchmarks/penn_treebank.py b/torchrecurrent/benchmarks/penn_treebank.py new file mode 100644 index 0000000..5451145 --- /dev/null +++ b/torchrecurrent/benchmarks/penn_treebank.py @@ -0,0 +1,139 @@ +from collections import Counter +from pathlib import Path +from typing import Dict, NamedTuple, Union + +import torch +from torch.utils.data import DataLoader, Dataset + + +class PennTreebankCorpus(NamedTuple): + """Data loaders and vocabulary for the Penn Treebank language model task.""" + + train: DataLoader + validation: DataLoader + test: DataLoader + vocabulary: Dict[str, int] + + +class _BatchedLanguageModelDataset(Dataset): + def __init__(self, tokens: torch.Tensor, batch_size: int, sequence_length: int): + stream_length = tokens.numel() // batch_size + if stream_length < 2: + raise ValueError("each corpus split must contain more tokens than batch_size") + self.tokens = tokens[: batch_size * stream_length].view(batch_size, stream_length) + self.sequence_length = sequence_length + + def __len__(self): + return (self.tokens.shape[1] - 1 + self.sequence_length - 1) // self.sequence_length + + def __getitem__(self, index): + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + start = index * self.sequence_length + length = min(self.sequence_length, self.tokens.shape[1] - start - 1) + inputs = self.tokens[:, start : start + length] + targets = self.tokens[:, start + 1 : start + length + 1] + return inputs, targets + + +def _read_tokens(path: Union[str, Path], eos_token: str): + tokens = [] + with Path(path).open("r", encoding="utf-8") as file: + for line in file: + tokens.extend(line.split()) + tokens.append(eos_token) + return tokens + + +def _build_vocabulary(tokens, max_vocab_size: int, eos_token: str, unk_token: str): + counts = Counter(tokens) + ordered_tokens = sorted(counts, key=lambda token: (-counts[token], token)) + ordered_tokens = [ + token for token in ordered_tokens if token not in (unk_token, eos_token) + ] + vocabulary_tokens = [unk_token, eos_token] + vocabulary_tokens.extend(ordered_tokens[: max_vocab_size - 2]) + return {token: index for index, token in enumerate(vocabulary_tokens)} + + +def _encode(tokens, vocabulary, unk_token: str, device): + unknown_index = vocabulary[unk_token] + return torch.tensor( + [vocabulary.get(token, unknown_index) for token in tokens], + dtype=torch.long, + device=device, + ) + + +def penn_treebank( + train_file: Union[str, Path], + validation_file: Union[str, Path], + test_file: Union[str, Path], + batch_size: int = 20, + sequence_length: int = 35, + *, + max_vocab_size: int = 10_000, + eos_token: str = "", + unk_token: str = "", + device: torch.device = None, +): + """Prepare the standard word-level Penn Treebank language-model benchmark. + + This adapter expects the three whitespace-tokenized text files from the + Mikolov-preprocessed Penn Treebank corpus. It appends ```` at each + newline, builds a training-only vocabulary, maps unseen validation and test + words to ````, and creates contiguous streams for truncated BPTT. The + default batch size and unroll length follow Zaremba et al. (2014) + (https://arxiv.org/abs/1409.2329). + + Penn Treebank data is not downloaded or redistributed by this package. The + original corpus is described by Marcus et al. (1993) + (https://aclanthology.org/J93-2004/). + + Args: + train_file: Path to the preprocessed training split. + validation_file: Path to the preprocessed validation split. + test_file: Path to the preprocessed test split. + batch_size: Number of independent contiguous token streams per batch. + sequence_length: Maximum truncated-BPTT unroll length. + max_vocab_size: Maximum number of tokens, including ``unk_token``. + eos_token: Token appended at the end of every input line. + unk_token: Token used for words outside the training vocabulary. + device: Device on which to store encoded token tensors. + + Returns: + :class:`PennTreebankCorpus` containing train, validation, and test data + loaders plus the token-to-index vocabulary. Each loader yields + ``(inputs, targets)`` tensors shaped ``(batch_size, time_steps)``. The + targets are the inputs shifted forward by one token. + """ + if batch_size < 1: + raise ValueError("batch_size must be positive") + if sequence_length < 1: + raise ValueError("sequence_length must be positive") + if max_vocab_size < 2: + raise ValueError("max_vocab_size must be at least 2") + if eos_token == unk_token: + raise ValueError("eos_token and unk_token must be different") + + train_tokens = _read_tokens(train_file, eos_token) + validation_tokens = _read_tokens(validation_file, eos_token) + test_tokens = _read_tokens(test_file, eos_token) + if not train_tokens: + raise ValueError("training split must not be empty") + + vocabulary = _build_vocabulary(train_tokens, max_vocab_size, eos_token, unk_token) + + def make_loader(tokens): + encoded = _encode(tokens, vocabulary, unk_token, device) + dataset = _BatchedLanguageModelDataset(encoded, batch_size, sequence_length) + return DataLoader(dataset, batch_size=None, shuffle=False) + + return PennTreebankCorpus( + train=make_loader(train_tokens), + validation=make_loader(validation_tokens), + test=make_loader(test_tokens), + vocabulary=vocabulary, + ) diff --git a/torchrecurrent/benchmarks/sequential_cifar10.py b/torchrecurrent/benchmarks/sequential_cifar10.py new file mode 100644 index 0000000..372436d --- /dev/null +++ b/torchrecurrent/benchmarks/sequential_cifar10.py @@ -0,0 +1,91 @@ +import torch +from torch.utils.data import DataLoader, TensorDataset + + +def sequential_cifar10( + images: torch.Tensor, + targets: torch.Tensor, + permutation: torch.Tensor = None, + return_dataloader: bool = True, + batch_size: int = 64, + shuffle: bool = True, + *, + normalize: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + generator: torch.Generator = None, + **dataloader_kwargs, +): + """Convert CIFAR-10 tensors into the sequential or permuted-CIFAR task. + + Every 32 by 32 RGB image becomes a sequence of 1024 three-channel pixel + inputs in raster-scan order. With a permutation, the same fixed pixel + ordering is applied to every image, mirroring the permuted-MNIST protocol + of Arjovsky et al. (2016), Section 5.3 + (https://proceedings.mlr.press/v48/arjovsky16.html). This function extends + that pixel-by-pixel/permuted protocol to CIFAR-10; it is not itself drawn + from a specific paper's CIFAR-10 experiment. The caller supplies the + CIFAR-10 tensors so this package does not require a dataset-download + dependency. + + Args: + images: CIFAR-10 images with shape ``(N, 32, 32, 3)`` or + ``(N, 3, 32, 32)``. + targets: Class indices with shape ``(N,)``. + permutation: Optional permutation of the integers from 0 through 1023. + Reuse the same tensor for training and test data. + return_dataloader: Return a data loader when true, otherwise tensors. + batch_size: Batch size of the returned data loader. + shuffle: Whether the returned data loader shuffles samples. + normalize: Convert integer pixels from ``[0, 255]`` to ``[0, 1]``. + Floating-point inputs are assumed to be normalized already. + dtype: Floating-point dtype of the returned inputs. Defaults to the + current PyTorch default dtype. + device: Device on which to place inputs and targets. + generator: Optional generator used by the data loader when shuffling. + **dataloader_kwargs: Additional arguments passed to + :class:`torch.utils.data.DataLoader`. + + Returns: + A data loader, or ``(sequences, targets)`` when + ``return_dataloader=False``. Sequences have shape ``(N, 1024, 3)`` and + targets have shape ``(N,)`` with dtype :class:`torch.long`. + """ + if images.ndim == 4 and images.shape[1] == 3: + images = images.permute(0, 2, 3, 1) + if images.ndim != 4 or images.shape[1:] != (32, 32, 3): + raise ValueError("images must have shape (N, 32, 32, 3) or (N, 3, 32, 32)") + if targets.ndim != 1 or targets.shape[0] != images.shape[0]: + raise ValueError("targets must have shape (N,) matching images") + + if dtype is None: + dtype = torch.get_default_dtype() + if not dtype.is_floating_point: + raise TypeError("dtype must be a floating-point dtype") + + integer_pixels = not images.is_floating_point() + sequences = images.reshape(images.shape[0], 1024, 3) + sequences = sequences.to(device=device, dtype=dtype) + if normalize and integer_pixels: + sequences = sequences / 255 + + if permutation is not None: + if permutation.ndim != 1 or permutation.numel() != 1024: + raise ValueError("permutation must have shape (1024,)") + permutation = permutation.to(device=sequences.device, dtype=torch.long) + expected = torch.arange(1024, device=sequences.device) + if not torch.equal(torch.sort(permutation).values, expected): + raise ValueError("permutation must contain every index from 0 through 1023") + sequences = sequences.index_select(1, permutation) + + targets = targets.to(device=device, dtype=torch.long) + if not return_dataloader: + return sequences, targets + + return DataLoader( + TensorDataset(sequences, targets), + batch_size=batch_size, + shuffle=shuffle, + generator=generator, + **dataloader_kwargs, + ) diff --git a/torchrecurrent/benchmarks/sequential_mnist.py b/torchrecurrent/benchmarks/sequential_mnist.py new file mode 100644 index 0000000..65a9f27 --- /dev/null +++ b/torchrecurrent/benchmarks/sequential_mnist.py @@ -0,0 +1,92 @@ +import torch +from torch.utils.data import DataLoader, TensorDataset + + +def sequential_mnist( + images: torch.Tensor, + targets: torch.Tensor, + permutation: torch.Tensor = None, + return_dataloader: bool = True, + batch_size: int = 64, + shuffle: bool = True, + *, + normalize: bool = True, + dtype: torch.dtype = None, + device: torch.device = None, + generator: torch.Generator = None, + **dataloader_kwargs, +): + """Convert MNIST tensors into the sequential or permuted-MNIST task. + + Every 28 by 28 image becomes a sequence of 784 scalar inputs in row-major + raster order. With a permutation, the same fixed pixel ordering is applied + to every image, matching the pixel-by-pixel and permuted task setup of + Arjovsky et al. (2016), Section 5.3 + (https://proceedings.mlr.press/v48/arjovsky16.html); the specific base scan + direction is an arbitrary fixed convention and does not affect task + difficulty. The caller supplies the MNIST tensors so this package does not + require a dataset-download dependency. + + Args: + images: MNIST images with shape ``(N, 28, 28)`` or ``(N, 1, 28, 28)``. + targets: Digit class indices with shape ``(N,)``. + permutation: Optional permutation of the integers from 0 through 783. + Reuse the same tensor for training and test data. + return_dataloader: Return a data loader when true, otherwise tensors. + batch_size: Batch size of the returned data loader. + shuffle: Whether the returned data loader shuffles samples. + normalize: Convert integer pixels from ``[0, 255]`` to ``[0, 1]``. + Floating-point inputs are assumed to be normalized already. + dtype: Floating-point dtype of the returned inputs. Defaults to the + current PyTorch default dtype. + device: Device on which to place inputs and targets. + generator: Optional generator used by the data loader when shuffling. + **dataloader_kwargs: Additional arguments passed to + :class:`torch.utils.data.DataLoader`. + + Returns: + A data loader, or ``(sequences, targets)`` when + ``return_dataloader=False``. Sequences have shape ``(N, 784, 1)`` and + targets have shape ``(N,)`` with dtype :class:`torch.long`. + """ + if images.ndim == 4: + if images.shape[1] != 1: + raise ValueError("four-dimensional images must have one channel") + images = images.squeeze(1) + if images.ndim != 3 or images.shape[1:] != (28, 28): + raise ValueError("images must have shape (N, 28, 28) or (N, 1, 28, 28)") + if targets.ndim != 1 or targets.shape[0] != images.shape[0]: + raise ValueError("targets must have shape (N,) matching images") + + if dtype is None: + dtype = torch.get_default_dtype() + if not dtype.is_floating_point: + raise TypeError("dtype must be a floating-point dtype") + + integer_pixels = not images.is_floating_point() + sequences = images.reshape(images.shape[0], 784) + sequences = sequences.to(device=device, dtype=dtype) + if normalize and integer_pixels: + sequences = sequences / 255 + + if permutation is not None: + if permutation.ndim != 1 or permutation.numel() != 784: + raise ValueError("permutation must have shape (784,)") + permutation = permutation.to(device=sequences.device, dtype=torch.long) + expected = torch.arange(784, device=sequences.device) + if not torch.equal(torch.sort(permutation).values, expected): + raise ValueError("permutation must contain every index from 0 through 783") + sequences = sequences.index_select(1, permutation) + + sequences = sequences.unsqueeze(-1) + targets = targets.to(device=device, dtype=torch.long) + if not return_dataloader: + return sequences, targets + + return DataLoader( + TensorDataset(sequences, targets), + batch_size=batch_size, + shuffle=shuffle, + generator=generator, + **dataloader_kwargs, + ) diff --git a/torchrecurrent/benchmarks/timit.py b/torchrecurrent/benchmarks/timit.py new file mode 100644 index 0000000..68f766c --- /dev/null +++ b/torchrecurrent/benchmarks/timit.py @@ -0,0 +1,140 @@ +from typing import NamedTuple, Sequence + +import torch +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import DataLoader, Dataset + + +class TIMITBatch(NamedTuple): + """A padded batch of acoustic features, frame labels, and valid lengths.""" + + features: torch.Tensor + targets: torch.Tensor + lengths: torch.Tensor + + +class TIMITDataset(Dataset): + """Validated variable-length TIMIT feature and phone-state sequences.""" + + def __init__( + self, + features: Sequence[torch.Tensor], + targets: Sequence[torch.Tensor], + feature_size: int = 120, + num_classes: int = 180, + ): + if feature_size < 1: + raise ValueError("feature_size must be positive") + if num_classes < 2: + raise ValueError("num_classes must be at least 2") + if len(features) != len(targets): + raise ValueError( + "features and targets must contain the same number of utterances" + ) + if not features: + raise ValueError("features and targets must not be empty") + + self.features = [] + self.targets = [] + for index, (utterance, labels) in enumerate(zip(features, targets)): + if utterance.ndim != 2 or utterance.shape[1] != feature_size: + raise ValueError( + f"features[{index}] must have shape (frames, {feature_size})" + ) + if not utterance.is_floating_point(): + raise TypeError(f"features[{index}] must have a floating-point dtype") + if labels.ndim != 1 or labels.shape[0] != utterance.shape[0]: + raise ValueError( + f"targets[{index}] must have one label for every feature frame" + ) + if utterance.shape[0] == 0: + raise ValueError(f"features[{index}] must contain at least one frame") + labels = labels.to(dtype=torch.long) + if torch.any(labels < 0) or torch.any(labels >= num_classes): + raise ValueError( + f"targets[{index}] must contain class indices in [0, {num_classes})" + ) + self.features.append(utterance) + self.targets.append(labels) + + def __len__(self): + return len(self.features) + + def __getitem__(self, index): + return self.features[index], self.targets[index] + + +class _PadTIMITBatch: + def __init__(self, padding_value: float, target_padding_value: int): + self.padding_value = padding_value + self.target_padding_value = target_padding_value + + def __call__(self, samples): + features, targets = zip(*samples) + lengths = torch.tensor([sequence.shape[0] for sequence in features]) + padded_features = pad_sequence( + features, batch_first=True, padding_value=self.padding_value + ) + padded_targets = pad_sequence( + targets, batch_first=True, padding_value=self.target_padding_value + ) + return TIMITBatch(padded_features, padded_targets, lengths) + + +def timit( + features: Sequence[torch.Tensor], + targets: Sequence[torch.Tensor], + batch_size: int = 32, + shuffle: bool = True, + *, + feature_size: int = 120, + num_classes: int = 180, + padding_value: float = 0.0, + target_padding_value: int = -100, + generator: torch.Generator = None, + **dataloader_kwargs, +): + """Prepare aligned TIMIT features for frame-level phone-state recognition. + + The default contract follows the TIMIT experiment of Le et al. (2015), + Section 4.4 (https://arxiv.org/abs/1504.00941): each frame contains 40 log-Mel + filterbank coefficients, their deltas, and accelerations (120 features), and + is aligned to one of 180 Kaldi phone states. + + This function does not download the licensed TIMIT corpus or reproduce the + external Kaldi alignment pipeline. It batches user-provided aligned feature + and target tensors for direct use with batch-first recurrent layers. + + Args: + features: Utterance tensors shaped ``(frames, feature_size)``. + targets: Corresponding frame-label tensors shaped ``(frames,)``. + batch_size: Number of utterances per batch. + shuffle: Whether to shuffle utterances between epochs. + feature_size: Expected number of acoustic features per frame. + num_classes: Number of valid phone-state target classes. + padding_value: Value used to pad acoustic feature sequences. + target_padding_value: Value used to pad target sequences. The default + matches :class:`torch.nn.CrossEntropyLoss`'s ``ignore_index``. + generator: Optional generator used for reproducible shuffling. + **dataloader_kwargs: Additional arguments passed to + :class:`torch.utils.data.DataLoader`. + + Returns: + A data loader yielding :class:`TIMITBatch` objects. Features have shape + ``(batch, max_frames, feature_size)``, targets have shape + ``(batch, max_frames)``, and lengths has shape ``(batch,)``. + """ + if batch_size < 1: + raise ValueError("batch_size must be positive") + if 0 <= target_padding_value < num_classes: + raise ValueError("target_padding_value must not be a valid class index") + + dataset = TIMITDataset(features, targets, feature_size, num_classes) + return DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle, + generator=generator, + collate_fn=_PadTIMITBatch(padding_value, target_padding_value), + **dataloader_kwargs, + )