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

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

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

  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)