2.3  Synthetic Regression Data

Before we can train a model we need data. Real datasets are what we ultimately care about, but they conflate three separate sources of failure: a misspecified model, a flawed optimization algorithm, and pathological data. When a method performs poorly on real data, all three explanations remain on the table at once. Synthetic data removes this ambiguity by construction. If we know the data-generating process exactly (the true weights \(\mathbf{w}^*\), the true bias \(b^*\), and the noise distribution), then any systematic failure to recover them (beyond the irreducible noise) points to the algorithm or implementation. This is why synthetic datasets are the first test for any new learning method. We confirm that it solves a problem with a known answer before we ever hand it a real one.

%matplotlib inline
from d2l import torch as d2l
import torch
import random
%matplotlib inline
from d2l import tensorflow as d2l
import tensorflow as tf
import random
%matplotlib inline
from d2l import jax as d2l
import jax
from jax import numpy as jnp
import numpy as np
import random
import tensorflow as tf
%matplotlib inline
from d2l import mxnet as d2l
from mxnet import np, npx, gluon
import random
npx.set_np()

2.3.1 Generating the Dataset

For this example, we will work in low dimension for succinctness. The following code snippet generates 1000 examples with 2-dimensional features drawn from a standard normal distribution. The resulting design matrix \(\mathbf{X}\) belongs to \(\mathbb{R}^{1000 \times 2}\). We generate each label by applying a ground truth linear function, corrupting them via additive noise \(\boldsymbol{\epsilon}\), drawn independently and identically for each example:

\[\mathbf{y}= \mathbf{X} \mathbf{w}^* + b^* + \boldsymbol{\epsilon}. \tag{2.3.1}\]

For convenience we assume that \(\boldsymbol{\epsilon}\) is drawn from a normal distribution with mean \(\mu= 0\) and standard deviation \(\sigma = 0.01\). We put the generation code in the __init__ method of a subclass of d2l.DataModule (introduced in Section 2.2.3), calling save_hyperparameters() so that every constructor argument (the parameters w and b, the noise level, the split sizes, and batch_size) is stored as an attribute and the dataset stays introspectable.

class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = d2l.randn(n, len(w))
        eps = d2l.randn(n, 1) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = tf.random.normal((n, w.shape[0]))
        eps = tf.random.normal((n, 1)) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000,
                 batch_size=32, key=None):
        super().__init__()
        self.save_hyperparameters()
        # Resolve the key at call time rather than reusing a key in the signature.
        key = jax.random.key(0) if key is None else key
        n = num_train + num_val
        key1, key2 = jax.random.split(key)
        self.X = jax.random.normal(key1, (n, w.shape[0]))
        eps = jax.random.normal(key2, (n, 1)) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps
class SyntheticRegressionData(d2l.DataModule):
    """Synthetic data for linear regression."""
    def __init__(self, w, b, noise=0.01, num_train=1000, num_val=1000, 
                 batch_size=32):
        super().__init__()
        self.save_hyperparameters()
        n = num_train + num_val
        self.X = d2l.randn(n, len(w))
        eps = d2l.randn(n, 1) * noise
        self.y = d2l.matmul(self.X, d2l.reshape(w, (-1, 1))) + b + eps

Below, we set the true parameters to \(\mathbf{w}^* = [2, -3.4]^\top\) and \(b^* = 4.2\). Later, we can check our estimated parameters against these ground truth values.

data = SyntheticRegressionData(w=d2l.tensor([2, -3.4]), b=4.2)

Each row of data.X is a feature vector in \(\mathbb{R}^2\) and each row of data.y is a scalar label. Let’s have a look at the first entry.

print('features:', data.X[0],'\nlabel:', data.y[0])
features: tensor([-1.2309,  0.6609]) 
label: tensor([-0.5205])
features: tf.Tensor([-0.70784295 -0.2909985 ], shape=(2,), dtype=float32) 
label: tf.Tensor([3.7740877], shape=(1,), dtype=float32)
features: [ 1.0040143 -0.9063372] 
label: [9.265151]
features: [ 1.2680211 -0.9800729] 
label: [10.066712]

2.3.2 Reading the Dataset

Training machine learning models often requires multiple passes over a dataset, grabbing one minibatch of examples at a time. This data is then used to update the model. To illustrate how this works, we implement the get_dataloader method, registering it in the SyntheticRegressionData class via add_to_class (introduced in Section 2.2.1). It takes a batch size, a matrix of features, and a vector of labels, and generates minibatches of size batch_size. As such, each minibatch consists of a tuple of features and labels. Note that we need to be mindful of whether we’re in training or validation mode: in the former, we will want to read the data in random order, whereas for the latter, being able to read data in a pre-defined order may be important for debugging purposes.

@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        j = tf.constant(indices[i : i+self.batch_size])
        yield tf.gather(self.X, j), tf.gather(self.y, j)
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    if train:
        indices = list(range(0, self.num_train))
        # The examples are read in random order
        random.shuffle(indices)
    else:
        indices = list(range(self.num_train, self.num_train+self.num_val))
    for i in range(0, len(indices), self.batch_size):
        batch_indices = d2l.tensor(indices[i: i+self.batch_size])
        yield self.X[batch_indices], self.y[batch_indices]

To build some intuition, let’s inspect the first minibatch of data. Each minibatch of features provides us with both its size and the dimensionality of input features. Likewise, our minibatch of labels will have a matching shape given by batch_size.

X, y = next(iter(data.train_dataloader()))
print('X shape:', X.shape, '\ny shape:', y.shape)
X shape: torch.Size([32, 2]) 
y shape: torch.Size([32, 1])
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)

Iterating over data.train_dataloader() yields distinct minibatches until the dataset is exhausted (try it). Writing the loader by hand makes every step explicit, but it costs us in three ways: all of the data must fit in memory, the iteration is single-threaded Python looping over indices, and there is no prefetching to overlap data loading with computation on the previous batch. The data loaders built into a deep learning framework fix all three. They run several worker processes in parallel, prefetch the next batch while the current one trains, and stream from sources such as files, network streams, or generators that produce data on the fly. We now switch to the framework’s built-in loader, which presents an identical interface to the caller.

2.3.3 Concise Implementation of the Data Loader

Rather than writing our own iterator, we can call the existing API in a framework to load data. As before, we need a dataset with features X and labels y. Beyond that, we set batch_size in the built-in data loader and let it take care of shuffling examples efficiently.

JAX is all about NumPy like API with device acceleration and the functional transformations, so at least the current version doesn’t include data loading methods. With other libraries we already have great data loaders out there, and JAX suggests using them instead. Here we will grab TensorFlow’s data loader, and modify it slightly to make it work with JAX.

@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    dataset = torch.utils.data.TensorDataset(*tensors)
    return torch.utils.data.DataLoader(dataset, self.batch_size,
                                       shuffle=train)
@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    shuffle_buffer = tensors[0].shape[0] if train else 1
    return tf.data.Dataset.from_tensor_slices(tensors).shuffle(
        buffer_size=shuffle_buffer).batch(self.batch_size)
class TensorFlowDataLoader:
    """Expose a tf.data.Dataset as re-iterable NumPy batches."""
    def __init__(self, dataset):
        self.dataset = dataset

    def __iter__(self):
        return self.dataset.as_numpy_iterator()

    def __len__(self):
        return len(self.dataset)

@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    # Use TensorFlow's data loader. JAX and Flax do not provide data-loading
    # functionality. `drop_remainder=train` keeps every
    # *training* minibatch the same shape, so a `@jax.jit`'d step
    # function compiles once per epoch instead of recompiling for the
    # smaller last batch.
    shuffle_buffer = tensors[0].shape[0] if train else 1
    dataset = tf.data.Dataset.from_tensor_slices(tensors).shuffle(
        buffer_size=shuffle_buffer).batch(
            self.batch_size, drop_remainder=train)
    return TensorFlowDataLoader(dataset)
@d2l.add_to_class(d2l.DataModule)
def get_tensorloader(self, tensors, train, indices=slice(0, None)):
    tensors = tuple(a[indices] for a in tensors)
    dataset = gluon.data.ArrayDataset(*tensors)
    return gluon.data.DataLoader(dataset, self.batch_size,
                                 shuffle=train)
@d2l.add_to_class(SyntheticRegressionData)
def get_dataloader(self, train):
    i = slice(0, self.num_train) if train else slice(self.num_train, None)
    return self.get_tensorloader((self.X, self.y), train, i)

The new data loader behaves just like the previous one, except that it is more efficient and has some added functionality.

X, y = next(iter(data.train_dataloader()))
print('X shape:', X.shape, '\ny shape:', y.shape)
X shape: torch.Size([32, 2]) 
y shape: torch.Size([32, 1])
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)
X shape: (32, 2) 
y shape: (32, 1)

For instance, the data loader provided by the framework API supports the built-in __len__ method, so we can query its length, i.e., the number of batches.

len(data.train_dataloader())
32
32
31
32

With 1000 training examples and a batch size of 32, we expect \(\lceil 1000 / 32 \rceil = 32\) batches: 31 full ones and a final partial batch of 8 examples. Note also that the built-in training loader reshuffles the examples at the start of every epoch, just as our hand-rolled loader drew a fresh random order on each call; exercise 8 of Section 2.4 asks why this reshuffling matters.

You may notice that the JAX loader reports 31 batches rather than 32. This is because get_tensorloader passes drop_remainder=True when training: the final partial batch of 8 examples is discarded. We do this so that every training minibatch has an identical shape, which keeps a @jax.jit-compiled training step from being recompiled for the differently sized last batch (a recompilation that can cost minutes per epoch on larger datasets). The price is that we drop a handful of examples each epoch, which is negligible here. A loader that keeps the partial batch would report 32.

2.3.4 Summary

Synthetic data lets us check the recovered parameters against the truth we fixed: because we chose \(\mathbf{w}^*\) and \(b^*\) ourselves, we can see after training whether the estimates agree, which makes such datasets the first place to validate any new algorithm. The SyntheticRegressionData class introduced here packages this data-generating process as a DataModule subclass, separating where the batches come from from how a model consumes them. Along the way we implemented the same get_dataloader protocol twice: a transparent hand-rolled iterator that is easy to read but loads everything in memory and loops in Python, and a framework-native loader that shuffles, prefetches, and parallelizes for us. The hand-rolled version is there to teach; the framework version is what we use from here on.

2.3.5 Exercises

  1. When the number of examples is not divisible by the batch size, the loaders above keep the final partial batch. What does PyTorch’s drop_last argument (and its TensorFlow counterpart, batch(..., drop_remainder=...)) do, and when would you want to enable it?
  2. Suppose that we want to generate a huge dataset, where both the size of the parameter vector w and the number of examples num_examples are large.
    1. What happens if we cannot hold all data in memory?
    2. How would you shuffle the data if it is held on disk? Your task is to design an efficient algorithm that does not require too many random reads or writes. Hint: pseudorandom permutation generators allow you to design a reshuffle without the need to store the permutation table explicitly (Naor and Reingold 1999).
  3. Implement a data generator that produces new data on the fly, every time the iterator is called.
  4. (Reproducibility.) How would you design a random data generator that produces the same dataset every time it is called? Libraries differ here: some expose a single global seed, while others thread an explicit random state through every draw. For the threaded style, explain why fixing that state up front makes the result reproducible, and why re-using the same state for both \(\mathbf{X}\) and \(\boldsymbol{\epsilon}\) (rather than advancing it between the two draws) would be a bug.
  5. (Signal-to-noise and recovery.) Vary the noise standard deviation noise over \(\{0.001, 0.01, 0.1, 0.5, 1.0\}\). After fitting a linear model on each dataset (using the code from Section 2.4 or Section 2.5), how closely does the estimate \(\hat{\mathbf{w}}\) match the true \(\mathbf{w}^* = [2, -3.4]^\top\)? Plot the error \(\|\hat{\mathbf{w}} - \mathbf{w}^*\|_2\) as a function of \(\sigma\). How do you expect it to scale with \(\sigma\) and with the number of training examples, and does the experiment agree?