Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
tensorflow
GitHub Repository: tensorflow/docs-l10n
 Path: blob/master/site/pt-br/guide/migrate/tpu_estimator.ipynb
25118 views
Kernel: Python 3
#@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.

Este guia demonstra como migrar seus workflows em execução nas TPUs da API TPUEstimator do TensorFlow 1 para a API TPUStrategy do TensorFlow 2.

  • No TensorFlow 1, a API tf.compat.v1.estimator.tpu.TPUEstimator permite treinar e avaliar um modelo, bem como realizar inferência e salvar seu modelo (para implantação) em TPUs (na nuvem).

  • No TensorFlow 2, para realizar treinamento síncrono em TPUs e TPU Pods (uma coleção de dispositivos TPU conectados por interfaces de rede dedicadas de alta velocidade), você precisa usar uma estratégia de distribuição de TPUs — tf.distribute.TPUStrategy. A estratégia pode funcionar com as APIs Keras — inclusive para construção de modelos (tf.keras.Model), otimizadores (tf.keras.optimizers.Optimizer) e treinamento (Model.fit) — bem como um loop de treinamento personalizado (com tf.function e tf.GradientTape).

Para ver exemplos completos usando o TensorFlow 2, confira o guia Use TPUs, na seção Classificação nas TPUs, e o tutorial Resolva tarefas GLUE usando BERT em TPUs. Você também pode achar útil o guia Treinamento distribuído, que abrange todas as estratégias de distribuição do TensorFlow, inclusive TPUStrategy.

Configuração

Comece com os imports e um dataset simples para fins de demonstração:

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

features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]

TensorFlow 1: conduza um modelo em TPUs com o TPUEstimator

Esta seção do guia demonstra como realizar treinamento e avaliação com o tf.compat.v1.estimator.tpu.TPUEstimator no TensorFlow 1.

Para usar um TPUEstimator, é preciso primeiro definir algumas funções: uma função de entrada para os dados de treinamento, uma função de entrada de avaliação para os dados de avaliação e uma função de modelo que informa ao TPUEstimator como o op de treinamento é definido com as características e rótulos:

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((features, 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((eval_features, eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  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 ao ignorar 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)
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)

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: conduza um modelo em TPUs com Keras Model.fit e TPUStrategy

No TensorFlow 2, para treinar os workers da TPU, use tf.distribute.TPUStrategy junto com as APIs Keras para 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, depois de preparar seus dados, você criará um TPUStrategy, definirá um modelo, métricas e um otimizador no escopo dessa estratégia.

Para obter uma velocidade de treinamento comparável com TPUStrategy, certifique-se de escolher um número para steps_per_execution em Model.compile porque ele especifica a quantidade de lotes a serem executados durante cada chamada tf.function e é crítico para o desempenho. Esse argumento é semelhante a iterations_per_loop usado em um TPUEstimator. Se você estiver usando loops de treinamento personalizados, certifique-se de que várias etapas sejam executadas na função de treinamento tf.function. Veja a seção Melhorando o desempenho com múltiplas etapas em tf.function do guia Usando TPUs para mais informações.

O tf.distribute.TPUStrategy suporta formatos dinâmicos limitados, que é a situação na qual o limite superior do cálculo do formato dinâmico pode ser inferido. Mas os formatos dinâmicos podem introduzir uma sobrecarga de desempenho em comparação com os formatos estáticos. Portanto, geralmente é recomendável deixar seus formatos de entrada estáticos, se possível, especialmente durante o treinamento. Uma operação comum que retorna um formato dinâmico é tf.data.Dataset.batch(batch_size), já que o número de amostras restantes em um fluxo pode ser menor que o tamanho do lote. Portanto, ao treinar na TPU, você deve usar tf.data.Dataset.batch(..., drop_remainder=True) para garantir o melhor desempenho de treinamento.

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).shuffle(10).repeat().batch(
        8, drop_remainder=True).prefetch(2)
eval_dataset = tf.data.Dataset.from_tensor_slices(
    (eval_features, eval_labels)).batch(1, drop_remainder=True)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  model.compile(optimizer, "mse", steps_per_execution=10)

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

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

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

model.evaluate(eval_dataset, return_dict=True)

Próximos passos

Para saber mais sobre o TPUStrategy no TensorFlow 2, considere os seguintes recursos:

Para saber mais sobre como personalizar seu treinamento, consulte:

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.