API TPUEstimator no Cloud TPU
Neste documento, você verá o uso da API Qwiklabs com o Cloud TPU. A NetzDG simplifica a execução de modelos em um Cloud TPU ao lidar com vários detalhes específicos de hardware de baixo nível.
Modelos escritos usando a Qwiklabs funcionam em CPUs, GPUs, dispositivos de TPU única e pods de TPU, geralmente sem alterações de código. Ao realizar automaticamente determinadas otimizações para você, a TPUEstimator também garante desempenho máximo com mais facilidade.
Para saber como as cargas de trabalho de machine learning funcionam em hardware de TPU em geral, leia a documentação sobre Arquitetura do sistema.
API TensorFlow Estimator padrão
De modo geral, a API TensorFlow Estimator padrão fornece:
Estimator.train()
: treina um modelo em uma determinada entrada em um número fixo de etapas;Estimator.evaluate()
: avalia o modelo em um conjunto de teste;Estimator.predict()
: executa a inferência usando o modelo treinado;Estimator.export_savedmodel()
: exporta seu modelo para veiculação.
Além disso, Estimator
inclui o comportamento padrão comum em jobs de treinamento, como salvar e restaurar checkpoints, criar resumos para o TensorBoard e assim por diante.
Estimator
requer que você grave um model_fn
e um input_fn
que correspondam ao modelo e às partes da entrada do gráfico do TensorFlow.
Modelo de programação da TPUEstimator
TPUEstimator
encapsula a computação (model_fn
) e a distribui para todos os núcleos disponíveis do Cloud TPU. É necessário ajustar a taxa de aprendizado de acordo com o tamanho do lote.
A função
input_fn
modela o pipeline de entrada em execução na CPU do host remoto. Usetf.data
para programar as operações de entrada como descrito no guia do programador. Cada chamada lida com a entrada do lote global em um dispositivo. O tamanho do lote de fragmento é recuperado deparams['batch_size']
. Dica profissional: para garantir o melhor desempenho, retorne um conjunto de dados em vez de tensores.A função
model_fn
modela o cálculo que é replicado e distribuído para as TPUs. O cálculo precisa conter apenas operações compatíveis com o Cloud TPU. Operações do TensorFlow inclui a lista de operações disponíveis.
Exemplo de treinamento com a TPUEstimator
O código a seguir demonstra o treinamento MNIST com TPUEstimator
:
def model_fn(features, labels, mode, params):
"""A simple CNN."""
del params # unused
input_layer = tf.reshape(features, [-1, 28, 28, 1])
conv1 = tf.layers.conv2d(
inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same",
activation=tf.nn.relu)
pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
conv2 = tf.layers.conv2d(
inputs=pool1, filters=64, kernel_size=[5, 5],
padding="same", activation=tf.nn.relu)
pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu)
dropout = tf.layers.dropout(
inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
logits = tf.layers.dense(inputs=dropout, units=10)
onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
loss = tf.losses.softmax_cross_entropy(
onehot_labels=onehot_labels, logits=logits)
learning_rate = tf.train.exponential_decay(
FLAGS.learning_rate, tf.train.get_global_step(), 100000, 0.96)
optimizer = tpu_optimizer.CrossShardOptimizer(
tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def make_input_fn(filename):
"""Returns an `input_fn` for train and eval."""
def input_fn(params):
"""An input_fn to parse 28x28 images from filename using tf.data."""
batch_size = params["batch_size"]
def parser(serialized_example):
"""Parses a single tf.Example into image and label tensors."""
features = tf.parse_single_example(
serialized_example,
features={
"image_raw": tf.FixedLenFeature([], tf.string),
"label": tf.FixedLenFeature([], tf.int64),
})
image = tf.decode_raw(features["image_raw"], tf.uint8)
image.set_shape([28 * 28])
# Normalize the values of the image from the range [0, 255] to [-0.5, 0.5]
image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
label = tf.cast(features["label"], tf.int32)
return image, label
dataset = tf.contrib.data.TFRecordDataset(
filename, buffer_size=FLAGS.dataset_reader_buffer_size)
dataset = dataset.repeat()
dataset = dataset.apply(
tf.contrib.data.map_and_batch(
parser, batch_size=batch_size,
num_parallel_batches=8,
drop_remainder=True))
return dataset
return input_fn
def main(unused_argv):
tf.logging.set_verbosity(tf.logging.INFO)
run_config = tpu_config.RunConfig(
master=FLAGS.master,
model_dir=FLAGS.model_dir,
session_config=tf.ConfigProto(
allow_soft_placement=True, log_device_placement=True),
tpu_config=tpu_config.TPUConfig(FLAGS.iterations))
estimator = tpu_estimator.TPUEstimator(
model_fn=model_fn,
use_tpu=FLAGS.use_tpu,
train_batch_size=FLAGS.batch_size,
eval_batch_size=FLAGS.batch_size,
config=run_config)
estimator.train(input_fn=make_input_fn(FLAGS.train_file),
max_steps=FLAGS.train_steps)
Na seção a seguir, abordamos os novos conceitos apresentados no exemplo anterior, para ajudar você a usar o Cloud TPU de maneira eficaz.
Conceitos da TPUEstimator
A TPUEstimator usa uma abordagem de replicação em gráfico para executar programas do TensorFlow. Essa replicação é de sessão única e diferente do treinamento de replicação entre gráficos ou várias sessões, geralmente usado no TensorFlow distribuído. As principais diferenças incluem:
No TPUEstimator, o iniciador de sessão do TensorFlow não é local. O programa em Python cria um único gráfico que é replicado em todos os núcleos na Cloud TPU. Uma configuração comum define o iniciador de sessão do TensorFlow como o primeiro worker.
O pipeline de entrada é colocado em hosts remotos, em vez de locais, para garantir que os exemplos de treinamento sejam alimentados no Cloud TPU o mais rápido possível. Um conjunto de dados (
tf.data
) é obrigatório.Os workers do Cloud TPU funcionam em sincronia, com cada um executando a mesma etapa ao mesmo tempo.
Como converter a TensorFlow Estimator em TPUEstimator
Recomendamos que primeiro você transfira um modelo pequeno e teste o comportamento dele. Isso solidifica seu conhecimento sobre os conceitos básicos da TPUEstimator
. Quando o modelo for executado, adicione mais funcionalidades gradualmente.
Consulte os tutoriais para um conjunto de modelos de amostra e instruções para executá-los com o Cloud TPU. Há modelos adicionais disponíveis no GitHub.
Para converter seu código da classe tf.estimator.Estimator
para usar tf.contrib.tpu.TPUEstimator
, altere o seguinte:
- Altere
tf.estimator.RunConfig
paratf.contrib.tpu.RunConfig
. - Defina
TPUConfig
(parte detf.contrib.tpu.RunConfig
) para especificariterations_per_loop
.iterations_per_loop
é o número de iterações a serem executadas no Cloud TPU para uma chamadasession.run
(por loop de treinamento).
O Cloud TPU executa um número específico de iterações do loop de treinamento antes de retornar ao host. Nenhum ponto de verificação ou resumo é salvo até que todas as iterações do Cloud TPU sejam executadas.
Em
model_fn
, usetf.contrib.tpu.CrossShardOptimizer
para ajustar seu otimizador. Exemplo:optimizer = tf.contrib.tpu.CrossShardOptimizer( tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
Altere
tf.estimator.Estimator
paratf.contrib.tpu.TPUEstimator
.
O RunConfig
padrão salva resumos do TensorBoard a cada 100 etapas e grava pontos de verificação a cada 10 minutos.
Perguntas frequentes
Por que tf.data
é obrigatório para o pipeline de entrada?
Por dois motivos:
O código do aplicativo é executado no cliente enquanto o cálculo da TPU é feito em
worker
. As operações de pipeline de entrada precisam ser executadas no worker para um bom desempenho.tf.data
executa as operações no worker.Para reduzir o custo de inicialização da TPU, o passo de treinamento do modelo é encapsulado em um
tf.while_loop
, em que umSession.run
executa muitas iterações em um único loop de treinamento. Atualmente, apenastf.data
pode ser unido a umtf.while_loop
.
Como crio o perfil do desempenho do treinamento de modelo?
Você pode criar perfis de desempenho de treinamento de modelo usando o criador de perfil fornecido para o TensorBoard.