TensorFlow 2.x no TFX

O TensorFlow 2.0 foi lançado em 2019 , com forte integração do Keras , execução antecipada por padrão e execução da função Pythonic , entre outros novos recursos e melhorias .

Este guia fornece uma visão geral técnica abrangente do TF 2.x no TFX.

Qual versão usar?

O TFX é compatível com o TensorFlow 2.x, e as APIs de alto nível que existiam no TensorFlow 1.x (principalmente Estimators) continuam funcionando.

Inicie novos projetos no TensorFlow 2.x

Como o TensorFlow 2.x mantém os recursos de alto nível do TensorFlow 1.x, não há vantagem em usar a versão mais antiga em novos projetos, mesmo que você não planeje usar os novos recursos.

Portanto, se você estiver iniciando um novo projeto TFX, recomendamos o uso do TensorFlow 2.x. Você pode querer atualizar seu código mais tarde, pois o suporte completo para Keras e outros novos recursos se torna disponível, e o escopo das alterações será muito mais limitado se você começar com o TensorFlow 2.x, em vez de tentar atualizar do TensorFlow 1.x em o futuro.

Convertendo projetos existentes para o TensorFlow 2.x

O código escrito para o TensorFlow 1.x é amplamente compatível com o TensorFlow 2.xe continuará a funcionar no TFX.

No entanto, se você quiser aproveitar os aprimoramentos e novos recursos à medida que forem disponibilizados no TF 2.x, siga as instruções para migrar para o TF 2.x .

Estimador

A Estimator API foi mantida no TensorFlow 2.x, mas não é o foco de novos recursos e desenvolvimento. O código escrito no TensorFlow 1.x ou 2.x usando Estimators continuará funcionando conforme o esperado no TFX.

Aqui está um exemplo de TFX de ponta a ponta usando Estimator puro: exemplo de táxi (Estimator)

Keras com model_to_estimator

Os modelos Keras podem ser encapsulados com a função tf.keras.estimator.model_to_estimator , que permite que funcionem como se fossem estimadores. Para usar isso:

1. Construir um modelo Keras.
2. Passe o modelo compilado para model_to_estimator .
3. Use o resultado de model_to_estimator no Trainer, da mesma forma que você normalmente usaria um 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,
      ...
  }

Além do arquivo do módulo do usuário do Trainer, o restante do pipeline permanece inalterado.

Keras nativo (ou seja, Keras sem model_to_estimator )

Observação: o suporte completo para todos os recursos do Keras está em andamento; na maioria dos casos, o Keras no TFX funcionará conforme o esperado. Ele ainda não funciona com Sparse Features para FeatureColumns.

Exemplos e Colab

Aqui estão vários exemplos com Keras nativos:

• Pinguim ( arquivo de módulo ): exemplo de ponta a ponta 'Hello world'.
• MNIST ( arquivo de módulo ): exemplo ponta a ponta de imagem e TFLite.
• Táxi ( arquivo do módulo ): exemplo de ponta a ponta com uso avançado do Transform.

Também temos um Keras Colab por componente.

Componentes TFX

As seções a seguir explicam como os componentes TFX relacionados oferecem suporte ao Keras nativo.

Transformar

Transform atualmente tem suporte experimental para modelos Keras.

O próprio componente Transform pode ser usado para Keras nativo sem alterações. A definição preprocessing_fn permanece a mesma, usando TensorFlow e tf.Transform ops.

A função de serviço e a função de avaliação são alteradas para Keras nativo. Os detalhes serão discutidos nas seções a seguir do Instrutor e do Avaliador.

Observação: as transformações dentro do preprocessing_fn não podem ser aplicadas ao recurso de rótulo para treinamento ou avaliação.

Treinador

Para configurar o Keras nativo, o GenericExecutor precisa ser configurado para o componente Trainer para substituir o executor padrão baseado no Estimator. Para detalhes, por favor verifique aqui .

Arquivo de módulo Keras com Transform

O arquivo do módulo de treinamento deve conter um run_fn que será chamado pelo GenericExecutor , um run_fn típico do Keras ficaria assim:

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)

No run_fn acima, uma assinatura de serviço é necessária ao exportar o modelo treinado para que o modelo possa obter exemplos brutos para previsão. Uma função de serviço típica ficaria assim:

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

Na função de serviço acima, as transformações tf.Transform precisam ser aplicadas aos dados brutos para inferência, usando a camada tft.TransformFeaturesLayer . O _serving_input_receiver_fn anterior que era necessário para estimadores não será mais necessário com Keras.

Arquivo de módulo Keras sem Transform

Isso é semelhante ao arquivo de módulo mostrado acima, mas sem as transformações:

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.distribute.Strategy

No momento, o TFX suporta apenas estratégias de trabalho único (por exemplo, MirroredStrategy , OneDeviceStrategy ).

Para usar uma estratégia de distribuição, crie um tf.distribute.Strategy apropriado e mova a criação e compilação do modelo Keras dentro de um escopo de estratégia.

Por exemplo, substitua model = _build_keras_model() por:

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

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

Para verificar o dispositivo (CPU/GPU) usado por MirroredStrategy , habilite o log tensorflow de nível de informação:

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

e você deve ser capaz de ver Using MirroredStrategy with devices (...) no log.

Observação: a variável de ambiente TF_FORCE_GPU_ALLOW_GROWTH=true pode ser necessária para um problema de GPU sem memória. Para obter detalhes, consulte o guia de GPU do tensorflow .

avaliador

No TFMA v0.2x, ModelValidator e Evaluator foram combinados em um único novo componente Evaluator . O novo componente Avaliador pode executar avaliação de modelo único e também validar o modelo atual em comparação com modelos anteriores. Com essa alteração, o componente Pusher agora consome um resultado de bênção do Evaluator em vez do ModelValidator.

O novo Evaluator suporta modelos Keras, bem como modelos Estimator. O _eval_input_receiver_fn e o modelo salvo eval que eram necessários anteriormente não serão mais necessários com o Keras, pois o Evaluator agora é baseado no mesmo SavedModel usado para servir.

Consulte Avaliador para obter mais informações .