O Google I/O é um embrulho! Fique por dentro das sessões do TensorFlow Ver sessões

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

Veja no TensorFlow.org Executar no Google Colab Ver fonte no GitHub Baixar caderno

Para fazer o aprendizado de máquina no TensorFlow, é provável que você precise definir, salvar e restaurar um modelo.

Um modelo é, abstratamente:

  • Uma função que calcula algo em tensores (uma passagem para frente )
  • Algumas variáveis ​​que podem ser atualizadas em resposta ao treinamento

Neste guia, você irá abaixo da superfície do Keras para ver como os modelos do TensorFlow são definidos. Isso analisa como o TensorFlow coleta variáveis ​​e modelos, bem como como eles são salvos e restaurados.

Configurar

import tensorflow as tf
from datetime import datetime

%load_ext tensorboard

Definindo modelos e camadas no TensorFlow

A maioria dos modelos são feitos de camadas. Camadas são funções com uma estrutura matemática conhecida que podem ser reutilizadas e possuem variáveis ​​treináveis. No TensorFlow, a maioria das implementações de alto nível de camadas e modelos, como Keras ou Sonnet , são construídas na mesma classe fundamental: tf.Module .

Aqui está um exemplo de um tf.Module muito simples que opera 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))
<tf.Tensor: shape=(), dtype=float32, numpy=30.0>

Módulos e, por extensão, camadas são terminologia de aprendizado profundo para "objetos": eles têm estado interno e métodos que usam esse estado.

Não há nada de especial em __call__ exceto agir como um Python que pode ser chamado; você pode invocar seus modelos com quaisquer funções 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.

Ao subclassificar tf.Module , quaisquer instâncias de tf.Variable ou tf.Module atribuídas às propriedades deste objeto são coletadas automaticamente. Isso permite que você salve e carregue variáveis, e também crie coleções de tf.Module s.

# All trainable variables
print("trainable variables:", simple_module.trainable_variables)
# Every variable
print("all variables:", simple_module.variables)
trainable variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>,)
all variables: (<tf.Variable 'train_me:0' shape=() dtype=float32, numpy=5.0>, <tf.Variable 'do_not_train_me:0' shape=() dtype=float32, numpy=5.0>)
2021-10-26 01:29:45.284549: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.

Este é um exemplo de um modelo de camada linear de duas camadas feito de 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 então o modelo completo, que faz duas instâncias de camada 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]])))
Model results: tf.Tensor([[7.706234  3.0919805]], shape=(1, 2), dtype=float32)

As instâncias tf.Module automaticamente, recursivamente, quaisquer instâncias tf.Variable ou tf.Module atribuídas a ela. Isso permite que você gerencie coleções de tf.Module s com uma única instância de modelo e salve e carregue modelos inteiros.

print("Submodules:", my_model.submodules)
Submodules: (<__main__.Dense object at 0x7f7ab2391290>, <__main__.Dense object at 0x7f7b6869ea10>)
for var in my_model.variables:
  print(var, "\n")
<tf.Variable 'b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)> 

<tf.Variable 'w:0' shape=(3, 3) dtype=float32, numpy=
array([[ 0.05711935,  0.22440144,  0.6370985 ],
       [ 0.3136791 , -1.7006774 ,  0.7256515 ],
       [ 0.16120772, -0.8412193 ,  0.5250952 ]], dtype=float32)> 

<tf.Variable 'b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)> 

<tf.Variable 'w:0' shape=(3, 2) dtype=float32, numpy=
array([[-0.5353216 ,  1.2815404 ],
       [ 0.62764466,  0.47087234],
       [ 2.19187   ,  0.45777202]], dtype=float32)>

Aguardando para criar variáveis

Você deve ter notado aqui que você tem que definir os tamanhos de entrada e saída para a camada. Isso é para que a variável w tenha uma forma conhecida e possa ser alocada.

Ao adiar a criação da variável para a primeira vez que o módulo é chamado com uma forma de entrada específica, você não precisa especificar o tamanho da entrada antecipadamente.

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]])))
Model results: tf.Tensor([[4.0598335 0.       ]], shape=(1, 2), dtype=float32)

Essa flexibilidade é o motivo pelo qual as camadas do TensorFlow geralmente precisam especificar apenas a forma de suas saídas, como em tf.keras.layers.Dense , em vez do tamanho da entrada e da saída.

Economizando pesos

Você pode salvar um tf.Module como um ponto de verificação e um SavedModel .

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

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

Os pontos de verificação consistem em dois tipos de arquivos: os próprios dados e um arquivo de índice para metadados. O arquivo de índice mantém o controle do que é realmente salvo e a numeração dos pontos de verificação, enquanto os dados do ponto de verificação contêm os valores das variáveis ​​e seus caminhos de pesquisa de atributos.

ls my_checkpoint*
my_checkpoint.data-00000-of-00001  my_checkpoint.index

Você pode olhar dentro de um ponto de verificação para ter certeza de que toda a coleção de variáveis ​​foi salva, classificada pelo objeto Python que as contém.

tf.train.list_variables(chkp_path)
[('_CHECKPOINTABLE_OBJECT_GRAPH', []),
 ('model/dense_1/b/.ATTRIBUTES/VARIABLE_VALUE', [3]),
 ('model/dense_1/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 3]),
 ('model/dense_2/b/.ATTRIBUTES/VARIABLE_VALUE', [2]),
 ('model/dense_2/w/.ATTRIBUTES/VARIABLE_VALUE', [3, 2])]

Durante o treinamento distribuído (com várias máquinas), eles podem ser fragmentados, razão pela qual são numerados (por exemplo, '00000-de-00001'). Neste caso, porém, há apenas um fragmento.

Quando você carrega os modelos de volta, você sobrescreve os valores em seu objeto 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]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[4.0598335, 0.       ]], dtype=float32)>

Salvando funções

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

O TensorFlow precisa saber fazer os cálculos descritos em Python, mas sem o código original . Para fazer isso, você pode criar um gráfico , descrito no guia Introdução a gráficos e funções .

Este gráfico contém operações, ou ops , que implementam a função.

Você pode definir um gráfico no modelo acima adicionando o decorador @tf.function para indicar que esse código deve ser executado como um gráfico.

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 como antes. Cada assinatura única passada para a função cria um gráfico separado. Consulte o guia Introdução a gráficos e funções para obter detalhes.

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
tf.Tensor([[0.62891716 0.        ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[0.62891716 0.        ]
  [0.62891716 0.        ]]], shape=(1, 2, 2), dtype=float32)

Você pode visualizar o gráfico rastreando-o 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)
tf.Tensor([[0.         0.01750386]], shape=(1, 2), dtype=float32)

Inicie o TensorBoard para visualizar o rastreamento resultante:

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

Uma captura de tela do gráfico no TensorBoard

Criando um SavedModel

A maneira recomendada de compartilhar modelos completamente 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 forma:

tf.saved_model.save(my_model, "the_saved_model")
INFO:tensorflow:Assets written to: the_saved_model/assets
# Inspect the SavedModel in the directory
ls -l the_saved_model
total 24
drwxr-sr-x 2 kbuilder kokoro  4096 Oct 26 01:29 assets
-rw-rw-r-- 1 kbuilder kokoro 14702 Oct 26 01:29 saved_model.pb
drwxr-sr-x 2 kbuilder kokoro  4096 Oct 26 01:29 variables
# The variables/ directory contains a checkpoint of the variables
ls -l the_saved_model/variables
total 8
-rw-rw-r-- 1 kbuilder kokoro 408 Oct 26 01:29 variables.data-00000-of-00001
-rw-rw-r-- 1 kbuilder kokoro 356 Oct 26 01:29 variables.index

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

Modelos e camadas podem ser carregados a partir dessa representação sem realmente criar uma instância da classe que a criou. Isso é desejado em situações em que você não tem (ou deseja) um interpretador Python, como servir em 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 usar.

Você pode carregar o modelo como novo objeto:

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

new_model , criado a partir do carregamento de um modelo salvo, é um objeto de usuário interno do TensorFlow sem nenhum conhecimento de classe. Não é do tipo SequentialModule .

isinstance(new_model, SequentialModule)
False

Este novo modelo funciona nas assinaturas de entrada já definidas. Você não pode adicionar mais assinaturas a um modelo restaurado assim.

print(my_model([[2.0, 2.0, 2.0]]))
print(my_model([[[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]]))
tf.Tensor([[0.62891716 0.        ]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[0.62891716 0.        ]
  [0.62891716 0.        ]]], shape=(1, 2, 2), dtype=float32)

Assim, usando SavedModel , você pode salvar pesos e gráficos do TensorFlow usando tf.Module e carregá-los novamente.

Modelos e camadas Keras

Observe que até este ponto, não há menção a Keras. Você pode construir sua própria API de alto nível em cima de tf.Module , e as pessoas têm.

Nesta seção, você examinará como o Keras usa tf.Module . Um guia do usuário completo para os modelos Keras pode ser encontrado no guia Keras .

Camadas Keras

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

Você pode converter um módulo em uma camada Keras apenas trocando o pai e alterando __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 Keras têm seu próprio __call__ que faz alguma contabilidade descrita na próxima seção e então chama call() . Você não deve notar nenhuma alteração na funcionalidade.

simple_layer([[2.0, 2.0, 2.0]])
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=array([[0.      , 0.179402, 0.      ]], dtype=float32)>

A etapa build

Conforme observado, em muitos casos é conveniente esperar para criar variáveis ​​até ter certeza da forma de entrada.

As camadas Keras vêm com uma etapa extra do ciclo de vida que permite mais flexibilidade na forma como você define suas camadas. Isso é definido na função build .

build é chamado exatamente uma vez e é chamado com a forma da entrada. Geralmente é usado para criar variáveis ​​(pesos).

Você pode reescrever a camada MyDense acima para ser flexível ao tamanho de suas 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 ponto, o modelo não foi construído, 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]])))
Model results: tf.Tensor(
[[-1.6998017  1.6444504 -1.3103955]
 [-2.5497022  2.4666753 -1.9655929]], shape=(2, 3), dtype=float32)
flexible_dense.variables
[<tf.Variable 'flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
 array([[ 1.277462  ,  0.5399406 , -0.301957  ],
        [-1.6277349 ,  0.7374014 , -1.7651852 ],
        [-0.49962795, -0.45511687,  1.4119445 ]], dtype=float32)>,
 <tf.Variable 'flexible_dense/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>]

Como o build é chamado apenas uma vez, as entradas serão rejeitadas se a forma de entrada 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)
Failed: In[0] mismatch In[1] shape: 4 vs. 3: [1,4] [3,3] 0 0 [Op:MatMul]

As camadas Keras têm muito mais recursos extras, incluindo:

  • Perdas opcionais
  • Suporte para métricas
  • Suporte integrado para um argumento de training opcional para diferenciar entre treinamento e uso de inferência
  • métodos get_config e from_config que permitem armazenar configurações com precisão para permitir a clonagem de modelo em Python

Leia sobre eles no guia completo de camadas e modelos personalizados.

Modelos Keras

Você pode definir seu modelo como camadas Keras aninhadas.

No entanto, Keras também fornece uma classe de modelo completa chamada tf.keras.Model . Ele herda de tf.keras.layers.Layer , portanto, um modelo Keras pode ser usado, aninhado e salvo da mesma forma que as camadas Keras. Os modelos Keras vêm com funcionalidade extra que os torna fáceis de treinar, avaliar, carregar, salvar e até mesmo treinar em várias máquinas.

Você pode definir o SequentialModule acima com código quase idêntico, novamente convertendo __call__ para 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]])))
Model results: tf.Tensor([[5.5604653 3.3511646]], shape=(1, 2), dtype=float32)

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

my_sequential_model.variables
[<tf.Variable 'my_sequential_model/flexible_dense_1/w:0' shape=(3, 3) dtype=float32, numpy=
 array([[ 0.05627853, -0.9386015 , -0.77410126],
        [ 0.63149   ,  1.0802224 , -0.37785745],
        [-0.24788402, -1.1076807 , -0.5956209 ]], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_1/b:0' shape=(3,) dtype=float32, numpy=array([0., 0., 0.], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_2/w:0' shape=(3, 2) dtype=float32, numpy=
 array([[-0.93912166,  0.77979285],
        [ 1.4049559 , -1.9380962 ],
        [-2.6039495 ,  0.30885765]], dtype=float32)>,
 <tf.Variable 'my_sequential_model/flexible_dense_2/b:0' shape=(2,) dtype=float32, numpy=array([0., 0.], dtype=float32)>]
my_sequential_model.submodules
(<__main__.FlexibleDense at 0x7f7b48525550>,
 <__main__.FlexibleDense at 0x7f7b48508d10>)

Substituir tf.keras.Model é uma abordagem muito Pythonica para construir modelos do TensorFlow. Se você estiver migrando modelos de outras estruturas, isso pode ser muito simples.

Se você estiver construindo modelos que são montagens simples de camadas e entradas existentes, poderá economizar tempo e espaço usando a API funcional , que vem com recursos adicionais sobre reconstrução e arquitetura de modelos.

Aqui está 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()
Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 3)]               0         
_________________________________________________________________
flexible_dense_3 (FlexibleDe (None, 3)                 12        
_________________________________________________________________
flexible_dense_4 (FlexibleDe (None, 2)                 8         
=================================================================
Total params: 20
Trainable params: 20
Non-trainable params: 0
_________________________________________________________________
my_functional_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[8.219393, 4.511119]], dtype=float32)>

A principal diferença aqui é que a forma de entrada é especificada antecipadamente como parte do processo de construção funcional. O argumento input_shape neste caso não precisa ser completamente especificado; você pode deixar algumas dimensões como None .

Salvando modelos Keras

Os modelos Keras podem ser verificados, e isso terá a mesma aparência que tf.Module .

Os modelos Keras também podem ser salvos com tf.saved_model.save() , pois são módulos. No entanto, os modelos Keras têm métodos de conveniência e outras funcionalidades:

my_sequential_model.save("exname_of_file")
INFO:tensorflow:Assets written to: exname_of_file/assets

Com a mesma facilidade, eles podem ser carregados de volta em:

reconstructed_model = tf.keras.models.load_model("exname_of_file")
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.

Os Keras SavedModels também salvam estados de métrica, perda e otimizador.

Este modelo reconstruído pode ser usado e produzirá o mesmo resultado quando chamado nos mesmos dados:

reconstructed_model(tf.constant([[2.0, 2.0, 2.0]]))
<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[5.5604653, 3.3511646]], dtype=float32)>

Há mais para saber sobre como salvar e serializar modelos Keras, incluindo o fornecimento de métodos de configuração para camadas personalizadas para suporte a recursos. Confira o guia para salvar e serialização .

Qual é o próximo

Se você quiser saber mais detalhes sobre Keras, você pode seguir os guias Keras existentes aqui .

Outro exemplo de uma API de alto nível criada no tf.module é o Sonnet da DeepMind, que é abordado em seu site .