TensorFlow Dataset Pipelines With Python | Towards Data Science

Data pipelines are the less glamorous but still fundamental building blocks in any scalable, production-quality ML solution. Indeed, the vast majority of ML is actually data wrangling — so it makes sense that a strong pipeline is a big factor in building strong solutions.
Focusing on TensorFlow 2, we have a wonderful thing called a Dataset object built-in with the library. Using dataset objects, we can design efficient data pipelines with significantly less effort — the result is a cleaner, logical, and highly optimized pipeline.
We will be doing a deep-dive on the dataset object. Covering what they are, why we should use them, and how we use them.
What Are Datasets?  - Why Use Them?Reading Data  - In-memory  - From FileBasic Operations  - Batching and Shuffling  - Data Transformations with MapFeeding Datasets into a Model  - Multiple Inputs/Outputs  - Train-Validation Split
If you prefer video, I’ve covered everything in this article here too:
Datasets are an object contained within the tf.data API, which represents a sequence of elements. Each element is given a particular structure dependent on the format required by our models.
The tf.data API is TensorFlow’s built-in approach for building input data pipelines — providing methods for developing more efficient pipelines with less code.
The dataset abstraction makes data extraction, transformation, and loading incredibly easy.
For example, we can build a pipeline to iteratively load from file (rather than keeping everything in memory), perform some given transformations (like tokenization of text data), batch, shuffle, and load the data into our model for training.
All of this typically needs no more than a few lines of code. To batch and shuffle our dataset takes nothing more than dataset.shuffle(10000).batch(64).
As mentioned, we have two options regarding how we read data into our dataset, (1) from in-memory or (2) from disk.
The simplest approach is simply reading the data directly into Python using a dataframe, array, lists, or other data structure. This approach is absolutely okay to use but can quickly become unmanageable with larger datasets.
So, if working with smaller datasets — you have these options:
import tensorflow as tf# read from listsinputs = [0, 1, 2, 3]labels = [1, 0, 1, 0]dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))# the exact same approach for Numpy arraysinputs = np.asarray([0, 1, 2, 3])labels = np.asarray([1, 0, 1, 0])dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))# and from dataframesdf = pd.DataFrame({    'inputs': [0, 1, 2, 3],    'labels': [1, 0, 1, 0]})dataset = tf.data.Dataset.from_tensor_slices(df)
These methods produce the dataset structures shown above. Clearly, reading from in-memory objects is incredibly straightforward — but it’s not ideal for larger datasets.
The second approach is to read in the data from an out-of-memory source, such as a local hard-drive.
When taking this approach, we benefit from only loading the data we need in memory — meaning that for a 20M sample dataset, we can avoid loading the full dataset at once and iteratively load only what we need.
We will use the IMDB movie reviews dataset in this example. You can download it from Kaggle here.
Our training data is contained inside train.tsv, which we will use to train a sentiment classifier. For this, we would need to take the text contained in Content as our input feature and Sentiment as our target labels. We do not need any other columns.
dataset = tf.data.experimental.make_csv_dataset(    "train.tsv",    batch_size=8,    field_delim='\t',    select_columns=['Phrase', 'Sentiment'],    label_name='Sentiment')
With this approach, we can easily keep the data stored in-memory low and still shuffle and batch data with ease — it’s an incredibly convenient option.
Now we’ve covered the basics of reading and writing dataset objects; we can begin transforming our loaded dataset.
When reading from file, these operations are not needed as they are built-in to read functions like tf.data.experimental.make_csv_dataset — but this does not happen when we build the dataset from in-memory.
Fortunately, it’s still incredibly easy. We give a large number (10000 is pretty common) as the argument to shuffle and our training batch size to batch:
dataset = dataset.shuffle(10000).batch(64)
It takes nothing more than this single line to shuffle and batch our dataset!
We can use the map function to perform operations on each sample within xour dataset. For example, for predicting the next time-step in a sequence, we may want to train on input data, which consists of timesteps n to n+8, and output data consisting of timesteps n+1 to n+9.
Initially, our dataset may consist of many samples containing sequences of 10 time-steps. To train our model, we need to split these 10-long sequences into two sets of 9-long sequences — the input data, the other our target data.
To perform this operation, we use the map method, like so:
def Xy_split(x):    X = x[:-1]  # [0, 1, 2, 3, 4, 5, 6, 7, 8]    y = x[1:]   # [1, 2, 3, 4, 5, 6, 7, 8, 9]    return X, ydataset = dataset.map(Xy_split)
When feeding our datasets into a model for training the default format of (input, output) is expected. Meaning every record/batch contained in the dataset should contain an input tensor and label/output tensor.
With that format, we can pass our dataset object to our training method like so:
model.fit(dataset, epochs=2)
The single input/output format is not always the one we will need to use. For any restructuring, we can again rely on the map method.
We can map multiple inputs and outputs by mapping our dataset into a dictionary format where we have our input name (the key) pointing to the set of values we will be using (the value).
If we match up the input/output layer names to the input/output dataset dictionary keys, we can pass dataset to our model as per usual!
# current dataset format is (inputs1, inputs2, labels)def map_func(x):    return {        'inputs1': x[0],        'inputs2': x[1]    }, x[2]dataset = dataset.map(map_func)input_layer_1 = tf.keras.layers.Input(shape=(100,), name='inputs1')input_layer_2 = tf.keras.layers.Input(shape=(100,), name='inputs2')
A common use-case for this is creating two input layers for transformer models. Many of these models require an input ID tensor and corresponding attention mask tensor — where we can use the same logic:
def transformer_fmt(x):    return {        'input_ids': x[0],        'attention_mask': x[1],    }, x[2]  # x[2] are the output labelsdataset = dataset.map(transformer_fmt)input_ids = tf.keras.layers.Input(shape=(100,), name='input_ids')mask = tf.keras.layers.Input(shape=(100,), name='attention_mask')...model.fit(dataset, epochs=2)
Another important point here is the split of training and validation sets (test-sets too). The easiest way to implement a split in a dataset is to use the take and split methods.
To take a given number of records, we use take — like so:
first5 = dataset.take(5)
We can then skip a given number of records using skip, returning all that follow — like so:
miss5 = dataset.skip(5)
With a combination of both methods, we can create a split in our data. For example, for a 70–30% training-validation split, we do:
train = dataset.take(round(length*0.7))val = dataset.skip(round(length*0.7))
And create another split to add a test-set.
When we come to feeding the training and validations sets into our model for training, we do so like this:
history = model.fit(    train,    validation_data=val,    epochs=2)
We’ve covered how-to build cleaner, more efficient data input pipelines in TF2 using dataset objects! Including:
What they are and why we should use them
Loading in-memory data to datasets
Reading into datasets from file
Batching and shuffling
Transforming datasets with map
Feeding datasets into models
How to deal with multiple input/output layers
Train-val-test splits
Which is all we need to know for a big portion of ML input pipelines. Everything we’ve done is incredibly easy to perform, and lightning-fast thanks to the clear and efficient tf.data API.
I hope you’ve enjoyed this article! Let me know if you have any questions or ideas via Twitter or in the comments below.
Thanks for reading!