tf.data.experimental.get_single_element

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

tf.data.experimental.get_single_element(
    dataset
)

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

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

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 = tf.data.experimental.get_single_element(dataset)

В приведенном выше примере тензор raw_features длины BATCH_SIZE был преобразован в tf.data.Dataset. Затем каждый из raw_feature был сопоставлен с помощью preprocessing_fn, а обработанные признаки были объединены в один пакет. Итоговый dataset содержит только один элемент — пакет со всеми обработанными признаками.

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

Теперь вместо создания итератора для dataset и извлечения пакета признаков используется функция tf.data.experimental.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(tf.data.experimental.get_single_element(ds)),
      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})