I'm using tf.keras.Model.fit() on a GPU with input as tf.keras.utils.Sequence generating batches of Numpy arrays. When the generated input or target Numpy arrays are large I notice that the training slows down dramatically. Part of this is of course expected due to the host (CPU) to device (GPU) transfer, but the times I'm seeing suggest an additional overhead somewhere else.
Here is a minimal example reproducing the issue. To keep things simple, the Numpy arrays are created before data_sequence starts generating batches, so there is almost no work done by data_sequence[idx]. Also, the Model consists of a single linear activation, which should just translate to a NoOp.
import tqdm
import numpy as np
import tensorflow as tf
from tensorflow.keras.utils import Sequence
class NumpySequence(Sequence):
def __init__(self, x, y):
self.x = x
self.y = y
def __len__(self):
return len(self.x)
def __getitem__(self, idx):
return self.x[idx], self.y[idx]
workers = 1
steps_per_epoch, epochs = 100, 10
N, H, W, C = 4, 1000, 1000, 10
x_train = []
y_train = []
rng = np.random.default_rng()
for step in tqdm.tqdm(range(steps_per_epoch)):
x_train.append(rng.standard_normal(size=(N, H, W, C), dtype=np.float32))
y_train.append(rng.standard_normal(size=(N, H, W, C), dtype=np.float32))
data_sequence = NumpySequence(x_train, y_train)
model = tf.keras.models.Sequential([tf.keras.layers.Activation("linear")])
model.compile(optimizer='adam', loss=tf.keras.losses.CategoricalCrossentropy())
model.fit(data_sequence, epochs=epochs, workers=workers)
On an NVIDIA Tesla V100 GPU I get:
Epoch 1/10
100/100 [==============================] - 90s 899ms/step - loss: -3.4149e-04
Epoch 2/10
100/100 [==============================] - 91s 907ms/step - loss: -3.4149e-04
...
The input and target Numpy arrays per batch combined have a size of 2*N*H*W*C*sizeof(np.float32) = 320 MB. The CPU to GPU transfer bandwidth should be around 12.4 GB/s, e.g. according to the sample bandwidth test from NVIDIA:
Host to Device Bandwidth, 1 Device(s)
PINNED Memory Transfers
Transfer Size (Bytes) Bandwidth(GB/s)
32000000 12.4
so the CPU to GPU transfer cannot explain the almost 1s per batch alone. Increasing the number of workers does not lead to an improvement.
Is TensorFlow doing some additional serialization or memory copies of the Numpy arrays? If so, is there any way to prevent them?
I'm aware of tf.data.Dataset and the optimizations it applies. Unfortunately, my actual pre-processing pipeline, of which the above example is only a simple caricature, is too complex to port to tf.data.Dataset. Also, if I understand it correctly, tf.keras.Model.fit() internally transforms the tf.keras.utils.Sequence to tf.data.Dataset using the from_generator generator method.