intro_to_modules.ipynb

Kernel: Python 3

Copyright 2020 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.

Introdução aos módulos, camadas e modelos

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Para fazer aprendizado de máquina no TensorFlow, provavelmente você precisará definir, salvar e restaurar um modelo.

De forma abstrata, um modelo é:

• Uma função que faz determinadas computações com tensores (um passo para frente).

• Algumas variáveis que podem ser atualizadas em resposta ao treinamento.

Neste guia, veremos como os modelos do TensorFlow são definidos no Keras e como o TensorFlow coleta variáveis e modelos, além de como os modelos são salvos e restaurados.

Observação: se você quiser começar a usar o Keras imediatamente, confira a coleção de guias do Keras.

Configuração

import tensorflow as tf
from datetime import datetime

%load_ext tensorboard

Módulos do TensorFlow

A maioria dos modelos é feita de camadas, que são funções com uma estrutura matemática conhecida que podem ser reusadas e têm variáveis treináveis. No TensorFlow, a maioria das implementações de alto nível de camadas e modelos, como o Keras ou o Sonnet, são construídas usando-se a mesma classe fundamental: tf.Module.

Como construir módulos

Veja um exemplo de um tf.Module bem simples que faz operações em um tensor escalar:

class SimpleModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)
    self.a_variable = tf.Variable(5.0, name="train_me")
    self.non_trainable_variable = tf.Variable(5.0, trainable=False, name="do_not_train_me")
  def __call__(self, x):
    return self.a_variable * x + self.non_trainable_variable

simple_module = SimpleModule(name="simple")

simple_module(tf.constant(5.0))

Os módulos e, por extensão, as camadas são terminologias do aprendizado profundo para "objetos": têm um estado interno e métodos para acessar esse estado.

Não há nada de especial quanto a __call__, exceto pelo fato de funcionar como um callable do Python – você pode chamar seus modelos com qualquer função que desejar.

Você pode ativar e desativar a capacidade de treinamento de variáveis por qualquer motivo, incluindo o congelamento de camadas e variáveis durante o ajuste fino.

Observação: tf.Module é a classe básica tanto de tf.keras.layers.Layer quanto de tf.keras.Model, então tudo o que é abordado aqui aplica-se ao Keras. Por questões de compatibilidade histórica, as camadas do Keras não coletam variáveis dos modelos, então seus modelos devem usar apenas módulos ou apenas camadas do Keras. Entretanto, os métodos mostrados abaixo para inspecionar variáveis são os mesmos em ambos os casos.

Ao criar uma subclasse de tf.Module, qualquer instância de tf.Variable ou tf.Module atribuída às propriedades desse objeto é coletada automaticamente, o que permite salvar e carregar variáveis, bem como criar coleções de tf.Modules.

# All trainable variables
print("trainable variables:", simple_module.trainable_variables)
# Every variable
print("all variables:", simple_module.variables)

Este é um exemplo de um modelo linear de duas camadas feito com módulos.

Primeiro, uma camada densa (linear):

class Dense(tf.Module):
  def __init__(self, in_features, out_features, name=None):
    super().__init__(name=name)
    self.w = tf.Variable(
      tf.random.normal([in_features, out_features]), name='w')
    self.b = tf.Variable(tf.zeros([out_features]), name='b')
  def __call__(self, x):
    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)

E depois o modelo completo, que cria instâncias de duas camadas e as aplica:

class SequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = Dense(in_features=3, out_features=3)
    self.dense_2 = Dense(in_features=3, out_features=2)

  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a model!
my_model = SequentialModule(name="the_model")

# Call it, with random results
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))

As instâncias de tf.Module coletam de forma automática e recursiva qualquer instância de tf.Variable ou tf.Module atribuída a ele, o que permite gerenciar coleções de tf.Modules com uma única instância de modelo, além de salvar e carregar modelos inteiros.

print("Submodules:", my_model.submodules)

for var in my_model.variables:
  print(var, "\n")

Esperar para criar variáveis

Talvez você tenha notado que precisa definir os tamanhos de entrada e saída da camada. Isso ocorre para que a variável w tenha um formato conhecido e possa ser alocada.

Ao adiar a criação da variável para a primeira vez em que o módulo é chamado com um formato de entrada específico, você não precisa especificar o tamanho da entrada de antemão.

class FlexibleDenseModule(tf.Module):
  # Note: No need for `in_features`
  def __init__(self, out_features, name=None):
    super().__init__(name=name)
    self.is_built = False
    self.out_features = out_features

  def __call__(self, x):
    # Create variables on first call.
    if not self.is_built:
      self.w = tf.Variable(
        tf.random.normal([x.shape[-1], self.out_features]), name='w')
      self.b = tf.Variable(tf.zeros([self.out_features]), name='b')
      self.is_built = True

    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)

# Used in a module
class MySequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = FlexibleDenseModule(out_features=3)
    self.dense_2 = FlexibleDenseModule(out_features=2)

  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

my_model = MySequentialModule(name="the_model")
print("Model results:", my_model(tf.constant([[2.0, 2.0, 2.0]])))

É devido a essa flexibilidade que as camadas do TensorFlow geralmente só precisam especificar o formato das saídas, como em tf.keras.layers.Dense, em vez de especificar o tamanho tanto da entrada quanto da saída.

Como salvar pesos

É possível salvar um tf.Module tanto como um checkpoint quanto como um SavedModel.

Os checkpoints são apenas os pesos (ou seja, os valores do conjunto de variáveis dentro do módulo e de seus submódulos):

chkp_path = "my_checkpoint"
checkpoint = tf.train.Checkpoint(model=my_model)
checkpoint.write(chkp_path)

Os checkpoints são compostos por dois tipos de arquivo: os dados em si e um arquivo de índice para os metadados. O arquivo de índice mantém o controle do que é salvo e da numeração dos checkpoints, enquanto os dados do checkpoint contêm os valores das variáveis e seus caminhos de pesquisa de atributos.

!ls my_checkpoint*

Você pode avaliar o checkpoint para ter certeza de que a coleção inteira de variáveis tenha sido salva, ordenadas pelo objeto do Python que as contém.

tf.train.list_variables(chkp_path)

Durante o treinamento distribuído (com várias máquinas), elas podem ser fragmentadas, e é por isso que não numeradas (por exemplo, '00000-of-00001'). Porém, neste caso, só há um fragmento.

Quando você carrega os modelos de volta, sobrescreve os valores em seu objeto do Python.

new_model = MySequentialModule()
new_checkpoint = tf.train.Checkpoint(model=new_model)
new_checkpoint.restore("my_checkpoint")

# Should be the same result as above
new_model(tf.constant([[2.0, 2.0, 2.0]]))

Observação: os checkpoints estão no centro dos workflows de treinamento longos. tf.checkpoint.CheckpointManager é uma classe helper que facilita bastante o gerenciamento de criação de checkpoints. Confira mais informações no guia Treinando checkpoints.

Como salvar funções

O TensorFlow pode executar modelos sem os objetos originais do Python, conforme demonstrado pelo TensorFlow Serving e TensorFlow Lite, mesmo quando você baixa um modelo treinado no TensorFlow Hub.

O TensorFlow precisa saber como realizar as computações descritas no Python, mas sem o código original. Para fazer isso, você pode criar um grafo, que é descrito no guia Introdução aos grafos e funções.

Esse grafo contém operações, chamadas de ops, que implementam a função.

Para definir um grafo no modelo acima, basta adicionar o decorador @tf.function para indicar que esse código deve ser executado como grafo.

class MySequentialModule(tf.Module):
  def __init__(self, name=None):
    super().__init__(name=name)

    self.dense_1 = Dense(in_features=3, out_features=3)
    self.dense_2 = Dense(in_features=3, out_features=2)

  @tf.function
  def __call__(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a model with a graph!
my_model = MySequentialModule(name="the_model")

O módulo que você criou funciona exatamente da mesma forma que antes. Cada assinatura única passada para a função cria um grafo separado. Confira mais informações no Introdução aos grafos e funções .

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))

Para ver o grafo, basta traçá-lo em um resumo do TensorBoard.

# Set up logging.
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
logdir = "logs/func/%s" % stamp
writer = tf.summary.create_file_writer(logdir)

# Create a new model to get a fresh trace
# Otherwise the summary will not see the graph.
new_model = MySequentialModule()

# Bracket the function call with
# tf.summary.trace_on() and tf.summary.trace_export().
tf.summary.trace_on(graph=True)
tf.profiler.experimental.start(logdir)
# Call only one tf.function when tracing.
z = print(new_model(tf.constant([[2.0, 2.0, 2.0]])))
with writer.as_default():
  tf.summary.trace_export(
      name="my_func_trace",
      step=0,
      profiler_outdir=logdir)

Abra o TensorBoard para ver o gráfico resultante:

#docs_infra: no_execute
%tensorboard --logdir logs/func

Como criar um SavedModel

A maneira recomendada de compartilhar modelos totalmente treinados é usar SavedModel. SavedModel contém uma coleção de funções e uma coleção de pesos.

Você pode salvar o modelo que acabou de treinar da seguinte maneira:

tf.saved_model.save(my_model, "the_saved_model")

# Inspect the SavedModel in the directory
!ls -l the_saved_model

# The variables/ directory contains a checkpoint of the variables 
!ls -l the_saved_model/variables

O arquivo saved_model.pb é um buffer de protocolo que descreve o tf.Graph funcional.

Os modelos e as camadas podem ser carregados a partir dessa representação sem criar uma instância da classe que a criou, o que é desejado em situações quando você não tem (ou não quer ter) um interpretador do Python, como ao oferecer serviços em larga escala ou em um dispositivo de borda, ou em situações em que o código Python original não está disponível ou não é prático de se usar.

Você pode carregar o modelo como um novo objeto:

new_model = tf.saved_model.load("the_saved_model")

new_model, criado ao carregar um modelo salvo, é um objeto de usuário interno do TensorFlow, sem qualquer conhecimento de classe. Ele não é do tipo SequentialModule.

isinstance(new_model, SequentialModule)

Esse novo modelo funciona com as assinaturas de entrada já definidas. Não é possível adicionar mais assinaturas a um modelo restaurado dessa forma.

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))

Portanto, ao usar SavedModel, você consegue salvar os pesos e grafos do TensorFlow utilizando tf.Module, e depois pode carregá-los de volta.

Modelos e camadas do Keras

Observe que, até o momento, não houve menção ao Keras. É possível criar sua própria API de alto nível usando tf.Module, e isso já foi feito por outras pessoas.

Nesta seção, você verá como o Keras usa tf.Module. Confira instruções completas sobre o modelo do Keras no guia do Keras.

As camadas e modelos do Keras têm muitos outros recursos, como:

• Perdas opcionais

• Suporte a métricas

• Suporte integrado a um argumento de training (treinamento) opcional para diferenciar entre treinamento e inferência

• Salvar e restaurar objetos do Python em vez de funções que são uma caixa preta

• Métodos get_config e from_config, que permitem armazenar com precisão as configurações para permitir a clonagem de modelos no Python

Com esses recursos, é possível criar modelos muito mais complexos por meio da criação de subclasses, como uma GAN personalizada ou um modelo de AutoEncoder Variacional (VAE, na sigla em inglês). Saiba mais no guia completo sobre camadas e modelos personalizados.

Os modelos do Keras também têm outras funcionalidades que tornam fácil treinar, avaliar, carregar, salvar e até mesmo treinar modelos em várias máquinas.

Camadas do Keras

tf.keras.layers.Layer é a classe base de todas as camadas do Keras e herda de tf.Module.

Para converter um módulo em uma camada do Keras, basta trocar o pai e alterar __call__ para call:

class MyDense(tf.keras.layers.Layer):
  # Adding **kwargs to support base Keras layer arguments
  def __init__(self, in_features, out_features, **kwargs):
    super().__init__(**kwargs)

    # This will soon move to the build step; see below
    self.w = tf.Variable(
      tf.random.normal([in_features, out_features]), name='w')
    self.b = tf.Variable(tf.zeros([out_features]), name='b')
  def call(self, x):
    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)

simple_layer = MyDense(name="simple", in_features=3, out_features=3)

As camadas do Keras têm sua própria __call__, que faz alguns controles conforme descrito na próxima seção e depois chama call(). Você não deverá observar mudanças nessa funcionalidade.

simple_layer([[2.0, 2.0, 2.0]])

Passo build

Conforme vimos, em muitos casos, é conveniente esperar a criação de variáveis até saber com certeza o formato da entrada.

As camadas do Keras vêm com um passo extra de ciclo de vida que permite mais flexibilidade na definição das camadas. Isso é definido na função build.

A função build é chamada exatamente uma vez, e é chamada com o formato da entrada. Geralmente, é usada para criar variáveis (pesos).

Você pode reescrever a camada MyDense acima para flexibilizar o tamanho das entradas:

class FlexibleDense(tf.keras.layers.Layer):
  # Note the added `**kwargs`, as Keras supports many arguments
  def __init__(self, out_features, **kwargs):
    super().__init__(**kwargs)
    self.out_features = out_features

  def build(self, input_shape):  # Create the state of the layer (weights)
    self.w = tf.Variable(
      tf.random.normal([input_shape[-1], self.out_features]), name='w')
    self.b = tf.Variable(tf.zeros([self.out_features]), name='b')

  def call(self, inputs):  # Defines the computation from inputs to outputs
    return tf.matmul(inputs, self.w) + self.b

# Create the instance of the layer
flexible_dense = FlexibleDense(out_features=3)

Neste momento, o modelo ainda não foi criado, então não há variáveis:

flexible_dense.variables

Chamar a função aloca variáveis de tamanho apropriado.

# Call it, with predictably random results
print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0], [3.0, 3.0, 3.0]])))

flexible_dense.variables

Como a função build é chamada apenas uma vez, as entradas serão rejeitadas se o formato delas não for compatível com as variáveis da camada:

try:
  print("Model results:", flexible_dense(tf.constant([[2.0, 2.0, 2.0, 2.0]])))
except tf.errors.InvalidArgumentError as e:
  print("Failed:", e)

Modelos do Keras

É possível definir seu modelo como camadas aninhadas do Keras.

Porém, o Keras também conta com uma classe de modelo completa chamada tf.keras.Model, que herda de tf.keras.layers.Layer, então um modelo do Keras pode ser usado e aninhado da mesma forma que camadas do Keras. Os modelos do Keras contam com uma funcionalidade extra que tornam fácil treinar, avaliar, carregar, salvar e até mesmo treinar modelos em várias máquinas.

Você pode definir o SequentialModule acima com praticamente o mesmo código, novamente convertendo __call__ em call() e alterando o pai:

class MySequentialModel(tf.keras.Model):
  def __init__(self, name=None, **kwargs):
    super().__init__(**kwargs)

    self.dense_1 = FlexibleDense(out_features=3)
    self.dense_2 = FlexibleDense(out_features=2)
  def call(self, x):
    x = self.dense_1(x)
    return self.dense_2(x)

# You have made a Keras model!
my_sequential_model = MySequentialModel(name="the_model")

# Call it on a tensor, with random results
print("Model results:", my_sequential_model(tf.constant([[2.0, 2.0, 2.0]])))

Todos os mesmos recursos estão disponíveis, incluindo o rastreamento de variáveis e submódulos.

Observação: um tf.Module bruto aninhado dentro de uma camada ou modelo do Keras não tem suas variáveis coletadas para treinamento ou salvamento. Em vez disso, aninhe camadas do Keras dentro de camadas do Keras.

my_sequential_model.variables

my_sequential_model.submodules

Sobrescrever tf.keras.Model é uma estratégia típica do Python para criar modelos do TensorFlow. Se você estiver fazendo a migração de modelos de outros frameworks, isso poderá ser bem direto.

Se você estiver construindo modelos que são uma montagem simples de camadas e entradas existentes, pode economizar tempo e espaço utilizando a API funcional, que conta com recursos adicionais relacionados à arquitetura e à reconstrução do modelo.

Veja o mesmo modelo com a API funcional:

inputs = tf.keras.Input(shape=[3,])

x = FlexibleDense(3)(inputs)
x = FlexibleDense(2)(x)

my_functional_model = tf.keras.Model(inputs=inputs, outputs=x)

my_functional_model.summary()

my_functional_model(tf.constant([[2.0, 2.0, 2.0]]))

A principal diferença é que o formato da entrada é especificado de antemão como parte do processo de construção funcional. Neste caso, o argumento input_shape não precisa ser totalmente especificado. Você pode deixar algumas dimensões como None (nenhuma).

Observação: você não precisa especificar input_shape ou uma InputLayer em um modelo criado via subclasse. Esses argumentos e camadas serão ignorados.

Como salvar modelos do Keras

Os modelos do Keras têm seu próprio formato especializado de salvamento de arquivo ZIP, com extensão .keras. Ao chamar tf.keras.Model.save, acrescente a extensão .keras ao nome do arquivo. Veja um exemplo:

my_sequential_model.save("exname_of_file.keras")

Também é fácil carregar de volta:

reconstructed_model = tf.keras.models.load_model("exname_of_file.keras")

Os arquivos ZIP do Keras, arquivos .keras, também salvam métricas, perdas e estados de otimizadores.

Esse modelo reconstruído pode ser usado e gerará o mesmo resultado quando chamado com os mesmos dados:

reconstructed_model(tf.constant([[2.0, 2.0, 2.0]]))

Criação de checkpoints de modelos do Keras

Também é possível criar checkpoints de modelos do Keras, e o resultado será o mesmo de tf.Module.

O salvamento e a serialização de modelos do Keras não param por aí, incluindo métodos de configuração para camadas personalizadas com suporte a características. Confira o guia sobre salvamento e serialização.

Próximos passos

Se você quiser saber mais detalhes sobre o Keras, confira os guias existentes aqui.

Outro exemplo de uma API de alto nível construída com tf.module é o Sonnet da DeepMind, abordado neste site.