Introducción a módulos, capas y modelos

Ver en TensorFlow.org Ejecutar en Google Colab Ver fuente en GitHub Descargar libreta

Para realizar el aprendizaje automático en TensorFlow, es probable que deba definir, guardar y restaurar un modelo.

Un modelo es, en abstracto:

  • Una función que calcula algo en tensores (un pase hacia adelante )
  • Algunas variables que se pueden actualizar en respuesta al entrenamiento

En esta guía, irá debajo de la superficie de Keras para ver cómo se definen los modelos de TensorFlow. Esto analiza cómo TensorFlow recopila variables y modelos, así como también cómo se guardan y restauran.

Configuración

import tensorflow as tf
from datetime import datetime

%load_ext tensorboard

Definición de modelos y capas en TensorFlow

La mayoría de los modelos están hechos de capas. Las capas son funciones con una estructura matemática conocida que se pueden reutilizar y tienen variables entrenables. En TensorFlow, la mayoría de las implementaciones de capas y modelos de alto nivel, como Keras o Sonnet , se crean en la misma clase fundamental: tf.Module .

Aquí hay un ejemplo de un tf.Module muy simple que opera en un 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>

Los módulos y, por extensión, las capas son terminología de aprendizaje profundo para "objetos": tienen un estado interno y métodos que usan ese estado.

No hay nada especial en __call__ excepto que actúa como un Python invocable ; puede invocar sus modelos con las funciones que desee.

Puede activar y desactivar la entrenabilidad de las variables por cualquier motivo, incluida la congelación de capas y variables durante el ajuste fino.

Al subclasificar tf.Module , cualquier instancia de tf.Variable o tf.Module asignada a las propiedades de este objeto se recopila automáticamente. Esto le permite guardar y cargar variables, y también crear colecciones 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 es un ejemplo de un modelo de capa lineal de dos capas hecho de módulos.

Primero una capa densa (lineal):

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)

Y luego el modelo completo, que crea dos instancias de capa y las 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)

Las instancias de tf.Module recopilarán automáticamente, de forma recursiva, cualquier instancia de tf.Variable o tf.Module que se le haya asignado. Esto le permite administrar colecciones de tf.Module s con una sola instancia de modelo y guardar y cargar modelos completos.

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)>

Esperando para crear variables

Es posible que haya notado aquí que tiene que definir los tamaños de entrada y salida de la capa. Esto es para que la variable w tenga una forma conocida y se pueda asignar.

Al aplazar la creación de variables a la primera vez que se llama al módulo con una forma de entrada específica, no necesita especificar el tamaño de entrada por adelantado.

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)

Esta flexibilidad es la razón por la cual las capas de TensorFlow a menudo solo necesitan especificar la forma de sus salidas, como en tf.keras.layers.Dense , en lugar del tamaño de entrada y salida.

Ahorro de pesos

Puede guardar un tf.Module como punto de control y como modelo guardado .

Los puntos de control son solo los pesos (es decir, los valores del conjunto de variables dentro del módulo y sus submódulos):

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

Los puntos de control constan de dos tipos de archivos: los datos en sí y un archivo de índice para los metadatos. El archivo de índice realiza un seguimiento de lo que realmente se guarda y la numeración de los puntos de control, mientras que los datos del punto de control contienen los valores de las variables y sus rutas de búsqueda de atributos.

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

Puede mirar dentro de un punto de control para asegurarse de que se guarde toda la colección de variables, ordenadas por el objeto de Python que las contiene.

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 el entrenamiento distribuido (en varias máquinas), se pueden fragmentar, por lo que se numeran (p. ej., '00000-de-00001'). En este caso, sin embargo, solo hay un fragmento.

Cuando vuelve a cargar modelos, sobrescribe los valores en su objeto de 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)>

Guardar funciones

TensorFlow puede ejecutar modelos sin los objetos originales de Python, como lo demuestran TensorFlow Serving y TensorFlow Lite , incluso cuando descarga un modelo entrenado de TensorFlow Hub .

TensorFlow necesita saber cómo realizar los cálculos descritos en Python, pero sin el código original . Para hacer esto, puede hacer un gráfico , que se describe en la guía Introducción a gráficos y funciones .

Este gráfico contiene operaciones, u ops , que implementan la función.

Puede definir un gráfico en el modelo anterior agregando el decorador @tf.function para indicar que este código debe ejecutarse como un 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")

El módulo que ha creado funciona exactamente igual que antes. Cada firma única pasada a la función crea un gráfico separado. Consulte la guía Introducción a gráficos y funciones para obtener más detalles.

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)

Puede visualizar el gráfico trazándolo dentro de un resumen de 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 TensorBoard para ver el seguimiento resultante:

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

Una captura de pantalla del gráfico en TensorBoard

Crear un SavedModel

La forma recomendada de compartir modelos completamente entrenados es usar SavedModel . SavedModel contiene una colección de funciones y una colección de pesos.

Puede guardar el modelo que acaba de entrenar de la siguiente manera:

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

El archivo saved_model.pb es un búfer de protocolo que describe el tf.Graph funcional.

Los modelos y las capas se pueden cargar desde esta representación sin crear realmente una instancia de la clase que la creó. Esto es deseable en situaciones en las que no tiene (o no desea) un intérprete de Python, como servir a escala o en un dispositivo perimetral, o en situaciones en las que el código original de Python no está disponible o no es práctico de usar.

Puede cargar el modelo como nuevo objeto:

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

new_model , creado a partir de la carga de un modelo guardado, es un objeto de usuario interno de TensorFlow sin ningún conocimiento de la clase. No es del tipo SequentialModule .

isinstance(new_model, SequentialModule)
False

Este nuevo modelo funciona en las firmas de entrada ya definidas. No puede agregar más firmas a un modelo restaurado de esta manera.

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)

Por lo tanto, con SavedModel , puede guardar los pesos y gráficos de TensorFlow con tf.Module y luego volver a cargarlos.

Modelos y capas de Keras

Tenga en cuenta que hasta este punto, no se menciona a Keras. Puede crear su propia API de alto nivel sobre tf.Module , y la gente lo ha hecho.

En esta sección, examinará cómo utiliza Keras tf.Module . Puede encontrar una guía de usuario completa de los modelos Keras en la guía Keras .

Capas de Keras

tf.keras.layers.Layer es la clase base de todas las capas de Keras y hereda de tf.Module .

Puede convertir un módulo en una capa de Keras simplemente intercambiando el padre y luego cambiando __call__ a 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)

Las capas de Keras tienen su propia __call__ que lleva la contabilidad descrita en la siguiente sección y luego llama a call() . No debería notar ningún cambio en la funcionalidad.

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

El paso build

Como se indicó, en muchos casos es conveniente esperar para crear variables hasta que esté seguro de la forma de entrada.

Las capas de Keras vienen con un paso adicional del ciclo de vida que le permite una mayor flexibilidad en la forma en que define sus capas. Esto se define en la función de build .

build se llama exactamente una vez, y se llama con la forma de la entrada. Por lo general, se usa para crear variables (pesos).

Puede reescribir la capa MyDense anterior para que sea flexible al tamaño de sus 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)

En este punto, el modelo no se ha construido, por lo que no hay variables:

flexible_dense.variables
[]

Llamar a la función asigna variables de tamaño adecuado:

# 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)>]

Dado que build solo se llama una vez, las entradas se rechazarán si la forma de entrada no es compatible con las variables de la capa:

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]

Las capas de Keras tienen muchas más características adicionales que incluyen:

  • Pérdidas opcionales
  • Soporte para métricas
  • Compatibilidad integrada con un argumento de training opcional para diferenciar entre uso de entrenamiento e inferencia
  • métodos get_config y from_config que le permiten almacenar configuraciones con precisión para permitir la clonación de modelos en Python

Lea sobre ellos en la guía completa de capas y modelos personalizados.

Modelos Keras

Puede definir su modelo como capas de Keras anidadas.

Sin embargo, Keras también proporciona una clase de modelo con todas las funciones llamada tf.keras.Model . Hereda de tf.keras.layers.Layer , por lo que un modelo de Keras se puede usar, anidar y guardar de la misma manera que las capas de Keras. Los modelos Keras vienen con una funcionalidad adicional que los hace fáciles de entrenar, evaluar, cargar, guardar e incluso entrenar en varias máquinas.

Puede definir SequentialModule desde arriba con un código casi idéntico, convirtiendo nuevamente __call__ en call() y cambiando el padre:

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)

Todas las mismas funciones están disponibles, incluidas las variables de seguimiento y los 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>)

tf.keras.Model es un enfoque muy Pythonic para construir modelos de TensorFlow. Si está migrando modelos de otros marcos, esto puede ser muy sencillo.

Si está construyendo modelos que son ensamblajes simples de capas y entradas existentes, puede ahorrar tiempo y espacio utilizando la API funcional , que viene con funciones adicionales sobre la arquitectura y la reconstrucción del modelo.

Aquí está el mismo modelo con la 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)>

La principal diferencia aquí es que la forma de entrada se especifica por adelantado como parte del proceso de construcción funcional. El argumento input_shape en este caso no tiene que especificarse completamente; puede dejar algunas dimensiones como None .

Guardar modelos de Keras

Los modelos de Keras pueden tener puntos de control, y se verán igual que tf.Module .

Los modelos de Keras también se pueden guardar con tf.saved_model.save() , ya que son módulos. Sin embargo, los modelos Keras tienen métodos convenientes y otras funciones:

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

Con la misma facilidad, se pueden volver a cargar en:

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.

Los modelos guardados SavedModels también guardan estados de métrica, pérdida y optimizador.

Este modelo reconstruido se puede usar y producirá el mismo resultado cuando se invoque con los mismos datos:

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)>

Hay más que saber sobre el guardado y la serialización de modelos de Keras, incluido el suministro de métodos de configuración para capas personalizadas para compatibilidad con funciones. Consulte la guía para guardar y serializar .

Que sigue

Si quieres conocer más detalles sobre Keras, puedes seguir las guías de Keras existentes aquí .

Otro ejemplo de una API de alto nivel construida en tf.module es Sonnet de DeepMind, que se cubre en su sitio .