Comparativa entre tf.data.Dataset.from_generator y un generador en Python estándar

Quiero usar tf.data como una tubería de entrada para mi modelo de TensorFlow. Actualmente tengo un generador en Python que puede generar alrededor de 16000 muestras/seg. Si uso este generador con tf.data.Dataset.from_generator, el rendimiento disminuye a alrededor de 2500 muestras/seg. ¿De dónde proviene esta sobrecarga y hay alguna manera de reducir esta sobrecarga? Porque en este momento, estoy teniendo problemas para utilizar mi GPU.

Aquí hay un ejemplo súper simple:

import numpy as np
import tensorflow as tf
import torch
import tensorflow_datasets as tfds

def gen():
    while True:
        x = np.zeros(shape=(1, 20000), dtype=np.float32)
        yield x, x

gen_py = gen()
gen_tf = tf.data.Dataset.from_generator(
    gen, 
    output_signature=(
        tf.TensorSpec(shape=(1, 20000), dtype=tf.float32), 
        tf.TensorSpec(shape=(1, 20000), dtype=tf.float32)
    )
)

tfds.benchmark(gen_py, batch_size=1, num_iter=500000)  # esto crea 350000 muestras/seg
tfds.benchmark(gen_tf, batch_size=1, num_iter=500000)  # esto crea 5800 muestras/seg

Sé que el ejemplo es extremadamente simple, pero las formas de salida se ajustan aproximadamente a la salida del generador que uso para obtener los datos para mi modelo.