tpu_embedding.ipynb

Kernel: Python 3

Copyright 2021 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Como migrar de TPU embedding_columns para a camada TPUEmbedding

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Este guia mostra como migrar o treinamento de embedding em TPUs da API embedding_column do TensorFlow 1 com TPUEstimator, para a API da camada TPUEmbedding do TensorFlow 2 com TPUStrategy.

Embeddings são (grandes) matrizes. Eles são tabelas de associação que mapeiam de um espaço de características esparso para vetores densos. Embeddings fornecem representações eficientes e densas, capturando semelhanças complexas e relacionamentos entre características.

O TensorFlow inclui suporte especializado para treinar embeddings em TPUs. Esse suporte a embeddings, específico para TPUs, permite que você treine embeddings que são maiores do que a memória de um único dispositivo TPU e use entradas esparsas e irregulares em TPUs.

• No TensorFlow 1, tf.compat.v1.estimator.tpu.TPUEstimator é uma API de alto nível que encapsula treinamento, avaliação, previsão e exportação para servir com TPUs. Tem suporte especial para tf.compat.v1.tpu.experimental.embedding_column.

• Para implementar no TensorFlow 2, use a camada tfrs.layers.embedding.TPUEmbedding do TensorFlow Recommenders. Para treinamento e avaliação, use uma estratégia de distribuição de TPUs — tf.distribute.TPUStrategy — que é compatível com as APIs Keras para, por exemplo, construção de modelos (tf.keras.Model), otimizadores (tf.keras.optimizers.Optimizer), e treinamento com Model.fit ou um loop de treinamento personalizado com tf.function e tf.GradientTape.

Para informações adicionais, consulte a documentação da API da camada tfrs.layers.embedding.TPUEmbedding, bem como a documentação tf.tpu.experimental.embedding.TableConfig e tf.tpu.experimental.embedding.FeatureConfig para obter informações adicionais. Para obter uma visão geral de tf.distribute.TPUStrategy, confira o guia Treinamento distribuído e o guia Usando TPUs. Se você estiver migrando de TPUEstimator para TPUStrategy, confira o O guia de migração de TPUs.

Configuração

Comece instalando o TensorFlow Recommenders e importando alguns pacotes necessários:

!pip install tensorflow-recommenders

import tensorflow as tf
import tensorflow.compat.v1 as tf1

# TPUEmbedding layer is not part of TensorFlow.
import tensorflow_recommenders as tfrs

E prepare um dataset simples para fins de demonstração:

features = [[1., 1.5]]
embedding_features_indices = [[0, 0], [0, 1]]
embedding_features_values = [0, 5]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_embedding_features_indices = [[0, 0], [0, 1]]
eval_embedding_features_values = [4, 3]
eval_labels = [[0.8]]

TensorFlow 1: treine incorporações em TPUs com o TPUEstimator

No TensorFlow 1, você configura os embeddings de TPU usando a API tf.compat.v1.tpu.experimental.embedding_column e treina/avalia o modelo em TPUs com tf.compat.v1.estimator.tpu.TPUEstimator.

As entradas são números inteiros que variam de zero ao tamanho do vocabulário para a tabela de embeddings da TPU. Comece codificando as entradas para um ID categórico com tf.feature_column.categorical_column_with_identity. Use "sparse_feature" para o parâmetro key, já que os recursos de entrada são de valor inteiro, enquanto num_buckets é o tamanho do vocabulário para a tabela de embeddings (10).

embedding_id_column = (
      tf1.feature_column.categorical_column_with_identity(
          key="sparse_feature", num_buckets=10))

Em seguida, converta as entradas categóricas esparsas numa representação densa com tpu.experimental.embedding_column, onde dimension é a largura da tabela de embedding. Ela armazenará um vetor de embedding para cada um dos num_buckets.

embedding_column = tf1.tpu.experimental.embedding_column(
    embedding_id_column, dimension=5)

Agora, defina a configuração de incorporação específica para a TPU através de tf.estimator.tpu.experimental.EmbeddingConfigSpec. Você vai passá-la depois para tf.estimator.tpu.TPUEstimator como um parâmetro embedding_config_spec.

embedding_config_spec = tf1.estimator.tpu.experimental.EmbeddingConfigSpec(
    feature_columns=(embedding_column,),
    optimization_parameters=(
        tf1.tpu.experimental.AdagradParameters(0.05)))

Em seguida, para usar um TPUEstimator, defina:

• Uma função de entrada para os dados de treinamento

• Uma função de entrada de avaliação para os dados de avaliação

• Uma função de modelo para instruir o TPUEstimator como a operação de treinamento será definida com as características e rótulos

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf1.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf1.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  embedding_features = tf1.keras.layers.DenseFeatures(embedding_column)(features)
  concatenated_features = tf1.keras.layers.Concatenate(axis=1)(
      [embedding_features, features["dense_feature"]])
  logits = tf1.layers.Dense(1)(concatenated_features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  optimizer = tf1.tpu.CrossShardOptimizer(optimizer)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

Com essas funções definidas, crie um tf.distribute.cluster_resolver.TPUClusterResolver que forneça as informações do cluster e um objeto tf.compat.v1.estimator.tpu.RunConfig.

Junto com a função de modelo que você definiu, agora você pode criar um TPUEstimator. Aqui, você simplificará o fluxo ignorando o salvamento de checkpoints. Em seguida, você especificará o tamanho do lote para treinamento e avaliação para o TPUEstimator.

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))

tpu_config = tf1.estimator.tpu.TPUConfig(
    iterations_per_loop=10,
    per_host_input_for_training=tf1.estimator.tpu.InputPipelineConfig
          .PER_HOST_V2)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn, config=config, train_batch_size=8, eval_batch_size=8,
    embedding_config_spec=embedding_config_spec)

Chame TPUEstimator.train para começar a treinar o modelo:

estimator.train(_input_fn, steps=1)

Em seguida, chame TPUEstimator.evaluate para avaliar o modelo usando os dados de avaliação:

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: treine incorporações em TPUs com o TPUStrategy

No TensorFlow 2, para treinar os workers da TPU, use tf.distribute.TPUStrategy junto com as APIs Keras para a definição do modelo e do treinamento/avaliação. (Consulte o guia Usando TPUs para mais exemplos de treinamento com Keras Model.fit e um loop de treinamento personalizado (com tf.function e tf.GradientTape).)

Já que você precisa realizar alguma inicialização para se conectar ao cluster remoto e inicializar os workers da TPU, comece criando um TPUClusterResolver para fornecer as informações do cluster e conectar-se ao cluster. (Saiba mais na seção Inicialização da TPU do guia Usando TPUs.)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

Em seguida, prepare seus dados. Isto é semelhante a como você criou um dataset no exemplo do TensorFlow 1, exceto que a função do dataset agora é transmitida a um objeto tf.distribute.InputContext em vez de um dict params. Você pode usar esse objeto para determinar o tamanho do lote local (e para qual host esse pipeline se destina, para que você possa particionar seus dados adequadamente).

• Ao usar a API tfrs.layers.embedding.TPUEmbedding, é importante incluir a opção drop_remainder=True ao agrupar o dataset com Dataset.batch, pois TPUEmbedding requer um tamanho de lote fixo.

• Além disso, o mesmo tamanho de lote deve ser usado para avaliação e treinamento se estiverem ocorrendo no mesmo conjunto de dispositivos.

• Por fim, você deve usar tf.keras.utils.experimental.DatasetCreator junto com a opção de entrada especial — experimental_fetch_to_device=False — em tf.distribute.InputOptions (que contém configurações específicas da estratégia). Isto é demonstrado a seguir:

global_batch_size = 8

def _input_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": features,
       "sparse_feature": tf.SparseTensor(
           embedding_features_indices,
           embedding_features_values, [1, 2])},
           labels))
  dataset = dataset.shuffle(10).repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

def _eval_dataset(context: tf.distribute.InputContext):
  dataset = tf.data.Dataset.from_tensor_slices((
      {"dense_feature": eval_features,
       "sparse_feature": tf.SparseTensor(
           eval_embedding_features_indices,
           eval_embedding_features_values, [1, 2])},
           eval_labels))
  dataset = dataset.repeat()
  dataset = dataset.batch(
      context.get_per_replica_batch_size(global_batch_size),
      drop_remainder=True)
  return dataset.prefetch(2)

input_options = tf.distribute.InputOptions(
    experimental_fetch_to_device=False)

input_dataset = tf.keras.utils.experimental.DatasetCreator(
    _input_dataset, input_options=input_options)

eval_dataset = tf.keras.utils.experimental.DatasetCreator(
    _eval_dataset, input_options=input_options)

Em seguida, depois de preparar os dados, você criará um TPUStrategy e definirá um modelo, métricas e um otimizador no escopo dessa estratégia (Strategy.scope).

Você deve escolher um número para steps_per_execution em Model.compile, pois ele especifica a quantidade de lotes a serem executados durante cada chamada tf.function e é crítico para o desempenho. Esse argumento é semelhante ao iterations_per_loop usado em TPUEstimator.

As características e a configuração da tabela que foram especificadas no TensorFlow 1 através de tf.tpu.experimental.embedding_column (e tf.tpu.experimental.shared_embedding_column) podem ser especificados diretamente no TensorFlow 2 através de um par de objetos de configuração:

• tf.tpu.experimental.embedding.FeatureConfig

• tf.tpu.experimental.embedding.TableConfig

(Consulte a documentação da API associada para mais detalhes.)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  if hasattr(tf.keras.optimizers, "legacy"):
    optimizer = tf.keras.optimizers.legacy.Adagrad(learning_rate=0.05)
  else:
    optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  dense_input = tf.keras.Input(shape=(2,), dtype=tf.float32, batch_size=global_batch_size)
  sparse_input = tf.keras.Input(shape=(), dtype=tf.int32, batch_size=global_batch_size)
  embedded_input = tfrs.layers.embedding.TPUEmbedding(
      feature_config=tf.tpu.experimental.embedding.FeatureConfig(
          table=tf.tpu.experimental.embedding.TableConfig(
              vocabulary_size=10,
              dim=5,
              initializer=tf.initializers.TruncatedNormal(mean=0.0, stddev=1)),
          name="sparse_input"),
      optimizer=optimizer)(sparse_input)
  input = tf.keras.layers.Concatenate(axis=1)([dense_input, embedded_input])
  result = tf.keras.layers.Dense(1)(input)
  model = tf.keras.Model(inputs={"dense_feature": dense_input, "sparse_feature": sparse_input}, outputs=result)
  model.compile(optimizer, "mse", steps_per_execution=10)

Com isso, você está pronto para treinar o modelo com o dataset de treinamento:

model.fit(input_dataset, epochs=5, steps_per_epoch=10)

Por fim, avalie o modelo usando o dataset de avaliação:

model.evaluate(eval_dataset, steps=1, return_dict=True)

Próximos passos

Saiba mais sobre como configurar embeddings específicos para TPUs na documentação da API:

• tfrs.layers.embedding.TPUEmbedding: particularmente sobre configuração de características e tabelas, configuração do otimizador, criação de um modelo (usando a API funcional Keras ou através de uma subclasse de tf.keras.Model), treinamento/avaliação, e servndo modelos com tf.saved_model

• tf.tpu.experimental.embedding.TableConfig

• tf.tpu.experimental.embedding.FeatureConfig

Para mais informações sobre TPUStrategy no TensorFlow 2, considere os seguintes recursos:

• Guia: Usando TPUs (cobre treinamento com Keras Model.fit / um loop de treinamento personalizado com tf.distribute.TPUStrategy, bem como dicas sobre como melhorar o desempenho com tf.function)

• Guia: Treinamento distribuído com o TensorFlow

• Como migrar de TPUEstimator para TPUStrategy

Para saber mais sobre como personalizar seu treinamento, consulte:

• Guia: Personalize o que acontece em Model.fit

• Guia: Escrevendo um loop de treinamento do zero

As TPUs, os ASICs especializados do Google para aprendizado de máquina, estão disponíveis através do Google Colab, TPU Research Cloud e Cloud TPU.