TensorFlow
tf.data
Keras
machine learning
multiple inputs and outputs

tf.data with multiple inputs / outputs in Keras

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

In the realm of deep learning, handling datasets efficiently becomes a critical aspect of model training and evaluation. TensorFlow's tf.data API is a potent tool that facilitates the creation of complex input pipelines from simple, reusable pieces. One of its advanced features includes the capability to manage multiple inputs and outputs in the context of training models in Keras. This article delves deep into using tf.data for handling such multifaceted data structures in Keras.

Introduction to tf.data

The tf.data API enables us to build performant, complex input pipelines seamlessly. By streaming data directly from storage through transformations, tf.data can deliver batches of data ready for model consumption efficiently. Key strengths include parallelism, prefetching, caching, and managing complex data shapes more intuitively.

Handling Multiple Inputs and Outputs

Use Case Scenarios

  • Image and Text Pairs: A model that inputs both an image and its description to produce a single label or a new text description.
  • Multi-task Learning: Models that learn distinct tasks concurrently, potentially producing various outputs.
  • Multi-modal Data: Applications requiring data from different modalities, such as visual, textual, and auditory inputs.

tf.data Pipeline for Multiple Inputs/Outputs

Let's break down a general approach using an example: a model with two inputs - images and tabular data for predicting two outputs - a regression and a classification task.

python
1import tensorflow as tf
2
3# Define feature descriptions for parsing
4feature_description = {
5    'image_raw': tf.io.FixedLenFeature([], tf.string),
6    'tabular': tf.io.VarLenFeature(tf.float32),
7    'label_regression': tf.io.FixedLenFeature([], tf.float32),
8    'label_classification': tf.io.FixedLenFeature([], tf.int64)
9}
10
11def _parse_function(proto):
12    # Parse the input `tf.train.Example` proto using the dictionary above.
13    parsed_features = tf.io.parse_single_example(proto, feature_description)
14
15    # Decode the image; here we assume a footprint of 28x28 grayscale for simplicity
16    image = tf.io.decode_raw(parsed_features['image_raw'], tf.uint8)
17    image = tf.reshape(image, [28, 28, 1])
18
19    # Retrieve sparse tabular data and convert to dense
20    tabular_features = tf.sparse.to_dense(parsed_features['tabular'], default_value=0)
21    
22    # Extract labels
23    label_regression = parsed_features['label_regression']
24    label_classification = tf.cast(parsed_features['label_classification'], tf.int32)
25
26    return {'image': image, 'tabular': tabular_features}, {'regression': label_regression, 'classification': label_classification}
27
28# Creating tf.data.Dataset
29filenames = ['data/file1.tfrecord', 'data/file2.tfrecord']
30raw_dataset = tf.data.TFRecordDataset(filenames)
31
32# Apply the processing function
33parsed_dataset = raw_dataset.map(_parse_function)

Integrating with Keras

To connect the tf.data pipeline into Keras, pass the parsed_dataset directly to the model's fit method.

python
1# Define Keras model
2image_input = tf.keras.layers.Input(shape=(28, 28, 1), name='image')
3tabular_input = tf.keras.layers.Input(shape=(10,), name='tabular')  # Assuming 10 tabular features
4
5x1 = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(image_input)
6x1 = tf.keras.layers.MaxPooling2D((2, 2))(x1)
7x1 = tf.keras.layers.Flatten()(x1)
8
9x2 = tf.keras.layers.Dense(64, activation='relu')(tabular_input)
10
11combined = tf.keras.layers.concatenate([x1, x2])
12
13output_regression = tf.keras.layers.Dense(1, name='regression')(combined)
14output_classification = tf.keras.layers.Dense(3, activation='softmax', name='classification')(combined)  # Assuming 3 classes
15
16model = tf.keras.models.Model(inputs=[image_input, tabular_input], outputs=[output_regression, output_classification])
17
18model.compile(optimizer='adam', 
19              loss={'regression': 'mse', 'classification': 'sparse_categorical_crossentropy'},
20              metrics={'regression': 'mae', 'classification': 'accuracy'})
21
22model.fit(parsed_dataset.batch(32), epochs=10)

Best Practices and Tips

AspectDescription
Data AugmentationApply within the map method to augment your dataset dynamically.
PrefetchingUse prefetch to prepare data in advance, reducing latency.
ParallelismMaximize map performance by using num_parallel_calls=tf.data.AUTOTUNE.
CachingUtilize cache for dataset caching, beneficial for repetitive iteration.
Mixed PrecisionLeverage TensorFlow’s mixed-precision to speed up training when applicable.

Handling Complex Data Flows

Advanced Data Processing

If your application involves advanced scenarios like NLP along with vision, consider leveraging tf.data.experimental. The make_batched_features_dataset utility, the assert_cardinality function for dynamic pipelines, and even dealing with variable-length sequences can all integrate well into a comprehensive data strategy.

Efficient Storage Considerations

Store your data efficiently using GCS or AWS S3 for scale and I/O efficiency. Formats like TFRecords can store serialized tf.train.Example byte sequences, ideal for large and complex datasets.

By proficiently utilizing tf.data capabilities, models can scale effectively while managing intricate data workflows, granting higher flexibility, optimized performance, and simplified handling of multiple inputs and outputs.

This powerful mechanism unlocks new possibilities across research domains, specific use cases, and industrial applications alike, further enriching our deep learning endeavors with TensorFlow and Keras.


Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.