Использование Datasets с TensorFlow

Этот документ представляет собой краткое введение к использованию datasets с TensorFlow, с особым акцентом на то, как получать объекты tf.Tensor из наших датасетов и как передавать данные из объектов Hugging Face Dataset в методы Keras, такие как model.fit().

Формат набора данных

По умолчанию наборы данных возвращают обычные объекты Python: целые числа, числа с плавающей точкой, строки, списки и т. д.

Чтобы получать тензоры TensorFlow, можно задать формат набора данных как 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]])>}

Объект Dataset представляет собой обёртку над таблицей Arrow, которая позволяет быстро читать данные из массивов набора в тензоры TensorFlow.

Это может быть полезно для преобразования набора данных в словарь объектов Tensor или для создания генератора, загружающего из него образцы TF. Если вы хотите преобразовать весь набор данных в Tensor, просто запросите полный набор данных:

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

Многомерные массивы

Если ваш набор данных состоит из N-мерных массивов, вы увидите, что по умолчанию они считаются одним и тем же тензором, если форма фиксирована:

>>> 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]])>}

В противном случае датасет в формате TensorFlow выводит RaggedTensor вместо одиночного тензора:

>>> 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]]>}

Однако такая логика часто требует медленных сравнений форм и копирования данных. Чтобы избежать этого, необходимо явно использовать тип фичи Array и указать форму ваших тензоров:

>>> 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]]])>}

Другие типы признаков

Данные ClassLabel корректно преобразуются в тензоры:

>>> 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])>}

Поддерживаются также строки и двоичные объекты:

>>> 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])>}

Вы также можете явно отформатировать определённые столбцы, оставив остальные без форматирования:

>>> 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']}

Строковые и двоичные объекты остаются без изменений, так как PyTorch поддерживает только числа.

Также поддерживаются типы функций Image и Audio.

Чтобы использовать тип функции Image, необходимо установить дополнительный пакет vision командой 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)>}

Чтобы использовать тип функции Audio, необходимо установить дополнительный пакет audio командой 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>

Загрузка данных

Хотя вы можете загружать отдельные примеры и батчи, просто обращаясь к элементам набора данных по индексу, это не сработает, если вы хотите использовать методы Keras, такие как fit() и predict(). Вы могли бы написать функцию-генератор, которая перемешивает и загружает батчи из вашего набора данных, и вызвать fit() на ней, но это звучит как много ненужной работы. Вместо этого, если вы хотите стримить данные из вашего набора данных на лету, мы рекомендуем преобразовать ваш набор данных в tf.data.Dataset с помощью метода to_tf_dataset().

Класс tf.data.Dataset покрывает широкий спектр сценариев использования — он часто создаётся из тензоров в памяти или с помощью функции загрузки для чтения файлов с диска или внешнего хранилища. Набор данных можно произвольно преобразовывать с помощью метода map(), а также использовать методы вроде batch() и shuffle() для создания набора данных, готового к обучению. Эти методы никак не изменяют сохранённые данные — вместо этого они строят граф конвейера данных, который выполняется при итерации по набору данных, обычно во время обучения или инференса модели. Это отличается от метода map() объектов Hugging Face Dataset, который выполняет функцию map немедленно и сохраняет новые или изменённые столбцы.

Поскольку весь конвейер предобработки данных может быть скомпилирован в tf.data.Dataset, этот подход позволяет осуществлять массово-параллельную асинхронную загрузку данных и обучение. Однако требование к компиляции графа может стать ограничением, особенно для токенизаторов Hugging Face, которые обычно (пока!) не могут быть скомпилированы как часть графа TF. В результате мы обычно рекомендуем предварительно обрабатывать датасет как датасет Hugging Face, где можно использовать произвольные функции Python, а затем конвертировать его в tf.data.Dataset с помощью to_tf_dataset(), чтобы получить датасет с батчами, готовый к обучению. Чтобы увидеть примеры этого подхода, обратитесь к примерам или ноутбукам для transformers.

Использование to_tf_dataset()

Использование to_tf_dataset() не представляет сложности. Как только ваш набор данных предобработан и готов, просто вызовите его следующим образом:

>>> 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
            )

Возвращаемый здесь объект tf_ds полностью готов для обучения и может быть передан непосредственно в model.fit(). Обратите внимание, что размер батча задаётся при создании набора данных, поэтому указывать его при вызове fit() не нужно:

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

Полное описание аргументов приведено в документации to_tf_dataset(). Во многих случаях вам также потребуется добавить collate_fn в ваш вызов. Это функция, которая принимает несколько элементов набора данных и объединяет их в один батч. Если все элементы имеют одинаковую длину, встроенного коллаборатора по умолчанию будет достаточно, но для более сложных задач может потребоваться пользовательский коллаборатор. В частности, во многих задачах образцы имеют разную длину последовательностей, что требует data collator, способного корректно дополнять батчи. Примеры этого можно найти в transformers NLP examples и notebooks, где переменные длины последовательностей встречаются очень часто.

Если загрузка с помощью to_tf_dataset выполняется медленно, можно также использовать аргумент num_workers. Он запускает несколько подпроцессов для параллельной загрузки данных. Эта функция появилась недавно и пока остаётся экспериментальной — пожалуйста, создайте issue, если вы столкнётесь с какими-либо ошибками при её использовании!

Когда использовать to_tf_dataset

Внимательный читатель, вероятно, уже заметил, что мы предложили два подхода для достижения одной и той же цели: если вы хотите передать ваш набор данных в модель TensorFlow, вы можете либо преобразовать набор данных в Tensor или dict из Tensors с помощью .with_format('tf'), либо преобразовать его в tf.data.Dataset с помощью to_tf_dataset(). Любой из этих вариантов можно передать в model.fit(), так какой же из них следует выбрать?

Ключевой момент, который нужно учитывать: при преобразовании всего датасета в Tensor он становится статическим и полностью загружается в оперативную память. Это просто и удобно, но если выполняется хотя бы одно из следующих условий, лучше использовать to_tf_dataset():

• Ваш набор данных слишком большой, чтобы поместиться в оперативную память. to_tf_dataset() обрабатывает данные по одному батчу за раз, поэтому с помощью этого метода можно работать даже с очень большими наборами данных.
• Вы хотите применять случайные преобразования с помощью dataset.with_transform() или collate_fn. Это распространённый подход в нескольких модальностях, например, аугментация изображений при обучении моделей компьютерного зрения или случайное маскирование при обучении моделей с маскированным языковым моделированием. Использование to_tf_dataset() применяет эти преобразования в момент загрузки батча, что означает, что одни и те же примеры будут получать разные аугментации каждый раз при загрузке. Обычно именно этого и требуется.
• Ваши данные имеют переменную размерность, например, входные тексты в NLP, состоящие из разного количества токенов. При создании батча из образцов с переменной размерностью стандартным решением является дополнение более коротких образцов до длины самого длинного. При потоковой передаче образцов из набора данных с помощью to_tf_dataset вы можете применить такое дополнение к каждому батчу через вашу collate_fn. Однако, если вы хотите преобразовать такой набор данных в плотные Tensor, вам придется дополнять образцы до длины самого длинного образца в всем наборе данных! Это может привести к огромному количеству дополнений, что расходует память и снижает скорость вашей модели.

Ограничения

Сейчас to_tf_dataset() всегда возвращает датасет с батчами — поддержка датасетов без батчей будет добавлена в ближайшее время!