TensorFlow 2.x en TFX

TensorFlow 2.0 se lanzó en 2019 , con una estrecha integración de Keras , ejecución entusiasta de forma predeterminada y ejecución de funciones Pythonic , entre otras características y mejoras nuevas .

Esta guía proporciona una descripción técnica completa de TF 2.x en TFX.

¿Qué versión usar?

TFX es compatible con TensorFlow 2.x y las API de alto nivel que existían en TensorFlow 1.x (particularmente Estimators) siguen funcionando.

Iniciar nuevos proyectos en TensorFlow 2.x

Dado que TensorFlow 2.x conserva las capacidades de alto nivel de TensorFlow 1.x, no hay ninguna ventaja en usar la versión anterior en nuevos proyectos, incluso si no planea usar las nuevas funciones.

Por lo tanto, si está comenzando un nuevo proyecto TFX, le recomendamos que use TensorFlow 2.x. Es posible que desee actualizar su código más tarde a medida que esté disponible la compatibilidad completa con Keras y otras características nuevas, y el alcance de los cambios será mucho más limitado si comienza con TensorFlow 2.x, en lugar de intentar actualizar desde TensorFlow 1.x en el futuro.

Conversión de proyectos existentes a TensorFlow 2.x

El código escrito para TensorFlow 1.x es en gran medida compatible con TensorFlow 2.x y seguirá funcionando en TFX.

Sin embargo, si desea aprovechar las mejoras y las nuevas funciones a medida que estén disponibles en TF 2.x, puede seguir las instrucciones para migrar a TF 2.x.

Estimador

La API de Estimator se ha conservado en TensorFlow 2.x, pero no es el foco de las nuevas características y el desarrollo. El código escrito en TensorFlow 1.x o 2.x con Estimators seguirá funcionando como se esperaba en TFX.

Aquí hay un ejemplo de TFX de extremo a extremo usando Estimador puro: Ejemplo de taxi (Estimador)

Keras con model_to_estimator

Los modelos de Keras se pueden envolver con la función tf.keras.estimator.model_to_estimator , que les permite trabajar como si fueran Estimadores. Para usar esto:

1. Construya un modelo de Keras.
2. Pase el modelo compilado a model_to_estimator .
3. Use el resultado de model_to_estimator en Trainer, de la forma en que normalmente usaría un Estimator.

# Build a Keras model.
def _keras_model_builder():
  """Creates a Keras model."""
  ...

  model = tf.keras.Model(inputs=inputs, outputs=output)
  model.compile()

  return model


# Write a typical trainer function
def trainer_fn(trainer_fn_args, schema):
  """Build the estimator, using model_to_estimator."""
  ...

  # Model to estimator
  estimator = tf.keras.estimator.model_to_estimator(
      keras_model=_keras_model_builder(), config=run_config)

  return {
      'estimator': estimator,
      ...
  }

Aparte del archivo del módulo de usuario de Trainer, el resto de la canalización permanece sin cambios.

Keras nativo (es decir, Keras sin model_to_estimator )

Nota: El soporte completo para todas las funciones de Keras está en progreso; en la mayoría de los casos, Keras en TFX funcionará como se esperaba. Todavía no funciona con funciones dispersas para FeatureColumns.

Ejemplos y Colaboración

Aquí hay varios ejemplos con Keras nativos:

• Penguin ( archivo de módulo ): ejemplo de extremo a extremo de 'Hola mundo'.
• MNIST ( archivo de módulo ): imagen y ejemplo de extremo a extremo de TFLite.
• Taxi ( archivo de módulo ): ejemplo de extremo a extremo con uso avanzado de Transform.

También tenemos un Keras Colab por componente.

Componentes TFX

Las siguientes secciones explican cómo los componentes TFX relacionados admiten Keras nativo.

Transformar

Transform actualmente tiene soporte experimental para modelos Keras.

El componente Transform en sí mismo se puede usar para Keras nativos sin cambios. La definición de preprocessing_fn sigue siendo la misma, usando TensorFlow y tf.Transform ops.

La función de servicio y la función de evaluación se cambian para Keras nativo. Los detalles se discutirán en las siguientes secciones para Entrenador y Evaluador.

Nota: Las transformaciones dentro de preprocessing_fn no se pueden aplicar a la función de etiqueta para entrenamiento o evaluación.

Entrenador

Para configurar Keras nativo, se debe configurar GenericExecutor para que el componente Trainer reemplace el ejecutor predeterminado basado en Estimator. Para obtener más información, consulte aquí .

Archivo del módulo Keras con Transform

El archivo del módulo de entrenamiento debe contener un run_fn que será llamado por GenericExecutor , un run_fn típico de Keras se vería así:

def run_fn(fn_args: TrainerFnArgs):
  """Train the model based on given args.

  Args:
    fn_args: Holds args used to train the model as name/value pairs.
  """
  tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)

  # Train and eval files contains transformed examples.
  # _input_fn read dataset based on transformed schema from tft.
  train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor,
                            tf_transform_output.transformed_metadata.schema)
  eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor,
                           tf_transform_output.transformed_metadata.schema)

  model = _build_keras_model()

  model.fit(
      train_dataset,
      steps_per_epoch=fn_args.train_steps,
      validation_data=eval_dataset,
      validation_steps=fn_args.eval_steps)

  signatures = {
      'serving_default':
          _get_serve_tf_examples_fn(model,
                                    tf_transform_output).get_concrete_function(
                                        tf.TensorSpec(
                                            shape=[None],
                                            dtype=tf.string,
                                            name='examples')),
  }
  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

En el run_fn anterior, se necesita una firma de publicación al exportar el modelo entrenado para que el modelo pueda tomar ejemplos sin procesar para la predicción. Una función de servicio típica se vería así:

def _get_serve_tf_examples_fn(model, tf_transform_output):
  """Returns a function that parses a serialized tf.Example."""

  # the layer is added as an attribute to the model in order to make sure that
  # the model assets are handled correctly when exporting.
  model.tft_layer = tf_transform_output.transform_features_layer()

  @tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    """Returns the output to be used in the serving signature."""
    feature_spec = tf_transform_output.raw_feature_spec()
    feature_spec.pop(_LABEL_KEY)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)

    transformed_features = model.tft_layer(parsed_features)

    return model(transformed_features)

  return serve_tf_examples_fn

En la función de servicio anterior, las transformaciones tf.Transform deben aplicarse a los datos sin procesar para la inferencia, utilizando la capa tft.TransformFeaturesLayer . El _serving_input_receiver_fn anterior que se requería para Estimators ya no se necesitará con Keras.

Archivo del módulo Keras sin transformación

Esto es similar al archivo de módulo que se muestra arriba, pero sin las transformaciones:

def _get_serve_tf_examples_fn(model, schema):

  @tf.function
  def serve_tf_examples_fn(serialized_tf_examples):
    feature_spec = _get_raw_feature_spec(schema)
    feature_spec.pop(_LABEL_KEY)
    parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
    return model(parsed_features)

  return serve_tf_examples_fn


def run_fn(fn_args: TrainerFnArgs):
  schema = io_utils.parse_pbtxt_file(fn_args.schema_file, schema_pb2.Schema())

  # Train and eval files contains raw examples.
  # _input_fn reads the dataset based on raw data schema.
  train_dataset = _input_fn(fn_args.train_files, fn_args.data_accessor, schema)
  eval_dataset = _input_fn(fn_args.eval_files, fn_args.data_accessor, schema)

  model = _build_keras_model()

  model.fit(
      train_dataset,
      steps_per_epoch=fn_args.train_steps,
      validation_data=eval_dataset,
      validation_steps=fn_args.eval_steps)

  signatures = {
      'serving_default':
          _get_serve_tf_examples_fn(model, schema).get_concrete_function(
              tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')),
  }
  model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

tf.distribuir.Estrategia

En este momento, TFX solo admite estrategias de un solo trabajador (por ejemplo, MirroredStrategy , OneDeviceStrategy ).

Para usar una estrategia de distribución, cree un tf.distribute.Strategy apropiado y mueva la creación y compilación del modelo Keras dentro del alcance de una estrategia.

Por ejemplo, reemplace model = _build_keras_model() con:

  mirrored_strategy = tf.distribute.MirroredStrategy()
  with mirrored_strategy.scope():
    model = _build_keras_model()

  # Rest of the code can be unchanged.
  model.fit(...)

Para verificar el dispositivo (CPU/GPU) utilizado por MirroredStrategy , habilite el registro de tensorflow de nivel de información:

import logging
logging.getLogger("tensorflow").setLevel(logging.INFO)

y debería poder ver Using MirroredStrategy with devices (...) en el registro.

Nota: La variable de entorno TF_FORCE_GPU_ALLOW_GROWTH=true podría ser necesaria para un problema de GPU sin memoria. Para obtener más información, consulte la guía de GPU de tensorflow .

evaluador

En TFMA v0.2x, ModelValidator y Evaluator se combinaron en un único componente Evaluator nuevo . El nuevo componente Evaluator puede realizar la evaluación de un solo modelo y también validar el modelo actual en comparación con los modelos anteriores. Con este cambio, el componente Pusher ahora consume un resultado de bendición de Evaluator en lugar de ModelValidator.

El nuevo Evaluator es compatible con los modelos Keras y los modelos Estimator. El modelo guardado _eval_input_receiver_fn y eval que se requerían anteriormente ya no serán necesarios con Keras, ya que Evaluator ahora se basa en el mismo SavedModel que se usa para servir.

Consulte Evaluador para obtener más información .