sonnet
v2.0.2
Documentación | Ejemplos
Sonnet es una biblioteca construida sobre TensorFlow 2 diseñada para proporcionar abstracciones simples y componibles para la investigación del aprendizaje automático.
Sonnet ha sido diseñado y construido por investigadores de DeepMind. Se puede utilizar para construir redes neuronales para muchos propósitos diferentes (aprendizaje no supervisado, aprendizaje por refuerzo,...). Creemos que es una abstracción exitosa para nuestra organización, ¡usted también podría hacerlo!
Más específicamente, Sonnet proporciona un modelo de programación simple pero potente centrado en un único concepto: snt.Module . Los módulos pueden contener referencias a parámetros, otros módulos y métodos que aplican alguna función en la entrada del usuario. Sonnet se entrega con muchos módulos predefinidos (por ejemplo, snt.Linear , snt.Conv2D , snt.BatchNorm ) y algunas redes predefinidas de módulos (por ejemplo, snt.nets.MLP ), pero también se anima a los usuarios a crear sus propios módulos.
A diferencia de muchos marcos, Sonnet no tiene opiniones sobre cómo utilizará sus módulos. Los módulos están diseñados para ser autónomos y completamente desacoplados entre sí. Sonnet no se entrega con un marco de capacitación y se anima a los usuarios a crear el suyo propio o adoptar los creados por otros.
Sonnet también está diseñado para ser fácil de entender, nuestro código es (¡con suerte!) claro y enfocado. Cuando hemos elegido valores predeterminados (por ejemplo, valores predeterminados para los valores de los parámetros iniciales), intentamos señalar por qué.
La forma más sencilla de probar Sonnet es utilizar Google Colab, que ofrece un cuaderno Python gratuito conectado a una GPU o TPU.
snt.distribute Para comenzar, instale TensorFlow 2.0 y Sonnet 2:
$ pip install tensorflow tensorflow-probability
$ pip install dm-sonnetPuede ejecutar lo siguiente para verificar que todo esté instalado correctamente:
import tensorflow as tf
import sonnet as snt
print ( "TensorFlow version {}" . format ( tf . __version__ ))
print ( "Sonnet version {}" . format ( snt . __version__ )) Sonnet se envía con una serie de módulos integrados que puedes utilizar de forma trivial. Por ejemplo, para definir un MLP, podemos usar el módulo snt.Sequential para llamar a una secuencia de módulos, pasando la salida de un módulo determinado como entrada para el siguiente módulo. Podemos usar snt.Linear y tf.nn.relu para definir nuestro cálculo:
mlp = snt . Sequential ([
snt . Linear ( 1024 ),
tf . nn . relu ,
snt . Linear ( 10 ),
]) Para utilizar nuestro módulo necesitamos "llamarlo". El módulo Sequential (y la mayoría de los módulos) definen un método __call__ que significa que puedes llamarlos por su nombre:
logits = mlp ( tf . random . normal ([ batch_size , input_size ]))También es muy común solicitar todos los parámetros de tu módulo. La mayoría de los módulos en Sonnet crean sus parámetros la primera vez que se llaman con alguna entrada (ya que en la mayoría de los casos la forma de los parámetros es función de la entrada). Los módulos Sonnet proporcionan dos propiedades para acceder a los parámetros.
La propiedad variables devuelve todos tf.Variable s a los que hace referencia el módulo dado:
all_variables = mlp . variables Vale la pena señalar que tf.Variable s no solo se usan para los parámetros de su modelo. Por ejemplo, se utilizan para mantener el estado en las métricas utilizadas en snt.BatchNorm . En la mayoría de los casos, los usuarios recuperan las variables del módulo para pasarlas a un optimizador para su actualización. En este caso, las variables que no se pueden entrenar normalmente no deberían estar en esa lista, ya que se actualizan mediante un mecanismo diferente. TensorFlow tiene un mecanismo integrado para marcar variables como "entrenables" (parámetros de su modelo) frente a no entrenables (otras variables). Sonnet proporciona un mecanismo para recopilar todas las variables entrenables de su módulo, que probablemente sea lo que desea pasar a un optimizador:
model_parameters = mlp . trainable_variables Sonnet recomienda encarecidamente a los usuarios que subclasen snt.Module para definir sus propios módulos. Comencemos creando una capa Linear simple llamada MyLinear :
class MyLinear ( snt . Module ):
def __init__ ( self , output_size , name = None ):
super ( MyLinear , self ). __init__ ( name = name )
self . output_size = output_size
@ snt . once
def _initialize ( self , x ):
initial_w = tf . random . normal ([ x . shape [ 1 ], self . output_size ])
self . w = tf . Variable ( initial_w , name = "w" )
self . b = tf . Variable ( tf . zeros ([ self . output_size ]), name = "b" )
def __call__ ( self , x ):
self . _initialize ( x )
return tf . matmul ( x , self . w ) + self . bUsar este módulo es trivial:
mod = MyLinear ( 32 )
mod ( tf . ones ([ batch_size , input_size ])) Al subclasificar snt.Module obtienes muchas propiedades interesantes de forma gratuita. Por ejemplo, una implementación predeterminada de __repr__ que muestra los argumentos del constructor (muy útil para la depuración y la introspección):
>> > print ( repr ( mod ))
MyLinear ( output_size = 10 ) También obtienes las variables y las propiedades trainable_variables :
>> > mod . variables
( < tf . Variable 'my_linear/b:0' shape = ( 10 ,) ...) > ,
< tf . Variable 'my_linear/w:0' shape = ( 1 , 10 ) ...) > ) Es posible que observe el prefijo my_linear en las variables anteriores. Esto se debe a que los módulos Sonnet también ingresan al alcance del nombre del módulo cada vez que se llaman métodos. Al ingresar el alcance del nombre del módulo, proporcionamos un gráfico mucho más útil para que lo consuman herramientas como TensorBoard (por ejemplo, todas las operaciones que ocurren dentro de my_linear estarán en un grupo llamado my_linear).
Además, su módulo ahora admitirá los puntos de control de TensorFlow y el modelo guardado, que son funciones avanzadas que se tratan más adelante.
Sonnet admite múltiples formatos de serialización. El formato más simple que admitimos es pickle de Python, y todos los módulos integrados se prueban para garantizar que se puedan guardar/cargar mediante pickle en el mismo proceso de Python. En general, desaconsejamos el uso de pickle, ya que muchas partes de TensorFlow no lo admiten bien y, según nuestra experiencia, puede ser bastante frágil.
Referencia: https://www.tensorflow.org/alpha/guide/checkpoints
Los puntos de control de TensorFlow se pueden utilizar para guardar el valor de los parámetros periódicamente durante el entrenamiento. Esto puede resultar útil para guardar el progreso del entrenamiento en caso de que su programa falle o se detenga. Sonnet está diseñado para funcionar limpiamente con los puntos de control de TensorFlow:
checkpoint_root = "/tmp/checkpoints"
checkpoint_name = "example"
save_prefix = os . path . join ( checkpoint_root , checkpoint_name )
my_module = create_my_sonnet_module () # Can be anything extending snt.Module.
# A `Checkpoint` object manages checkpointing of the TensorFlow state associated
# with the objects passed to it's constructor. Note that Checkpoint supports
# restore on create, meaning that the variables of `my_module` do **not** need
# to be created before you restore from a checkpoint (their value will be
# restored when they are created).
checkpoint = tf . train . Checkpoint ( module = my_module )
# Most training scripts will want to restore from a checkpoint if one exists. This
# would be the case if you interrupted your training (e.g. to use your GPU for
# something else, or in a cloud environment if your instance is preempted).
latest = tf . train . latest_checkpoint ( checkpoint_root )
if latest is not None :
checkpoint . restore ( latest )
for step_num in range ( num_steps ):
train ( my_module )
# During training we will occasionally save the values of weights. Note that
# this is a blocking call and can be slow (typically we are writing to the
# slowest storage on the machine). If you have a more reliable setup it might be
# appropriate to save less frequently.
if step_num and not step_num % 1000 :
checkpoint . save ( save_prefix )
# Make sure to save your final values!!
checkpoint . save ( save_prefix )Referencia: https://www.tensorflow.org/alpha/guide/saved_model
Los modelos guardados de TensorFlow se pueden usar para guardar una copia de su red que está desacoplada de la fuente de Python. Esto se habilita guardando un gráfico de TensorFlow que describe el cálculo y un punto de control que contiene el valor de los pesos.
Lo primero que debe hacer para crear un modelo guardado es crear un snt.Module que desea guardar:
my_module = snt . nets . MLP ([ 1024 , 1024 , 10 ])
my_module ( tf . ones ([ 1 , input_size ]))A continuación, necesitamos crear otro módulo que describa las partes específicas de nuestro modelo que queremos exportar. Le recomendamos hacer esto (en lugar de modificar el modelo original in situ) para tener un control detallado sobre lo que realmente se exporta. Por lo general, esto es importante para evitar la creación de modelos guardados muy grandes y de modo que solo comparta las partes de su modelo que desee (por ejemplo, solo desea compartir el generador para una GAN pero mantener privado el discriminador).
@ tf . function ( input_signature = [ tf . TensorSpec ([ None , input_size ])])
def inference ( x ):
return my_module ( x )
to_save = snt . Module ()
to_save . inference = inference
to_save . all_variables = list ( my_module . variables )
tf . saved_model . save ( to_save , "/tmp/example_saved_model" ) Ahora tenemos un modelo guardado en la carpeta /tmp/example_saved_model :
$ ls -lh /tmp/example_saved_model
total 24K
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:14 assets
-rw-rw-r-- 1 tomhennigan 154432098 14K Apr 28 00:15 saved_model.pb
drwxrwsr-t 2 tomhennigan 154432098 4.0K Apr 28 00:15 variablesCargar este modelo es simple y se puede hacer en una máquina diferente sin el código Python que creó el modelo guardado:
loaded = tf . saved_model . load ( "/tmp/example_saved_model" )
# Use the inference method. Note this doesn't run the Python code from `to_save`
# but instead uses the TensorFlow Graph that is part of the saved model.
loaded . inference ( tf . ones ([ 1 , input_size ]))
# The all_variables property can be used to retrieve the restored variables.
assert len ( loaded . all_variables ) > 0 Tenga en cuenta que el objeto cargado no es un módulo de Sonnet, es un objeto contenedor que tiene los métodos específicos (por ejemplo, inference ) y propiedades (por ejemplo, all_variables ) que agregamos en el bloque anterior.
Ejemplo: https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb
Sonnet admite capacitación distribuida utilizando estrategias de distribución personalizadas de TensorFlow.
Una diferencia clave entre Sonnet y el entrenamiento distribuido usando tf.keras es que los módulos y optimizadores de Sonnet no se comportan de manera diferente cuando se ejecutan bajo estrategias de distribución (por ejemplo, no promediamos sus gradientes ni sincronizamos sus estadísticas de normas por lotes). Creemos que los usuarios deben tener control total de estos aspectos de su capacitación y no deben incluirse en la biblioteca. La compensación aquí es que necesita implementar estas características en su script de entrenamiento (normalmente son solo 2 líneas de código para reducir todos sus gradientes antes de aplicar su optimizador) o intercambiar módulos que tengan en cuenta explícitamente la distribución (por ejemplo, snt.distribute.CrossReplicaBatchNorm ).
Nuestro ejemplo distribuido de Cifar-10 explica cómo realizar un entrenamiento de múltiples GPU con Sonnet.
