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

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

Para fazer o aprendizado de máquina no TensorFlow, provavelmente você precisará 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 de 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

Definição de 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 têm 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 criadas na mesma classe tf.Module : 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 chamável ; você pode invocar seus modelos com quaisquer funções que desejar.

Você pode definir a treinabilidade das variáveis ​​ativada e desativada por qualquer motivo, incluindo camadas de congelamento e variáveis ​​durante o ajuste fino.

Por subclasse tf.Module , nenhum tf.Variable ou tf.Module casos atribuídos a 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>)

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 cria 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([[6.061718  0.6471251]], shape=(1, 2), dtype=float32)

tf.Module casos será automaticamente coleta, de forma recursiva, nenhum tf.Variable ou tf.Module casos atribuído a ele. 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 0x7fc8f8303390>, <__main__.Dense object at 0x7fc914063f10>)
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.19300742,  1.3667928 , -0.29094702],
       [-0.56738806,  0.41195342, -0.8029899 ],
       [-2.249858  ,  0.21082571,  0.52548176]], 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([[-2.0251064 ,  1.2355493 ],
       [ 1.5233724 ,  0.16262923],
       [-1.4819455 , -0.39067182]], dtype=float32)>

Esperando para criar variáveis

Você deve ter notado aqui que é necessário definir os tamanhos de entrada e saída para a camada. Isso ocorre 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([[0.         0.08831326]], shape=(1, 2), dtype=float32)

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

Economizando peso

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

Os pontos de verificação 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 rastreia o que é realmente salvo e a numeração dos pontos de verificação, enquanto os dados dos pontos 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 (multi-máquina), eles podem ser fragmentados, e é por isso que são numerados (por exemplo, '00000-de-00001'). Nesse caso, porém, há apenas um fragmento.

Ao carregar 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([[0.        , 0.08831326]], dtype=float32)>

Salvar funções

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

O TensorFlow precisa saber como fazer os cálculos descritos em Python, mas sem o código original . Para fazer isso, você pode fazer um gráfico , que é descrito na Introdução aos gráficos e guia de 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 exclusiva passada para a função cria um gráfico separado. Consulte a Introdução aos gráficos e guia de 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. 0.]], shape=(1, 2), dtype=float32)
tf.Tensor(
[[[0. 0.]
  [0. 0.]]], shape=(1, 2, 2), dtype=float32)

Você pode visualizar o gráfico traçando-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([[4.6956697 0.       ]], shape=(1, 2), dtype=float32)

Inicie o TensorBoard para visualizar o trace resultante:

%tensorboard --logdir logs/func

Uma captura de tela do gráfico no TensorBoard

Criação de 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")
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 Jun  3 01:23 assets
-rw-rw-r-- 1 kbuilder kokoro 14343 Jun  3 01:23 saved_model.pb
drwxr-sr-x 2 kbuilder kokoro  4096 Jun  3 01:23 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 Jun  3 01:23 variables.data-00000-of-00001
-rw-rw-r-- 1 kbuilder kokoro 356 Jun  3 01:23 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 fazer uma instância da classe que os 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 ponta, 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 um 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 como este.

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

Assim, usando SavedModel , você pode salvar pesos e gráficos do tf.Module 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 do tf.Module , e as pessoas fizeram.

Nesta seção, você examinará como 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, em seguida, 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 alguns __call__ descritos na próxima seção e depois chama call() . Você não deve notar nenhuma mudança na funcionalidade.

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

A etapa de build

Conforme observado, é conveniente em muitos casos esperar para criar variáveis ​​até que você tenha certeza do formato de entrada.

As camadas Keras vêm com uma etapa extra do ciclo de vida que permite mais flexibilidade em como você define suas camadas. Isso é definido na função de 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(
[[ 2.704015   -0.16195035 -6.261377  ]
 [ 4.0560226  -0.24292576 -9.392065  ]], shape=(2, 3), dtype=float32)
flexible_dense.variables
[<tf.Variable 'flexible_dense/w:0' shape=(3, 3) dtype=float32, numpy=
 array([[-0.27628487,  0.13671786, -1.6931149 ],
        [ 1.4245744 , -1.238444  , -0.17901585],
        [ 0.20371802,  1.0207509 , -1.2585577 ]], 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 o uso de treinamento e inferência
  • métodos get_config e from_config que permitem armazenar com precisão as configurações para permitir a clonagem do modelo em Python

Leia sobre eles no guia completo para 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 com todos os recursos chamada tf.keras.Model . Ele herda de tf.keras.layers.Layer , portanto, um modelo Keras pode ser usado, aninhado e salvo da mesma maneira que as camadas Keras. Os modelos Keras vêm com funcionalidades extras que os tornam 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([[2.1485283 8.22457  ]], 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([[-1.35451   , -0.40638027, -0.40394917],
        [ 2.3224106 , -0.85175407, -1.6010288 ],
        [-0.5603658 , -0.3443388 ,  1.1965553 ]], 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([[ 2.0412993 , -0.30181855],
        [-0.54148334, -2.0181963 ],
        [ 0.77354205, -1.2384381 ]], 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 0x7fc9a0b9fe90>,
 <__main__.FlexibleDense at 0x7fc9a1109590>)

Substituir tf.keras.Model é uma abordagem muito pitônica 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, você pode economizar tempo e espaço usando a API funcional , que vem com recursos adicionais em torno da reconstrução e arquitetura do modelo.

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([[-1.8865316, -3.926752 ]], 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 terão a mesma tf.Module 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.

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([[2.1485283, 8.22457  ]], dtype=float32)>

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

Qual é o próximo

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

Outro exemplo de API de alto nível construída em tf.module é Sonnet da DeepMind, que é abordada em seu site .