image processing
data generator
computer vision
machine learning
image retrieval

Returning 3 images from data generator

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

Introduction

Returning three images from one generator usually means you are training a multi-input model. Common examples are triplet networks, anchor-positive-negative pipelines, or models that compare three related images in one training step. The generator must return tensors in the exact structure your model expects.

Match Generator Output to Model Input

The key rule is simple: if the model has three inputs, the generator must yield three image batches in the same order. In Keras, that usually means returning either a list, a tuple, or a dictionary keyed by input layer names.

For a triplet-style setup, the shape is often:

  • anchor batch
  • positive batch
  • negative batch
  • optional labels or dummy targets

If you get shape or unpacking errors, the problem is usually a mismatch between the generator's return value and the model definition.

Example with tf.keras.utils.Sequence

Sequence is a good fit because it is deterministic and works well with multiprocessing. This example returns three image tensors and a dummy target array.

python
1import math
2import numpy as np
3import tensorflow as tf
4
5
6class TripletSequence(tf.keras.utils.Sequence):
7    def __init__(self, anchors, positives, negatives, batch_size=4, image_shape=(64, 64, 3)):
8        self.anchors = anchors
9        self.positives = positives
10        self.negatives = negatives
11        self.batch_size = batch_size
12        self.image_shape = image_shape
13
14    def __len__(self):
15        return math.ceil(len(self.anchors) / self.batch_size)
16
17    def __getitem__(self, index):
18        start = index * self.batch_size
19        end = start + self.batch_size
20
21        a = np.array(self.anchors[start:end], dtype="float32")
22        p = np.array(self.positives[start:end], dtype="float32")
23        n = np.array(self.negatives[start:end], dtype="float32")
24
25        y = np.zeros((len(a), 1), dtype="float32")
26        return (a, p, n), y
27
28
29anchors = [np.random.rand(64, 64, 3) for _ in range(10)]
30positives = [np.random.rand(64, 64, 3) for _ in range(10)]
31negatives = [np.random.rand(64, 64, 3) for _ in range(10)]
32
33seq = TripletSequence(anchors, positives, negatives)
34batch_inputs, batch_targets = seq[0]
35print([x.shape for x in batch_inputs], batch_targets.shape)

This structure is directly consumable by a model with three image inputs.

Example Multi-Input Model

The model below accepts three images and embeds them through a shared encoder. The exact loss is not the point here; the important part is that the generator output structure matches the input structure.

python
1import tensorflow as tf
2
3encoder_input = tf.keras.Input(shape=(64, 64, 3))
4x = tf.keras.layers.Flatten()(encoder_input)
5x = tf.keras.layers.Dense(32, activation="relu")(x)
6encoder = tf.keras.Model(encoder_input, x, name="encoder")
7
8anchor_in = tf.keras.Input(shape=(64, 64, 3), name="anchor")
9positive_in = tf.keras.Input(shape=(64, 64, 3), name="positive")
10negative_in = tf.keras.Input(shape=(64, 64, 3), name="negative")
11
12anchor_vec = encoder(anchor_in)
13positive_vec = encoder(positive_in)
14negative_vec = encoder(negative_in)
15
16merged = tf.keras.layers.Concatenate()([anchor_vec, positive_vec, negative_vec])
17output = tf.keras.layers.Dense(1, activation="sigmoid")(merged)
18
19model = tf.keras.Model(
20    inputs=[anchor_in, positive_in, negative_in],
21    outputs=output,
22)
23model.compile(optimizer="adam", loss="binary_crossentropy")

If your generator returns (a, p, n), y, Keras will map those arrays to the three inputs in order.

Use Dictionaries When Order Is Risky

If the model has many inputs or the code is easy to misread, returning a dictionary is safer than relying on positional order.

python
1def __getitem__(self, index):
2    start = index * self.batch_size
3    end = start + self.batch_size
4
5    inputs = {
6        "anchor": np.array(self.anchors[start:end], dtype="float32"),
7        "positive": np.array(self.positives[start:end], dtype="float32"),
8        "negative": np.array(self.negatives[start:end], dtype="float32"),
9    }
10    y = np.zeros((len(inputs["anchor"]), 1), dtype="float32")
11    return inputs, y

Named inputs reduce bugs during refactors.

Common Pitfalls

  • Returning three single images instead of three batches, which breaks batch semantics.
  • Mixing the order of anchor, positive, and negative arrays relative to the model inputs.
  • Forgetting that Sequence must return NumPy arrays or tensors with consistent shapes.
  • Using labels shaped incorrectly for the compiled loss.
  • Debugging the model first when the real issue is generator structure.

Summary

  • A generator can return three images by yielding three batched tensors per step.
  • The return structure must match the model's input structure exactly.
  • 'tf.keras.utils.Sequence is a solid choice for multi-input image pipelines.'
  • Dictionaries keyed by input names are safer than positional tuples in larger models.
  • Most errors come from shape mismatches and incorrect ordering, not from the idea of multi-image generation itself.

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.