tf.data.Dataset

Представляет потенциально большой набор элементов.

tf.data.Dataset(
    variant_tensor
)

API tf.data.Dataset поддерживает создание описательных и эффективных входных конвейеров. Dataset использование следует общей схеме:

1. Создайте исходный набор данных из ваших входных данных.
2. Примените преобразования набора данных для предобработки данных.
3. Пройдите по набору данных и обработайте элементы.

Итерация происходит потоковым образом, поэтому весь набор данных не должен помещаться в оперативную память.

Исходные наборы данных:

Самый простой способ создать набор данных — создать его из Python list:

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
for element in dataset:
  print(element)
tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)
tf.Tensor(3, shape=(), dtype=int32)

Для обработки строк из файлов используйте tf.data.TextLineDataset:

dataset = tf.data.TextLineDataset(["file1.txt", "file2.txt"])

Для обработки записей, записанных в формате TFRecord , используйте TFRecordDataset:

dataset = tf.data.TFRecordDataset(["file1.tfrecords", "file2.tfrecords"])

Для создания набора данных всех файлов, соответствующих шаблону, используйте tf.data.Dataset.list_files:

dataset = tf.data.Dataset.list_files("/path/*.txt")

См. tf.data.FixedLengthRecordDataset и tf.data.Dataset.from_generator для получения дополнительных способов создания наборов данных.

Преобразования:

После получения набора данных вы можете применить преобразования для подготовки данных к вашей модели:

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
dataset = dataset.map(lambda x: x*2)
list(dataset.as_numpy_iterator())
[2, 4, 6]

Общие термины:

Элемент: единственный результат вызова next() для итератора набора данных. Элементы могут быть вложенными структурами, содержащими несколько компонентов. Например, элемент (1, (3, "apple")) содержит один кортеж, вложенный в другой кортеж. Компонентами являются 1, 3, и "apple".

Компонент: лист вложенной структуры элемента.

Поддерживаемые типы:

Элементы могут быть вложенными структурами кортежей, именованных кортежей и словарей. Обратите внимание, что списки Python не обрабатываются как вложенные структуры компонентов. Вместо этого списки преобразуются в тензоры и обрабатываются как компоненты. Например, элемент (1, [1, 2, 3]) имеет только два компонента; тензор 1 и тензор [1, 2, 3]. Компоненты элементов могут быть любого типа, представимого tf.TypeSpec, включая tf.Tensor, tf.data.Dataset, tf.sparse.SparseTensor, tf.RaggedTensor и tf.TensorArray.

a = 1 # Integer element
b = 2.0 # Float element
c = (1, 2) # Tuple element with 2 components
d = {"a": (2, 2), "b": 3} # Dict element with 3 components
Point = collections.namedtuple("Point", ["x", "y"])
e = Point(1, 2) # Named tuple
f = tf.data.Dataset.range(10) # Dataset element

Дополнительную информацию см. в данном руководстве.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
dataset.element_spec
TensorSpec(shape=(), dtype=tf.int32, name=None)

Методы

apply

Просмотреть исходный код

apply(
    transformation_func
)

Применяет функцию преобразования к этому набору данных.

apply позволяет объединять пользовательские Dataset преобразования, которые представлены как функции, принимающие один Dataset аргумент и возвращающие преобразованный Dataset.

dataset = tf.data.Dataset.range(100)
def dataset_fn(ds):
  return ds.filter(lambda x: x < 5)
dataset = dataset.apply(dataset_fn)
list(dataset.as_numpy_iterator())
[0, 1, 2, 3, 4]

as_numpy_iterator

Просмотреть исходный код

as_numpy_iterator()

Возвращает итератор, который преобразует все элементы набора данных в NumPy.

Используйте as_numpy_iterator для просмотра содержимого вашего набора данных. Чтобы увидеть формы и типы элементов, выведите элементы набора данных напрямую вместо использования as_numpy_iterator.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
for element in dataset:
  print(element)
tf.Tensor(1, shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)
tf.Tensor(3, shape=(), dtype=int32)

Этот метод требует, чтобы вы работали в режиме eager, а элемент_spec набора данных содержал только TensorSpec компоненты.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
for element in dataset.as_numpy_iterator():
  print(element)
1
2
3

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
print(list(dataset.as_numpy_iterator()))
[1, 2, 3]

as_numpy_iterator() сохранит вложенную структуру элементов набора данных.

dataset = tf.data.Dataset.from_tensor_slices({'a': ([1, 2], [3, 4]),
                                              'b': [5, 6]})
list(dataset.as_numpy_iterator()) == [{'a': (1, 3), 'b': 5},
                                      {'a': (2, 4), 'b': 6}]
True

batch

Просмотреть исходный код

batch(
    batch_size,
    drop_remainder=False,
    num_parallel_calls=None,
    deterministic=None,
    name=None
)

Объединяет последовательные элементы этого набора данных в пакеты.

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3)
list(dataset.as_numpy_iterator())
[array([0, 1, 2]), array([3, 4, 5]), array([6, 7])]

dataset = tf.data.Dataset.range(8)
dataset = dataset.batch(3, drop_remainder=True)
list(dataset.as_numpy_iterator())
[array([0, 1, 2]), array([3, 4, 5])]

Компоненты результирующего элемента будут иметь дополнительное внешнее измерение, которое будет batch_size (или N % batch_size для последнего элемента, если batch_size не делит количество входных элементов N равномерно и drop_remainder равно False). Если ваша программа зависит от того, что пакеты имеют одинаковое внешнее измерение, вы должны установить аргумент drop_remainder в True , чтобы предотвратить создание меньшего пакета.

Примечание: Если ваша программа требует, чтобы данные имели статически известную форму (например, при использовании XLA), вы должны использовать drop_remainder=True. Без drop_remainder=True форма выходного набора данных будет иметь неизвестное ведущее измерение из-за возможности меньшего конечного пакета.

bucket_by_sequence_length

Просмотреть исходный код

bucket_by_sequence_length(
    element_length_func,
    bucket_boundaries,
    bucket_batch_sizes,
    padded_shapes=None,
    padding_values=None,
    pad_to_bucket_boundary=False,
    no_padding=False,
    drop_remainder=False,
    name=None
)

Преобразование, группирующее элементы в Dataset по длине.

Элементы Dataset группируются по длине, а затем заполняются и объединяются в пакеты.

Это полезно для задач с последовательностями, у которых элементы имеют переменную длину. Группировка элементов с аналогичными длинами уменьшает общую долю заполнения в пакете, что повышает эффективность шагов обучения.

Ниже приведён пример букетизации входных данных в 3 пакета "[0, 3), [3, 5), [5, inf)" на основе длины последовательности с размером пакета 2.

elements = [
  [0], [1, 2, 3, 4], [5, 6, 7],
  [7, 8, 9, 10, 11], [13, 14, 15, 16, 19, 20], [21, 22]]
dataset = tf.data.Dataset.from_generator(
    lambda: elements, tf.int64, output_shapes=[None])
dataset = dataset.bucket_by_sequence_length(
        element_length_func=lambda elem: tf.shape(elem)[0],
        bucket_boundaries=[3, 5],
        bucket_batch_sizes=[2, 2, 2])
for elem in dataset.as_numpy_iterator():
  print(elem)
[[1 2 3 4]
[5 6 7 0]]
[[ 7  8  9 10 11  0]
[13 14 15 16 19 20]]
[[ 0  0]
[21 22]]

cache

Посмотреть исходный код

cache(
    filename='', name=None
)

Кэширует элементы в этом наборе данных.

В первый раз при итерации по набору данных его элементы будут кэшированы либо в указанном файле, либо в оперативной памяти. При последующих итерациях будет использоваться кэшированная информация.

Примечание: Для завершения кэширования входной набор данных должен быть пройден целиком. В противном случае последующие итерации не будут использовать кэшированные данные.

dataset = tf.data.Dataset.range(5)
dataset = dataset.map(lambda x: x**2)
dataset = dataset.cache()
# The first time reading through the data will generate the data using
# `range` and `map`.
list(dataset.as_numpy_iterator())
[0, 1, 4, 9, 16]
# Subsequent iterations read from the cache.
list(dataset.as_numpy_iterator())
[0, 1, 4, 9, 16]

При кэшировании в файл кэшированные данные будут сохраняться между запусками. Даже первая итерация данных будет читать данные из файла кэша. Изменение входного потока перед вызовом .cache() не будет иметь эффекта, пока файл кэша не будет удален или имя файла не будет изменено.

dataset = tf.data.Dataset.range(5)
dataset = dataset.cache("/path/to/file")
list(dataset.as_numpy_iterator())
# [0, 1, 2, 3, 4]
dataset = tf.data.Dataset.range(10)
dataset = dataset.cache("/path/to/file")  # Same file!
list(dataset.as_numpy_iterator())
# [0, 1, 2, 3, 4]

Примечание: cache будет генерировать точно такие же элементы при каждой итерации по набору данных. Если вы хотите случайный порядок итерации, убедитесь, что вы вызываете shuffle *после* вызова cache.

cardinality

Посмотреть исходный код

cardinality()

Возвращает мощность набора данных, если она известна.

cardinality может вернуть tf.data.INFINITE_CARDINALITY, если набор данных содержит бесконечное количество элементов, или tf.data.UNKNOWN_CARDINALITY, если анализ не смог определить количество элементов в наборе данных (например, когда источником набора данных является файл).

dataset = tf.data.Dataset.range(42)
print(dataset.cardinality().numpy())
42
dataset = dataset.repeat()
cardinality = dataset.cardinality()
print((cardinality == tf.data.INFINITE_CARDINALITY).numpy())
True
dataset = dataset.filter(lambda x: True)
cardinality = dataset.cardinality()
print((cardinality == tf.data.UNKNOWN_CARDINALITY).numpy())
True

choose_from_datasets

Посмотреть исходный код

@staticmethod
choose_from_datasets(
    datasets, choice_dataset, stop_on_empty_dataset=True
)

Создаёт набор данных, который детерминированно выбирает элементы из datasets.

Например, при следующих наборах данных:

datasets = [tf.data.Dataset.from_tensors("foo").repeat(),
            tf.data.Dataset.from_tensors("bar").repeat(),
            tf.data.Dataset.from_tensors("baz").repeat()]

# Define a dataset containing `[0, 1, 2, 0, 1, 2, 0, 1, 2]`.
choice_dataset = tf.data.Dataset.range(3).repeat(3)

result = tf.data.Dataset.choose_from_datasets(datasets, choice_dataset)

Элементы набора данных result будут:

"foo", "bar", "baz", "foo", "bar", "baz", "foo", "bar", "baz"

concatenate

Посмотреть исходный код

concatenate(
    dataset, name=None
)

Создает набор данных путем конкатенации заданного набора данных с этим набором данных.

a = tf.data.Dataset.range(1, 4)  # ==> [ 1, 2, 3 ]
b = tf.data.Dataset.range(4, 8)  # ==> [ 4, 5, 6, 7 ]
ds = a.concatenate(b)
list(ds.as_numpy_iterator())
[1, 2, 3, 4, 5, 6, 7]
# The input dataset and dataset to be concatenated should have
# compatible element specs.
c = tf.data.Dataset.zip((a, b))
a.concatenate(c)
Traceback (most recent call last):
TypeError: Two datasets to concatenate have different types
<dtype: 'int64'> and (tf.int64, tf.int64)
d = tf.data.Dataset.from_tensor_slices(["a", "b", "c"])
a.concatenate(d)
Traceback (most recent call last):
TypeError: Two datasets to concatenate have different types
<dtype: 'int64'> and <dtype: 'string'>

enumerate

Посмотреть исходный код

enumerate(
    start=0, name=None
)

Нумерует элементы в этом наборе данных.

Аналогично функции enumerate в Python.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
dataset = dataset.enumerate(start=5)
for element in dataset.as_numpy_iterator():
  print(element)
(5, 1)
(6, 2)
(7, 3)

# The (nested) structure of the input dataset determines the
# structure of elements in the resulting dataset.
dataset = tf.data.Dataset.from_tensor_slices([(7, 8), (9, 10)])
dataset = dataset.enumerate()
for element in dataset.as_numpy_iterator():
  print(element)
(0, array([7, 8], dtype=int32))
(1, array([ 9, 10], dtype=int32))

filter

Посмотреть исходный код

filter(
    predicate, name=None
)

Фильтрует этот набор данных в соответствии с predicate.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
dataset = dataset.filter(lambda x: x < 3)
list(dataset.as_numpy_iterator())
[1, 2]
# `tf.math.equal(x, y)` is required for equality comparison
def filter_fn(x):
  return tf.math.equal(x, 1)
dataset = dataset.filter(filter_fn)
list(dataset.as_numpy_iterator())
[1]

flat_map

Просмотр исходного кода

flat_map(
    map_func, name=None
)

Применяет map_func к этому набору данных и сплющивает результат.

Тип подписи:

def flat_map(
  self: Dataset[T],
  map_func: Callable[[T], Dataset[S]]
) -> Dataset[S]

Используйте flat_map, если хотите убедиться, что порядок вашего набора данных сохранится. Например, чтобы сплющить набор данных из пакетов в набор данных их элементов:

dataset = tf.data.Dataset.from_tensor_slices(
    [[1, 2, 3], [4, 5, 6], [7, 8, 9]])
dataset = dataset.flat_map(
    lambda x: tf.data.Dataset.from_tensor_slices(x))
list(dataset.as_numpy_iterator())
[1, 2, 3, 4, 5, 6, 7, 8, 9]

tf.data.Dataset.interleave() является обобщением flat_map, так как flat_map создаёт тот же вывод, что и tf.data.Dataset.interleave(cycle_length=1)

from_generator

Просмотр исходного кода

@staticmethod
from_generator(
    generator,
    output_types=None,
    output_shapes=None,
    args=None,
    output_signature=None,
    name=None
)

Создаёт набор данных, элементы которого генерируются с помощью generator. (устаревшие аргументы)

Примечание: Текущая реализация Dataset.from_generator() использует tf.numpy_function и унаследовала те же ограничения. В частности, она требует, чтобы операции, связанные с набором данных и итератором, были размещены на устройстве в том же процессе, что и программа Python, которая вызвала Dataset.from_generator(). В частности, использование from_generator воспрепятствует использованию tf.data service для масштабирования обработки набора данных. Тело generator не будет сериализовано в GraphDef, и вам не следует использовать этот метод, если вам нужно сериализовать вашу модель и восстановить её в другой среде.

Аргумент generator должен быть вызываемым объектом, который возвращает объект, поддерживающий протокол iter() (например, функцию-генератор).

Элементы, генерируемые generator, должны быть совместимы либо с заданным аргументом output_signature, либо с заданными аргументами output_types и (необязательно) output_shapes, в зависимости от того, какой из них был указан.

Рекомендуемый способ вызова from_generator — использовать аргумент output_signature. В этом случае предполагается, что вывод будет состоять из объектов с классами, формами и типами, определёнными объектами tf.TypeSpec из аргумента output_signature.

def gen():
  ragged_tensor = tf.ragged.constant([[1, 2], [3]])
  yield 42, ragged_tensor

dataset = tf.data.Dataset.from_generator(
     gen,
     output_signature=(
         tf.TensorSpec(shape=(), dtype=tf.int32),
         tf.RaggedTensorSpec(shape=(2, None), dtype=tf.int32)))

list(dataset.take(1))
[(<tf.Tensor: shape=(), dtype=int32, numpy=42>,
<tf.RaggedTensor [[1, 2], [3]]>)]

Также существует устаревший способ вызова from_generator либо с использованием только аргумента output_types, либо вместе с аргументом output_shapes. В этом случае предполагается, что вывод функции будет состоять из объектов tf.Tensor с типами, определёнными output_types, и с формами, которые либо неизвестны, либо определены аргументом output_shapes.

Примечание: Если generator зависит от изменяемых глобальных переменных или другого внешнего состояния, имейте в виду, что среда выполнения может вызвать generator несколько раз (для поддержки повторения Dataset) и в любое время между вызовом Dataset.from_generator() и созданием первого элемента из генератора. Изменение глобальных переменных или внешнего состояния может привести к неопределённому поведению, поэтому рекомендуется кэшировать любое внешнее состояние в generator до вызова Dataset.from_generator().

Примечание: Хотя параметр output_signature позволяет передавать элементы Dataset, область действия Dataset.from_generator() должна быть ограничена логикой, которую нельзя выразить через операции tf.data. Использование операций tf.data внутри функции-генератора — это антипаттерн и может привести к постепенному росту потребления памяти.

from_tensor_slices

Просмотр исходного кода

@staticmethod
from_tensor_slices(
    tensors, name=None
)

Создаёт набор данных, элементы которого представляют собой слайсы заданных тензоров.

Заданные тензоры нарезаются по первому измерению. Эта операция сохраняет структуру входных тензоров, удаляя первое измерение каждого тензора и используя его как измерение набора данных. Все входные тензоры должны иметь одинаковый размер в своих первых измерениях.

# Slicing a 1D tensor produces scalar tensor elements.
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
list(dataset.as_numpy_iterator())
[1, 2, 3]

# Slicing a 2D tensor produces 1D tensor elements.
dataset = tf.data.Dataset.from_tensor_slices([[1, 2], [3, 4]])
list(dataset.as_numpy_iterator())
[array([1, 2], dtype=int32), array([3, 4], dtype=int32)]

# Slicing a tuple of 1D tensors produces tuple elements containing
# scalar tensors.
dataset = tf.data.Dataset.from_tensor_slices(([1, 2], [3, 4], [5, 6]))
list(dataset.as_numpy_iterator())
[(1, 3, 5), (2, 4, 6)]

# Dictionary structure is also preserved.
dataset = tf.data.Dataset.from_tensor_slices({"a": [1, 2], "b": [3, 4]})
list(dataset.as_numpy_iterator()) == [{'a': 1, 'b': 3},
                                      {'a': 2, 'b': 4}]
True

# Two tensors can be combined into one Dataset object.
features = tf.constant([[1, 3], [2, 1], [3, 3]]) # ==> 3x2 tensor
labels = tf.constant(['A', 'B', 'A']) # ==> 3x1 tensor
dataset = Dataset.from_tensor_slices((features, labels))
# Both the features and the labels tensors can be converted
# to a Dataset object separately and combined after.
features_dataset = Dataset.from_tensor_slices(features)
labels_dataset = Dataset.from_tensor_slices(labels)
dataset = Dataset.zip((features_dataset, labels_dataset))
# A batched feature and label set can be converted to a Dataset
# in similar fashion.
batched_features = tf.constant([[[1, 3], [2, 3]],
                                [[2, 1], [1, 2]],
                                [[3, 3], [3, 2]]], shape=(3, 2, 2))
batched_labels = tf.constant([['A', 'A'],
                              ['B', 'B'],
                              ['A', 'B']], shape=(3, 2, 1))
dataset = Dataset.from_tensor_slices((batched_features, batched_labels))
for element in dataset.as_numpy_iterator():
  print(element)
(array([[1, 3],
       [2, 3]], dtype=int32), array([[b'A'],
       [b'A']], dtype=object))
(array([[2, 1],
       [1, 2]], dtype=int32), array([[b'B'],
       [b'B']], dtype=object))
(array([[3, 3],
       [3, 2]], dtype=int32), array([[b'A'],
       [b'B']], dtype=object))

Обратите внимание, что если tensors содержит массив NumPy, а выполнение в режиме eager отключено, значения будут встроены в граф как одна или несколько операций tf.constant. Для больших наборов данных (> 1 ГБ) это может привести к потере памяти и проблемам с ограничениями на размер сериализации графа. Если tensors содержит один или несколько больших массивов NumPy, рассмотрите альтернативный подход, описанный в этом руководстве.

from_tensors

Просмотр исходного кода

@staticmethod
from_tensors(
    tensors, name=None
)

Создаёт набор данных с одним элементом, содержащим заданные тензоры.

from_tensors создаёт набор данных, содержащий только один элемент. Чтобы нарезать входной тензор на несколько элементов, используйте from_tensor_slices.

dataset = tf.data.Dataset.from_tensors([1, 2, 3])
list(dataset.as_numpy_iterator())
[array([1, 2, 3], dtype=int32)]
dataset = tf.data.Dataset.from_tensors(([1, 2, 3], 'A'))
list(dataset.as_numpy_iterator())
[(array([1, 2, 3], dtype=int32), b'A')]

# You can use `from_tensors` to produce a dataset which repeats
# the same example many times.
example = tf.constant([1,2,3])
dataset = tf.data.Dataset.from_tensors(example).repeat(2)
list(dataset.as_numpy_iterator())
[array([1, 2, 3], dtype=int32), array([1, 2, 3], dtype=int32)]

Обратите внимание, что если tensors содержит массив NumPy, а выполнение в режиме eager отключено, значения будут встроены в граф как одна или несколько операций tf.constant. Для больших наборов данных (> 1 ГБ) это может привести к потере памяти и проблемам с ограничениями на размер сериализации графа. Если tensors содержит один или несколько больших массивов NumPy, рассмотрите альтернативный подход, описанный в этом руководстве.

get_single_element

Просмотр исходного кода

get_single_element(
    name=None
)

Возвращает единственный элемент dataset.

Функция позволяет использовать набор данных tf.data.Dataset в бессостоятельном выражении «тензор-в-тензор-на-выходе», не создавая итератор. Это облегчает преобразование данных над тензорами, используя оптимизированную абстракцию tf.data.Dataset над ними.

Например, рассмотрим набор данных, который будет принимать на вход исходные признаки и возвращать обработанный признак вместе с его меткой.

def preprocessing_fn(raw_feature):
  # ... the raw_feature is preprocessed as per the use-case
  return feature

raw_features = ...  # input batch of BATCH_SIZE elements.
dataset = (tf.data.Dataset.from_tensor_slices(raw_features)
          .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
          .batch(BATCH_SIZE))

processed_features = dataset.get_single_element()

В приведенном выше примере тензор raw_features длины=BATCH_SIZE был преобразован в tf.data.Dataset. Далее каждый из raw_feature был преобразован с помощью preprocessing_fn, а обработанные признаки были объединены в одну партию. Конечный dataset содержит только один элемент, который представляет собой пакет всех обработанных признаков.

Примечание: dataset должен содержать только один элемент.

Теперь, вместо создания итератора для dataset и извлечения пакета признаков, используется функция tf.data.get_single_element(), которая пропускает процесс создания итератора и напрямую выводит пакет признаков.

Это может быть особенно полезно, когда преобразования тензоров выражены как операции tf.data.Dataset, и вы хотите использовать эти преобразования при обслуживании вашей модели.

Keras

model = ... # A pre-built or custom model

class PreprocessingModel(tf.keras.Model):
  def __init__(self, model):
    super().__init__(self)
    self.model = model

  @tf.function(input_signature=[...])
  def serving_fn(self, data):
    ds = tf.data.Dataset.from_tensor_slices(data)
    ds = ds.map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
    ds = ds.batch(batch_size=BATCH_SIZE)
    return tf.argmax(self.model(ds.get_single_element()), axis=-1)

preprocessing_model = PreprocessingModel(model)
your_exported_model_dir = ... # save the model to this path.
tf.saved_model.save(preprocessing_model, your_exported_model_dir,
              signatures={'serving_default': preprocessing_model.serving_fn}
              )

Estimator

В случае оценщиков, вам обычно нужно определить serving_input_fn, который потребовал бы обработки признаков моделью во время вывода.

def serving_input_fn():

  raw_feature_spec = ... # Spec for the raw_features
  input_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
      raw_feature_spec, default_batch_size=None)
  )
  serving_input_receiver = input_fn()
  raw_features = serving_input_receiver.features

  def preprocessing_fn(raw_feature):
    # ... the raw_feature is preprocessed as per the use-case
    return feature

  dataset = (tf.data.Dataset.from_tensor_slices(raw_features)
            .map(preprocessing_fn, num_parallel_calls=BATCH_SIZE)
            .batch(BATCH_SIZE))

  processed_features = dataset.get_single_element()

  # Please note that the value of `BATCH_SIZE` should be equal to
  # the size of the leading dimension of `raw_features`. This ensures
  # that `dataset` has only element, which is a pre-requisite for
  # using `dataset.get_single_element()`.

  return tf.estimator.export.ServingInputReceiver(
      processed_features, serving_input_receiver.receiver_tensors)

estimator = ... # A pre-built or custom estimator
estimator.export_saved_model(your_exported_model_dir, serving_input_fn)

group_by_window

Просмотреть исходный код

group_by_window(
    key_func, reduce_func, window_size=None, window_size_func=None, name=None
)

Группирует окна элементов по ключу и выполняет их уменьшение.

Это преобразование сопоставляет каждый последовательный элемент в наборе данных с ключом, используя key_func, и группирует элементы по ключу. Затем оно применяет reduce_func к не более чем window_size_func(key) элементам, соответствующим одному ключу. Все, кроме конечного окна для каждого ключа, будут содержать window_size_func(key) элементов; конечное окно может быть меньше.

Вы можете предоставить либо постоянный window_size или размер окна, определяемый ключом с помощью window_size_func.

dataset = tf.data.Dataset.range(10)
window_size = 5
key_func = lambda x: x%2
reduce_func = lambda key, dataset: dataset.batch(window_size)
dataset = dataset.group_by_window(
          key_func=key_func,
          reduce_func=reduce_func,
          window_size=window_size)
for elem in dataset.as_numpy_iterator():
  print(elem)
[0 2 4 6 8]
[1 3 5 7 9]

interleave

Просмотреть исходный код

interleave(
    map_func,
    cycle_length=None,
    block_length=None,
    num_parallel_calls=None,
    deterministic=None,
    name=None
)

Применяет map_func к этому набору данных и вставляет результаты.

Тип подписи:

def interleave(
  self: Dataset[T],
  map_func: Callable[[T], Dataset[S]]
) -> Dataset[S]

Например, вы можете использовать Dataset.interleave() для одновременной обработки многих входных файлов:

# Preprocess 4 files concurrently, and interleave blocks of 16 records
# from each file.
filenames = ["/var/data/file1.txt", "/var/data/file2.txt",
             "/var/data/file3.txt", "/var/data/file4.txt"]
dataset = tf.data.Dataset.from_tensor_slices(filenames)
def parse_fn(filename):
  return tf.data.Dataset.range(10)
dataset = dataset.interleave(lambda x:
    tf.data.TextLineDataset(x).map(parse_fn, num_parallel_calls=1),
    cycle_length=4, block_length=16)

Аргументы cycle_length и block_length управляют порядком, в котором генерируются элементы. cycle_length контролирует количество входных элементов, которые обрабатываются одновременно. Если вы зададите cycle_length значение 1, это преобразование будет обрабатывать один входной элемент за раз и даст идентичные результаты tf.data.Dataset.flat_map. В общем, это преобразование будет применять map_func к cycle_length входным элементам, открывать итераторы на возвращённых Dataset объектах, перебирать их, производя block_length последовательных элементов из каждого итератора и потребляя следующий входной элемент каждый раз, когда достигает конца итератора.

Например:

dataset = Dataset.range(1, 6)  # ==> [ 1, 2, 3, 4, 5 ]
# NOTE: New lines indicate "block" boundaries.
dataset = dataset.interleave(
    lambda x: Dataset.from_tensors(x).repeat(6),
    cycle_length=2, block_length=4)
list(dataset.as_numpy_iterator())
[1, 1, 1, 1,
 2, 2, 2, 2,
 1, 1,
 2, 2,
 3, 3, 3, 3,
 4, 4, 4, 4,
 3, 3,
 4, 4,
 5, 5, 5, 5,
 5, 5]

Примечание: Порядок элементов, возвращаемых этим преобразованием, детерминирован, до тех пор, пока map_func является чистой функцией и deterministic=True. Если map_func содержит какие-либо состоятельные операции, порядок доступа к этому состоянию не определён.

Производительность часто может быть улучшена, задав num_parallel_calls, чтобы interleave использовало несколько потоков для извлечения элементов. Если детерминизм не требуется, производительность также может быть улучшена, задав deterministic=False.

filenames = ["/var/data/file1.txt", "/var/data/file2.txt",
             "/var/data/file3.txt", "/var/data/file4.txt"]
dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.interleave(lambda x: tf.data.TFRecordDataset(x),
    cycle_length=4, num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False)

list_files

Просмотреть исходный код

@staticmethod
list_files(
    file_pattern, shuffle=None, seed=None, name=None
)

Набор данных всех файлов, соответствующих одному или нескольким шаблонам glob.

Аргумент file_pattern должен содержать небольшое количество шаблонов glob. Если имена файлов уже были отфильтрованы glob, используйте Dataset.from_tensor_slices(filenames) вместо этого, так как повторное применение glob к каждому имени файла с list_files может привести к плохой производительности с удалёнными хранилищами.

Примечание: По умолчанию этот метод возвращает имена файлов в произвольном порядке без определения. Передайте seed или shuffle=False для получения результатов в определённом порядке.

Пример:

Если у нас есть следующие файлы на файловой системе:

• /path/to/dir/a.txt
• /path/to/dir/b.py
• /path/to/dir/c.py

Если мы передадим "/path/to/dir/*.py" в качестве директории, набор данных будет производить:

• /path/to/dir/b.py
• /path/to/dir/c.py

map

Просмотреть исходный код

map(
    map_func, num_parallel_calls=None, deterministic=None, name=None
)

Применяет функцию map_func ко всем элементам этого набора данных.

Это преобразование применяет map_func к каждому элементу этого набора данных и возвращает новый набор данных, содержащий преобразованные элементы в том же порядке, что и во входном наборе. map_func может быть использован для изменения как значений, так и структуры элементов набора данных. Поддерживаемые структуры элементов документированы здесь.

Например, map можно использовать для добавления 1 к каждому элементу или для проекции подмножества компонентов элемента.

dataset = Dataset.range(1, 6)  # ==> [ 1, 2, 3, 4, 5 ]
dataset = dataset.map(lambda x: x + 1)
list(dataset.as_numpy_iterator())
[2, 3, 4, 5, 6]

Подпись входных данных для map_func определяется структурой каждого элемента в этом наборе данных.

dataset = Dataset.range(5)
# `map_func` takes a single argument of type `tf.Tensor` with the same
# shape and dtype.
result = dataset.map(lambda x: x + 1)

# Each element is a tuple containing two `tf.Tensor` objects.
elements = [(1, "foo"), (2, "bar"), (3, "baz")]
dataset = tf.data.Dataset.from_generator(
    lambda: elements, (tf.int32, tf.string))
# `map_func` takes two arguments of type `tf.Tensor`. This function
# projects out just the first component.
result = dataset.map(lambda x_int, y_str: x_int)
list(result.as_numpy_iterator())
[1, 2, 3]

# Each element is a dictionary mapping strings to `tf.Tensor` objects.
elements =  ([{"a": 1, "b": "foo"},
              {"a": 2, "b": "bar"},
              {"a": 3, "b": "baz"}])
dataset = tf.data.Dataset.from_generator(
    lambda: elements, {"a": tf.int32, "b": tf.string})
# `map_func` takes a single argument of type `dict` with the same keys
# as the elements.
result = dataset.map(lambda d: str(d["a"]) + d["b"])

Значение или значения, возвращаемые map_func, определяют структуру каждого элемента в возвращаемом наборе данных.

dataset = tf.data.Dataset.range(3)
# `map_func` returns two `tf.Tensor` objects.
def g(x):
  return tf.constant(37.0), tf.constant(["Foo", "Bar", "Baz"])
result = dataset.map(g)
result.element_spec
(TensorSpec(shape=(), dtype=tf.float32, name=None), TensorSpec(shape=(3,), dtype=tf.string, name=None))
# Python primitives, lists, and NumPy arrays are implicitly converted to
# `tf.Tensor`.
def h(x):
  return 37.0, ["Foo", "Bar"], np.array([1.0, 2.0], dtype=np.float64)
result = dataset.map(h)
result.element_spec
(TensorSpec(shape=(), dtype=tf.float32, name=None), TensorSpec(shape=(2,), dtype=tf.string, name=None), TensorSpec(shape=(2,), dtype=tf.float64, name=None))
# `map_func` can return nested structures.
def i(x):
  return (37.0, [42, 16]), "foo"
result = dataset.map(i)
result.element_spec
((TensorSpec(shape=(), dtype=tf.float32, name=None),
  TensorSpec(shape=(2,), dtype=tf.int32, name=None)),
 TensorSpec(shape=(), dtype=tf.string, name=None))

map_func может принимать в качестве аргументов и возвращать любые типы элементов набора данных.

Обратите внимание, что независимо от контекста определения map_func (режим eager или граф), tf.data отслеживает функцию и выполняет её как граф. Для использования Python-кода внутри функции у вас есть несколько вариантов:

1) Используйте AutoGraph для преобразования Python-кода в эквивалентное вычисление графа. Недостатком этого подхода является то, что AutoGraph может преобразовать некоторые, но не все, фрагменты Python-кода.

2) Используйте tf.py_function, что позволяет писать произвольный Python-код, но обычно приводит к худшей производительности, чем 1). Например:

d = tf.data.Dataset.from_tensor_slices(['hello', 'world'])
# transform a string tensor to upper case string using a Python function
def upper_case_fn(t: tf.Tensor):
  return t.numpy().decode('utf-8').upper()
d = d.map(lambda x: tf.py_function(func=upper_case_fn,
          inp=[x], Tout=tf.string))
list(d.as_numpy_iterator())
[b'HELLO', b'WORLD']

3) Используйте tf.numpy_function, которое также позволяет писать произвольный Python-код. Обратите внимание, что tf.py_function принимает tf.Tensor, в то время как tf.numpy_function принимает массивы NumPy и возвращает только массивы NumPy. Например:

d = tf.data.Dataset.from_tensor_slices(['hello', 'world'])
def upper_case_fn(t: np.ndarray):
  return t.decode('utf-8').upper()
d = d.map(lambda x: tf.numpy_function(func=upper_case_fn,
          inp=[x], Tout=tf.string))
list(d.as_numpy_iterator())
[b'HELLO', b'WORLD']

Обратите внимание, что использование tf.numpy_function и tf.py_function в целом исключает возможность параллельного выполнения пользовательских преобразований (из-за блокировки интерпретатора Python).

Производительность часто можно улучшить, установив num_parallel_calls, чтобы map использовал несколько потоков для обработки элементов. Если порядок не важен, производительность также можно улучшить, установив deterministic=False.

dataset = Dataset.range(1, 6)  # ==> [ 1, 2, 3, 4, 5 ]
dataset = dataset.map(lambda x: x + 1,
    num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False)

Порядок элементов, возвращаемых этим преобразованием, является детерминированным, если deterministic=True. Если map_func содержит операции с состоянием и num_parallel_calls > 1, порядок доступа к этому состоянию не определён, поэтому значения элементов вывода могут не быть детерминированными независимо от значения флага deterministic.

options

Просмотреть исходный код

options()

Возвращает параметры этого набора данных и его входов.

padded_batch

Просмотреть исходный код

padded_batch(
    batch_size,
    padded_shapes=None,
    padding_values=None,
    drop_remainder=False,
    name=None
)

Объединяет последовательные элементы этого набора данных в запакованные пакеты.

Это преобразование объединяет несколько последовательных элементов входного набора данных в один элемент.

Как и tf.data.Dataset.batch, компоненты результирующего элемента будут иметь дополнительное внешнее измерение, которое будет batch_size (или N % batch_size для последнего элемента, если batch_size не делит количество входных элементов N равномерно и drop_remainder равно False). Если ваша программа зависит от того, чтобы пакеты имели одинаковое внешнее измерение, вы должны установить аргумент drop_remainder в True для предотвращения создания меньшего пакета.

В отличие от tf.data.Dataset.batch, входные элементы для формирования пакетов могут иметь разные формы, и это преобразование будет дополнять каждый компонент до соответствующей формы в padded_shapes. Аргумент padded_shapes определяет результирующую форму для каждого измерения каждого компонента в элементе вывода:

• Если размерность является константой, компонент будет дополнен до этой длины в этом измерении.
• Если размерность неизвестна, компонент будет дополнен до максимальной длины всех элементов в этом измерении.

A = (tf.data.Dataset
     .range(1, 5, output_type=tf.int32)
     .map(lambda x: tf.fill([x], x)))
# Pad to the smallest per-batch size that fits all elements.
B = A.padded_batch(2)
for element in B.as_numpy_iterator():
  print(element)
[[1 0]
 [2 2]]
[[3 3 3 0]
 [4 4 4 4]]
# Pad to a fixed size.
C = A.padded_batch(2, padded_shapes=5)
for element in C.as_numpy_iterator():
  print(element)
[[1 0 0 0 0]
 [2 2 0 0 0]]
[[3 3 3 0 0]
 [4 4 4 4 0]]
# Pad with a custom value.
D = A.padded_batch(2, padded_shapes=5, padding_values=-1)
for element in D.as_numpy_iterator():
  print(element)
[[ 1 -1 -1 -1 -1]
 [ 2  2 -1 -1 -1]]
[[ 3  3  3 -1 -1]
 [ 4  4  4  4 -1]]
# Components of nested elements can be padded independently.
elements = [([1, 2, 3], [10]),
            ([4, 5], [11, 12])]
dataset = tf.data.Dataset.from_generator(
    lambda: iter(elements), (tf.int32, tf.int32))
# Pad the first component of the tuple to length 4, and the second
# component to the smallest size that fits.
dataset = dataset.padded_batch(2,
    padded_shapes=([4], [None]),
    padding_values=(-1, 100))
list(dataset.as_numpy_iterator())
[(array([[ 1,  2,  3, -1], [ 4,  5, -1, -1]], dtype=int32),
  array([[ 10, 100], [ 11,  12]], dtype=int32))]
# Pad with a single value and multiple components.
E = tf.data.Dataset.zip((A, A)).padded_batch(2, padding_values=-1)
for element in E.as_numpy_iterator():
  print(element)
(array([[ 1, -1],
       [ 2,  2]], dtype=int32), array([[ 1, -1],
       [ 2,  2]], dtype=int32))
(array([[ 3,  3,  3, -1],
       [ 4,  4,  4,  4]], dtype=int32), array([[ 3,  3,  3, -1],
       [ 4,  4,  4,  4]], dtype=int32))

См. также tf.data.experimental.dense_to_sparse_batch, который объединяет элементы, которые могут иметь разные формы, в tf.sparse.SparseTensor.

prefetch

Просмотреть исходный код

prefetch(
    buffer_size, name=None
)

Создаёт Dataset для предварительной обработки элементов из этого набора данных.

Большинство конвейеров входных данных наборов данных должны заканчиваться вызовом prefetch. Это позволяет подготовить последующие элементы, пока текущий обрабатывается. Это часто улучшает задержку и пропускную способность, за счёт использования дополнительной памяти для хранения предварительно обработанных элементов.

Примечание: Как и другие методы Dataset, префич обрабатывает элементы входного набора данных. Он не имеет понятия об примерах vs. пакетах. examples.prefetch(2) предварительно обработает два элемента (2 примера), в то время как examples.batch(20).prefetch(2) предварительно обработает 2 элемента (2 пакета по 20 примеров каждый).

dataset = tf.data.Dataset.range(3)
dataset = dataset.prefetch(2)
list(dataset.as_numpy_iterator())
[0, 1, 2]

random

Просмотреть исходный код

@staticmethod
random(
    seed=None, name=None
)

Создаёт набор Dataset псевдослучайных значений.

Набор данных генерирует последовательность целых значений, равномерно распределённых.

ds1 = tf.data.Dataset.random(seed=4).take(10)
ds2 = tf.data.Dataset.random(seed=4).take(10)
print(list(ds2.as_numpy_iterator())==list(ds2.as_numpy_iterator()))
True

range

Просмотреть исходный код

@staticmethod
range(
    *args, **kwargs
)

Создаёт набор данных с диапазоном значений с заданным шагом.

list(Dataset.range(5).as_numpy_iterator())
[0, 1, 2, 3, 4]
list(Dataset.range(2, 5).as_numpy_iterator())
[2, 3, 4]
list(Dataset.range(1, 5, 2).as_numpy_iterator())
[1, 3]
list(Dataset.range(1, 5, -2).as_numpy_iterator())
[]
list(Dataset.range(5, 1).as_numpy_iterator())
[]
list(Dataset.range(5, 1, -2).as_numpy_iterator())
[5, 3]
list(Dataset.range(2, 5, output_type=tf.int32).as_numpy_iterator())
[2, 3, 4]
list(Dataset.range(1, 5, 2, output_type=tf.float32).as_numpy_iterator())
[1.0, 3.0]

reduce

Просмотреть исходный код

reduce(
    initial_state, reduce_func, name=None
)

Сводит набор данных к одному элементу.

Преобразование последовательно вызывает reduce_func для каждого элемента входного набора данных до тех пор, пока набор данных не будет исчерпан, агрегируя информацию во внутреннем состоянии. Аргумент initial_state используется для начального состояния, а конечное состояние возвращается в качестве результата.

tf.data.Dataset.range(5).reduce(np.int64(0), lambda x, _: x + 1).numpy()
5
tf.data.Dataset.range(5).reduce(np.int64(0), lambda x, y: x + y).numpy()
10

rejection_resample

Просмотреть исходный код

rejection_resample(
    class_func, target_dist, initial_dist=None, seed=None, name=None
)

Преобразование, перевыбирающее набор данных в соответствии с целевым распределением.

Рассмотрим пример, где набор данных с исходным распределением данных init_dist должен быть перевыбран в набор данных с распределением target_dist.

initial_dist = [0.6, 0.4]
num_classes = len(initial_dist)
num_samples = 1000
data_np = np.random.choice(num_classes, num_samples, p=initial_dist)
dataset = tf.data.Dataset.from_tensor_slices(data_np)

Значение x будет близко к {0: 50000, 1: 50000} в соответствии с распределением initial_dist.

target_dist = [0.5, 0.5]
resampled_dataset = dataset.rejection_resample(
   class_func=lambda x: x,
   target_dist=target_dist,
   initial_dist=initial_dist)
resampled_dataset = resampled_dataset.map(
    lambda class_func_result, data: data)

Распределение значений классов в resampled_distribution теперь будет близко к целевому распределению.

repeat

Просмотреть исходный код

repeat(
    count=None, name=None
)

Повторяет этот набор данных, чтобы каждое исходное значение отображалось count раз.

dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3])
dataset = dataset.repeat(3)
list(dataset.as_numpy_iterator())
[1, 2, 3, 1, 2, 3, 1, 2, 3]

Примечание: Если входной набор данных зависит от глобального состояния (например, генератора случайных чисел) или его вывод не детерминирован (например, из-за предыдущих shuffle), различные повторения могут генерировать различные элементы.

sample_from_datasets

Просмотреть исходный код

@staticmethod
sample_from_datasets(
    datasets, weights=None, seed=None, stop_on_empty_dataset=False
)

Выборка элементов случайным образом из наборов данных в datasets.

Создаёт набор данных путём чередования элементов datasets с вероятностью weight[i] выбора элемента из набора данных i. Выборка происходит без возвращения. Например, предположим, что у нас есть 2 набора данных:

dataset1 = tf.data.Dataset.range(0, 3)
dataset2 = tf.data.Dataset.range(100, 103)

Предположим, что мы выбираем из этих двух наборов данных с указанными весами:

sample_dataset = tf.data.Dataset.sample_from_datasets(
    [dataset1, dataset2], weights=[0.5, 0.5])

Один возможный результат элементов в sample_dataset:

print(list(sample_dataset.as_numpy_iterator()))
# [100, 0, 1, 101, 2, 102]

scan

Просмотреть исходный код

scan(
    initial_state, scan_func, name=None
)

Преобразование, которое применяет функцию к входному набору данных.

Это преобразование является состоятельным аналогом tf.data.Dataset.map. Помимо применения scan_func к элементам входного набора данных, scan() накапливает один или несколько тензоров состояния, начальные значения которых initial_state.

dataset = tf.data.Dataset.range(10)
initial_state = tf.constant(0, dtype=tf.int64)
scan_func = lambda state, i: (state + i, state + i)
dataset = dataset.scan(initial_state=initial_state, scan_func=scan_func)
list(dataset.as_numpy_iterator())
[0, 1, 3, 6, 10, 15, 21, 28, 36, 45]

shard

Просмотреть исходный код

shard(
    num_shards, index, name=None
)

Создает набор данных Dataset, включающий только 1/num_shards этого набора данных.

shard детерминировано. Набор данных, созданный A.shard(n, i), будет содержать все элементы A, для которых индекс mod n = i.

A = tf.data.Dataset.range(10)
B = A.shard(num_shards=3, index=0)
list(B.as_numpy_iterator())
[0, 3, 6, 9]
C = A.shard(num_shards=3, index=1)
list(C.as_numpy_iterator())
[1, 4, 7]
D = A.shard(num_shards=3, index=2)
list(D.as_numpy_iterator())
[2, 5, 8]

Этот оператор набора данных очень полезен при распределенном обучении, так как позволяет каждому рабочему процессу читать уникальный подмножество.

При чтении одного входного файла, вы можете разделить элементы следующим образом:

d = tf.data.TFRecordDataset(input_file)
d = d.shard(num_workers, worker_index)
d = d.repeat(num_epochs)
d = d.shuffle(shuffle_buffer_size)
d = d.map(parser_fn, num_parallel_calls=num_map_threads)

Важные замечания:

• Убедитесь, что вы разделили данные перед использованием любого оператора случайного выбора (например, shuffle).
• В целом лучше, если оператор shard используется на ранней стадии обработки набора данных. Например, при чтении из набора файлов TFRecord, разделите данные перед преобразованием набора данных в образцы. Это предотвращает чтение каждого файла на каждом рабочем процессе. Следующее — пример эффективной стратегии разбиения в полном конвейере:

d = Dataset.list_files(pattern)
d = d.shard(num_workers, worker_index)
d = d.repeat(num_epochs)
d = d.shuffle(shuffle_buffer_size)
d = d.interleave(tf.data.TFRecordDataset,
                 cycle_length=num_readers, block_length=1)
d = d.map(parser_fn, num_parallel_calls=num_map_threads)

shuffle

Просмотреть исходный код

shuffle(
    buffer_size, seed=None, reshuffle_each_iteration=None, name=None
)

Случайным образом перемешивает элементы этого набора данных.

Этот набор данных заполняет буфер buffer_size элементами, затем случайным образом выбирает элементы из этого буфера, заменяя выбранные элементы новыми элементами. Для идеального перемешивания требуется размер буфера, больший или равный полному размеру набора данных.

Например, если ваш набор данных содержит 10 000 элементов, но buffer_size задано значением 1000, тогда shuffle изначально выберет случайный элемент только из первых 1000 элементов в буфере. После выбора элемента его место в буфере заменяется следующим (т.е. 1001-м) элементом, поддерживая буфер из 1000 элементов.

reshuffle_each_iteration управляет тем, должна ли последовательность перемешивания быть разной для каждой эпохи. В TF 1.X, стандартный способ создания эпох был через преобразование repeat:

dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
dataset = dataset.repeat(2)
# [1, 0, 2, 1, 2, 0]

dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=False)
dataset = dataset.repeat(2)
# [1, 0, 2, 1, 0, 2]

В TF 2.0, объекты tf.data.Dataset являются Python-итерируемыми объектами, что позволяет создавать эпохи также с помощью Python-итерации:

dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
list(dataset.as_numpy_iterator())
# [1, 0, 2]
list(dataset.as_numpy_iterator())
# [1, 2, 0]

dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=False)
list(dataset.as_numpy_iterator())
# [1, 0, 2]
list(dataset.as_numpy_iterator())
# [1, 0, 2]

skip

Просмотреть исходный код

skip(
    count, name=None
)

Создает набор данных Dataset, пропускающий count элементов из этого набора данных.

dataset = tf.data.Dataset.range(10)
dataset = dataset.skip(7)
list(dataset.as_numpy_iterator())
[7, 8, 9]

snapshot

Просмотреть исходный код

snapshot(
    path,
    compression='AUTO',
    reader_func=None,
    shard_func=None,
    name=None
)

API для сохранения выходных данных входного набора данных.

API snapshot позволяет пользователям прозрачно сохранять выходные данные своего конвейера предобработки на диске и материализовать предобработанные данные в другом запуске обучения.

Этот API позволяет объединить повторяющиеся шаги предобработки и использовать уже обработанные данные, жертвуя хранением на диске и пропускной способностью сети для высвобождения большего количества ценных ресурсов ЦП и вычислительной мощности ускорителя.

https://github.com/tensorflow/community/blob/master/rfcs/20200107-tf-data-snapshot.md содержит подробную документацию по проектированию этой функции.

Пользователи могут указать различные параметры для управления поведением snapshot, включая чтение и запись снимков, передавая пользовательские функции параметрам reader_func и shard_func.

shard_func — это пользовательская функция, которая отображает входные элементы в фрагменты снимков.

Пользователям может потребоваться указать эту функцию для управления тем, как файлы снимков записываются на диск. Ниже приведен пример потенциальной реализации shard_func.

dataset = ...
dataset = dataset.enumerate()
dataset = dataset.snapshot("/path/to/snapshot/dir",
    shard_func=lambda x, y: x % NUM_SHARDS, ...)
dataset = dataset.map(lambda x, y: y)

reader_func — это пользовательская функция, которая принимает один аргумент: (1) набор данных наборов данных, каждый из которых представляет «разделение» элементов исходного набора данных. Мощность входного набора данных соответствует количеству фрагментов, указанных в shard_func (см. выше). Функция должна возвращать набор данных элементов исходного набора данных.

Пользователи могут указать эту функцию для управления тем, как файлы снимков должны считываться с диска, включая количество перемешивания и параллелизма.

Вот пример стандартной функции чтения, которую может определить пользователь. Эта функция позволяет как перемешивать набор данных, так и параллельно читать наборы данных:

def user_reader_func(datasets):
  # shuffle the datasets splits
  datasets = datasets.shuffle(NUM_CORES)
  # read datasets in parallel and interleave their elements
  return datasets.interleave(lambda x: x, num_parallel_calls=AUTOTUNE)

dataset = dataset.snapshot("/path/to/snapshot/dir",
    reader_func=user_reader_func)

По умолчанию snapshot параллелизует чтение количеством ядер, доступных в системе, но не будет пытаться перемешивать данные.

take

Просмотреть исходный код

take(
    count, name=None
)

Создаёт Dataset с не более чем count элементами из этого набора данных.

dataset = tf.data.Dataset.range(10)
dataset = dataset.take(3)
list(dataset.as_numpy_iterator())
[0, 1, 2]

take_while

Просмотреть исходный код

take_while(
    predicate, name=None
)

Преобразование, останавливающее итерацию набора данных на основе predicate.

dataset = tf.data.Dataset.range(10)
dataset = dataset.take_while(lambda x: x < 5)
list(dataset.as_numpy_iterator())
[0, 1, 2, 3, 4]

unbatch

Просмотреть исходный код

unbatch(
    name=None
)

Разбивает элементы набора данных на несколько элементов.

Например, если элементы набора данных имеют форму [B, a0, a1, ...], где B может отличаться для каждого входного элемента, то для каждого элемента в наборе данных, набор данных без объединения будет содержать B последовательных элементов формы [a0, a1, ...].

elements = [ [1, 2, 3], [1, 2], [1, 2, 3, 4] ]
dataset = tf.data.Dataset.from_generator(lambda: elements, tf.int64)
dataset = dataset.unbatch()
list(dataset.as_numpy_iterator())
[1, 2, 3, 1, 2, 1, 2, 3, 4]

Примечание: unbatch требует копирования данных для разделения объединённого тензора на меньшие, не объединённые тензоры. При оптимизации производительности старайтесь избегать ненужного использования unbatch.

unique

Просмотреть исходный код

unique(
    name=None
)

Преобразование, отбрасывающее дублирующие элементы из Dataset.

Используйте это преобразование, чтобы создать набор данных, содержащий один экземпляр каждого уникального элемента на входе. Например:

dataset = tf.data.Dataset.from_tensor_slices([1, 37, 2, 37, 2, 1])
dataset = dataset.unique()
sorted(list(dataset.as_numpy_iterator()))
[1, 2, 37]

Примечание: Это преобразование поддерживает только наборы данных, которые помещаются в память и имеют элементы типа tf.int32, tf.int64 или tf.string.

window

Просмотреть исходный код

window(
    size, shift=None, stride=1, drop_remainder=False, name=None
)

Возвращает набор данных "окон".

Каждое "окно" — это набор данных, содержащий подмножество элементов входного набора данных. Это конечные наборы данных размером size (или возможно меньше, если нет достаточно входных элементов для заполнения окна и drop_remainder оценивается как False).

Например:

dataset = tf.data.Dataset.range(7).window(3)
for window in dataset:
  print(window)
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>
<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int64, name=None)>

Поскольку окна являются наборами данных, их можно перебирать:

for window in dataset:
  print([item.numpy() for item in window])
[0, 1, 2]
[3, 4, 5]
[6]

Сдвиг

Аргумент shift определяет количество входных элементов, которые нужно сдвигать между началом каждого окна. Если номера окон и элементов начинаются с 0, первый элемент в окне k будет элементом k * shift входного набора данных. В частности, первый элемент первого окна всегда будет первым элементом входного набора данных.

dataset = tf.data.Dataset.range(7).window(3, shift=1,
                                          drop_remainder=True)
for window in dataset:
  print(list(window.as_numpy_iterator()))
[0, 1, 2]
[1, 2, 3]
[2, 3, 4]
[3, 4, 5]
[4, 5, 6]

Шаг

Аргумент stride определяет шаг между входными элементами внутри окна.

Вложенные элементы

Когда преобразование window применяется к набору данных, элементы которого являются вложенными структурами, оно генерирует набор данных, где элементы имеют ту же вложенную структуру, но каждый лист заменён окном. Другими словами, вложение применяется вне окон, а не внутри них.

Подпись типа:

def window(
    self: Dataset[Nest[T]], ...
) -> Dataset[Nest[Dataset[T]]]

Применение window к набору данных кортежей даёт кортеж окон:

dataset = tf.data.Dataset.from_tensor_slices(([1, 2, 3, 4, 5],
                                              [6, 7, 8, 9, 10]))
dataset = dataset.window(2)
windows = next(iter(dataset))
windows
(<...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int32, name=None)>,
 <...Dataset element_spec=TensorSpec(shape=(), dtype=tf.int32, name=None)>)

def to_numpy(ds):
  return list(ds.as_numpy_iterator())

for windows in dataset:
  print(to_numpy(windows[0]), to_numpy(windows[1]))
[1, 2] [6, 7]
[3, 4] [8, 9]
[5] [10]

Применение window к набору данных словарей даёт словарь Datasets:

dataset = tf.data.Dataset.from_tensor_slices({'a': [1, 2, 3],
                                              'b': [4, 5, 6],
                                              'c': [7, 8, 9]})
dataset = dataset.window(2)
def to_numpy(ds):
  return list(ds.as_numpy_iterator())

for windows in dataset:
  print(tf.nest.map_structure(to_numpy, windows))
{'a': [1, 2], 'b': [4, 5], 'c': [7, 8]}
{'a': [3], 'b': [6], 'c': [9]}

Сглаживание набора данных окон

Dataset.flat_map и Dataset.interleave методы могут быть использованы для сглаживания набора данных окон в один набор данных.

Аргументом flat_map является функция, которая берёт элемент из набора данных и возвращает Dataset. flat_map последовательно объединяет полученные наборы данных.

Например, чтобы преобразовать каждое окно в плотный тензор:

size = 3
dataset = tf.data.Dataset.range(7).window(size, shift=1,
                                          drop_remainder=True)
batched = dataset.flat_map(lambda x:x.batch(3))
for batch in batched:
  print(batch.numpy())
[0 1 2]
[1 2 3]
[2 3 4]
[3 4 5]
[4 5 6]

with_options

Просмотреть исходный код

with_options(
    options, name=None
)

Возвращает новый tf.data.Dataset с заданными параметрами.

Параметры "глобальны" в том смысле, что они применяются ко всему набору данных. Если параметры устанавливаются несколько раз, они объединяются, пока разные параметры не используют разные значения, отличные от по умолчанию.

ds = tf.data.Dataset.range(5)
ds = ds.interleave(lambda x: tf.data.Dataset.range(5),
                   cycle_length=3,
                   num_parallel_calls=3)
options = tf.data.Options()
# This will make the interleave order non-deterministic.
options.deterministic = False
ds = ds.with_options(options)

zip

Просмотреть исходный код

@staticmethod
zip(
    datasets, name=None
)

Создаёт Dataset, объединяя заданные наборы данных.

Этот метод имеет схожую семантику с встроенной функцией zip() в Python, с основным отличием, заключающимся в том, что аргумент datasets может быть (вложенной) структурой объектов Dataset. Поддерживаемые механизмы вложенности документированы здесь.

# The nested structure of the `datasets` argument determines the
# structure of elements in the resulting dataset.
a = tf.data.Dataset.range(1, 4)  # ==> [ 1, 2, 3 ]
b = tf.data.Dataset.range(4, 7)  # ==> [ 4, 5, 6 ]
ds = tf.data.Dataset.zip((a, b))
list(ds.as_numpy_iterator())
[(1, 4), (2, 5), (3, 6)]
ds = tf.data.Dataset.zip((b, a))
list(ds.as_numpy_iterator())
[(4, 1), (5, 2), (6, 3)]

# The `datasets` argument may contain an arbitrary number of datasets.
c = tf.data.Dataset.range(7, 13).batch(2)  # ==> [ [7, 8],
                                           #       [9, 10],
                                           #       [11, 12] ]
ds = tf.data.Dataset.zip((a, b, c))
for element in ds.as_numpy_iterator():
  print(element)
(1, 4, array([7, 8]))
(2, 5, array([ 9, 10]))
(3, 6, array([11, 12]))

# The number of elements in the resulting dataset is the same as
# the size of the smallest dataset in `datasets`.
d = tf.data.Dataset.range(13, 15)  # ==> [ 13, 14 ]
ds = tf.data.Dataset.zip((a, d))
list(ds.as_numpy_iterator())
[(1, 13), (2, 14)]

__bool__

Просмотреть исходный код

__bool__()

__iter__

Просмотреть исходный код

__iter__()

Создаёт итератор для элементов этого набора данных.

Возвращаемый итератор реализует протокол Python Iterator.

__len__

Просмотреть исходный код

__len__()

Возвращает длину набора данных, если она известна и конечна.

Этот метод требует, чтобы вы выполнялись в режиме Eager и длина набора данных была известна и не бесконечна. Когда длина может быть неизвестной или бесконечной, или если вы работаете в режиме графа, используйте tf.data.Dataset.cardinality вместо этого.

__nonzero__

Просмотреть исходный код

__nonzero__()