Introduction

Keras does not provide a built-in 3D equivalent of the classic 2D ImageDataGenerator, so training on 3D volumes usually means writing a custom generator or Sequence. The main idea is straightforward: load a batch of 3D arrays, apply augmentation consistently across the whole volume, and return tensors shaped for a 3D CNN.

Why 3D Needs a Custom Generator

A 3D volume is not just a stack of unrelated 2D images. If you rotate, flip, or crop slices independently, you destroy the spatial structure of the volume.

That means good 3D augmentation should preserve geometric consistency across all slices in the sample.

Typical 3D shapes look like:

• '(depth, height, width, channels)'
• or sometimes (height, width, depth, channels) depending on the pipeline

The generator must keep the shape convention consistent with the model.

A Simple Sequence for 3D Volumes

1import math
2import numpy as np
3import tensorflow as tf
4
5class VolumeSequence(tf.keras.utils.Sequence):
6    def __init__(self, x_paths, y, batch_size=4, augment=False, shuffle=True):
7        self.x_paths = list(x_paths)
8        self.y = np.array(y)
9        self.batch_size = batch_size
10        self.augment = augment
11        self.shuffle = shuffle
12        self.indices = np.arange(len(self.x_paths))
13        self.on_epoch_end()
14
15    def __len__(self):
16        return math.ceil(len(self.x_paths) / self.batch_size)
17
18    def __getitem__(self, index):
19        batch_ids = self.indices[index * self.batch_size:(index + 1) * self.batch_size]
20
21        x_batch = []
22        y_batch = self.y[batch_ids]
23
24        for i in batch_ids:
25            volume = np.load(self.x_paths[i]).astype(np.float32)
26
27            if self.augment:
28                volume = self.apply_augmentation(volume)
29
30            x_batch.append(volume)
31
32        return np.stack(x_batch), y_batch
33
34    def on_epoch_end(self):
35        if self.shuffle:
36            np.random.shuffle(self.indices)
37
38    def apply_augmentation(self, volume):
39        if np.random.rand() < 0.5:
40            volume = np.flip(volume, axis=1)
41        if np.random.rand() < 0.5:
42            volume = np.flip(volume, axis=2)
43        return volume.copy()

This example assumes each sample is stored as a NumPy file and that each volume already has the right size.

Important Rule for Augmentation

The transform must be applied to the whole volume, not to each slice independently. For example, flipping a volume across one axis is fine because it preserves 3D structure. Rotating each slice separately with unrelated angles is not.

Good augmentation ideas for 3D volumes include:

• flips along spatial axes
• small consistent rotations
• random cropping
• intensity shifts
• noise injection

The right choice depends heavily on the domain. Medical imaging, microscopy, and industrial scans all have different validity rules.

Example 3D CNN Model

1model = tf.keras.Sequential([
2    tf.keras.layers.Input(shape=(64, 128, 128, 1)),
3    tf.keras.layers.Conv3D(16, 3, activation="relu"),
4    tf.keras.layers.MaxPool3D(),
5    tf.keras.layers.Conv3D(32, 3, activation="relu"),
6    tf.keras.layers.MaxPool3D(),
7    tf.keras.layers.GlobalAveragePooling3D(),
8    tf.keras.layers.Dense(1, activation="sigmoid"),
9])
10
11model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

And training:

1train_gen = VolumeSequence(train_paths, train_labels, batch_size=2, augment=True)
2val_gen = VolumeSequence(val_paths, val_labels, batch_size=2, augment=False)
3
4model.fit(train_gen, validation_data=val_gen, epochs=10)

Memory Considerations

3D volumes are much heavier than 2D images, so batching needs more care. A batch size of 32 that is trivial for 2D images may be impossible for large 3D inputs.

That is why practical 3D pipelines often use:

• smaller batches
• preprocessed .npy volumes
• cached datasets
• mixed precision when appropriate

The generator should be written with memory pressure in mind from the start.

Common Pitfalls

One common mistake is trying to force 3D data through 2D image augmentation utilities. That usually breaks spatial consistency.

Another issue is applying different random transforms to different slices of the same volume.

A third pitfall is ignoring memory limits and discovering too late that volume loading plus augmentation makes the training job unstable.

Summary

• Keras 3D training usually needs a custom generator or Sequence.
• Apply augmentation consistently to the full volume, not slice by slice.
• Keep tensor shape conventions aligned with the model.
• Expect smaller batch sizes and higher memory usage than in 2D pipelines.
• Choose augmentation types that are valid for the specific 3D domain.