Uso de Datasets con TensorFlow

Este documento es una introducción rápida al uso de datasets con TensorFlow, con un enfoque particular en cómo obtener objetos tf.Tensor de nuestros datasets y cómo transmitir datos desde objetos Dataset de Hugging Face a métodos de Keras como model.fit().

Formato del conjunto de datos

Por defecto, los conjuntos de datos devuelven objetos de Python estándar: enteros, flotantes, cadenas, listas, etc.

Para obtener tensores de TensorFlow en su lugar, puede establecer el formato del conjunto de datos como tf:

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

Un objeto Dataset es un envoltorio de una tabla Arrow, que permite lecturas rápidas desde los arrays del dataset hacia tensores de TensorFlow.

Esto puede ser útil para convertir tu conjunto de datos en un diccionario de objetos Tensor, o para escribir un generador que cargue muestras de TF desde él. Si deseas convertir todo el conjunto de datos a Tensor, simplemente consulta el conjunto de datos completo:

>>> ds[:]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
       [3, 4]])>}

Arreglos de N dimensiones

Si su conjunto de datos consiste en matrices de N dimensiones, verá que por defecto se consideran el mismo tensor si la forma es fija:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}

En caso contrario, un conjunto de datos con formato de TensorFlow devuelve un RaggedTensor en lugar de un tensor único:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': <tf.RaggedTensor [[1, 2], [3]]>}

Sin embargo, esta lógica a menudo requiere comparaciones de forma y copias de datos lentas. Para evitarlo, debes usar explícitamente el tipo de función Array y especificar la forma de tus tensores:

>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
 array([[1, 2],
        [3, 4]])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2, 2), dtype=int64, numpy=
 array([[[1, 2],
         [3, 4]],
 
        [[5, 6],
         [7, 8]]])>}

Otros tipos de características

Los datos de ClassLabel se convierten correctamente en tensores:

>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[:3]
{'label': <tf.Tensor: shape=(3,), dtype=int64, numpy=array([0, 0, 1])>}

También se admiten cadenas de caracteres y objetos binarios:

>>> from datasets import Dataset, Features 
>>> text = ["foo", "bar"]
>>> data = [0, 1] 
>>> ds = Dataset.from_dict({"text": text, "data": data})  
>>> ds = ds.with_format("tf") 
>>> ds[:2]
{'text': <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'foo', b'bar'], dtype=object)>,
 'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>}

También puede formatear explícitamente ciertas columnas y dejar las demás sin formatear:

>>> ds = ds.with_format("tf", columns=["data"], output_all_columns=True)
>>> ds[:2]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>,
 'text': ['foo', 'bar']}

Los objetos de cadena y binarios no se modifican, ya que PyTorch solo admite números.

Los tipos de características Image y Audio también son compatibles.

Para usar el tipo de función Image, deberá instalar el extra vision como pip install datasets[vision].

>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]
{'image': <tf.Tensor: shape=(512, 512, 4), dtype=uint8, numpy=
 array([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=uint8)>}
>>> ds[:2]
{'image': <tf.Tensor: shape=(2, 512, 512, 4), dtype=uint8, numpy=
 array([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=uint8)>}

Para usar el tipo de función Audio, debes instalar el extra audio como pip install datasets[audio].

>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("tf")  
>>> ds[0]["audio"]["array"]
<tf.Tensor: shape=(202311,), dtype=float32, numpy=
array([ 6.1035156e-05,  1.5258789e-05,  1.6784668e-04, ...,
       -1.5258789e-05, -1.5258789e-05,  1.5258789e-05], dtype=float32)>
>>> ds[0]["audio"]["sampling_rate"]
<tf.Tensor: shape=(), dtype=int32, numpy=44100>

Carga de datos

Aunque puedes cargar muestras y lotes individuales simplemente indexando tu conjunto de datos, esto no funcionará si quieres usar métodos de Keras como fit() y predict(). Podrías escribir una función generadora que barajara y cargara lotes de tu conjunto de datos y fit() en ella, pero eso suena a mucho trabajo innecesario. En su lugar, si quieres transmitir datos de tu conjunto de datos al vuelo, recomendamos convertir tu conjunto de datos en un tf.data.Dataset usando el método to_tf_dataset().

La clase tf.data.Dataset cubre una amplia gama de casos de uso: a menudo se crea a partir de Tensors en memoria, o mediante una función de carga para leer archivos en disco o en almacenamiento externo. El conjunto de datos puede transformarse arbitrariamente con el método map(), o se pueden usar métodos como batch() y shuffle() para crear un conjunto de datos listo para el entrenamiento. Estos métodos no modifican los datos almacenados de ninguna manera; en su lugar, construyen un grafo de pipeline de datos que se ejecutará cuando se itere sobre el conjunto de datos, generalmente durante el entrenamiento o la inferencia del modelo. Esto es diferente del método map() de los objetos Dataset de Hugging Face, que ejecuta la función map de inmediato y guarda las nuevas columnas o las columnas modificadas.

Dado que todo el pipeline de preprocesamiento de datos puede compilarse en un tf.data.Dataset, este enfoque permite una carga de datos y un entrenamiento masivamente paralelos y asíncronos. Sin embargo, el requisito de compilación de grafos puede ser una limitación, particularmente para los tokenizadores de Hugging Face, que generalmente no son (¡todavía!) compilables como parte de un grafo de TF. Como resultado, normalmente recomendamos preprocesar el conjunto de datos como un dataset de Hugging Face, donde se pueden usar funciones de Python arbitrarias, y luego convertirlo a tf.data.Dataset después usando to_tf_dataset() para obtener un dataset por lotes listo para el entrenamiento. Para ver ejemplos de este enfoque, consulte los ejemplos o los cuadernos de transformers.

Usando to_tf_dataset()

Usar to_tf_dataset() es sencillo. Una vez que tu conjunto de datos esté preprocesado y listo, simplemente llámalo de la siguiente manera:

>>> from datasets import Dataset
>>> data = {"inputs": [[1, 2],[3, 4]], "labels": [0, 1]}
>>> ds = Dataset.from_dict(data)
>>> tf_ds = ds.to_tf_dataset(
            columns=["inputs"],
            label_cols=["labels"],
            batch_size=2,
            shuffle=True
            )

El objeto tf_ds devuelto aquí ya está completamente listo para el entrenamiento y puede pasarse directamente a model.fit(). Tenga en cuenta que estableció el tamaño del lote al crear el conjunto de datos, por lo que no necesita especificarlo al llamar a fit():

>>> model.fit(tf_ds, epochs=2)

Para una descripción completa de los argumentos, consulte la documentación de to_tf_dataset(). En muchos casos, también será necesario añadir un collate_fn a su llamada. Esta es una función que toma varios elementos del conjunto de datos y los combina en un único lote. Cuando todos los elementos tienen la misma longitud, el collator predeterminado integrado será suficiente, pero para tareas más complejas puede ser necesario un collator personalizado. En particular, muchas tareas tienen muestras con longitudes de secuencia variables, lo que requerirá un data collator que pueda rellenar los lotes correctamente. Puede ver ejemplos de esto en los examples y notebooks de NLP de transformers, donde las longitudes de secuencia variables son muy comunes.

Si notas que la carga con to_tf_dataset es lenta, también puedes usar el argumento num_workers. Esto inicia varios subprocesos para cargar los datos en paralelo. Esta función es reciente y aún es algo experimental: ¡abre un problema si encuentras algún error al usarla!

When to use to_tf_dataset

Es posible que el lector atento haya notado a estas alturas que hemos ofrecido dos enfoques para lograr el mismo objetivo: si desea pasar su conjunto de datos a un modelo de TensorFlow, puede convertir el conjunto de datos en un Tensor o un dict de Tensors usando .with_format('tf'), o bien puede convertir el conjunto de datos en un tf.data.Dataset con to_tf_dataset(). Cualquiera de estas opciones puede pasarse a model.fit(), así que, ¿cuál debería elegir?

Lo clave es reconocer que al convertir todo el conjunto de datos a Tensors, este se vuelve estático y se carga por completo en la RAM. Esto es simple y conveniente, pero si se aplica alguno de los siguientes casos, probablemente deba usar to_tf_dataset() en su lugar:

• Tu conjunto de datos es demasiado grande para caber en la RAM. to_tf_dataset() transmite solo un lote a la vez, por lo que incluso conjuntos de datos muy grandes pueden manejarse con este método.
• Deseas aplicar transformaciones aleatorias usando dataset.with_transform() o collate_fn. Esto es común en varias modalidades, como las aumentaciones de imagen al entrenar modelos de visión, o el enmascaramiento aleatorio al entrenar modelos de lenguaje enmascarados. Usar to_tf_dataset() aplicará esas transformaciones en el momento en que se carga un lote, lo que significa que las mismas muestras recibirán diferentes aumentaciones cada vez que se carguen. Esto es generalmente lo que deseas.
• Tus datos tienen una dimensión variable, como los textos de entrada en NLP que consisten en un número variable de tokens. Cuando creas un lote con muestras de dimensión variable, la solución estándar es rellenar las muestras más cortas hasta la longitud de la más larga. Cuando transmites muestras de un conjunto de datos con to_tf_dataset, puedes aplicar este relleno a cada lote a través de tu collate_fn. Sin embargo, si deseas convertir dicho conjunto de datos en Tensors densos, entonces tendrás que rellenar las muestras hasta la longitud de la muestra más larga de todo el conjunto de datos. Esto puede resultar en una gran cantidad de relleno, lo que desperdicia memoria y reduce la velocidad de tu modelo.

Limitaciones y advertencias

Por ahora, to_tf_dataset() siempre devuelve un dataset por lotes; ¡pronto añadiremos soporte para datasets sin lotes!