sonnet
v2.0.2
Documentação | Exemplos
Sonnet é uma biblioteca construída com base no TensorFlow 2, projetada para fornecer abstrações simples e combináveis para pesquisas de aprendizado de máquina.
Sonnet foi projetado e construído por pesquisadores da DeepMind. Ele pode ser usado para construir redes neurais para diversos fins (aprendizagem não supervisionada, aprendizagem por reforço, ...). Achamos que é uma abstração de sucesso para nossa organização, você também pode!
Mais especificamente, o Sonnet fornece um modelo de programação simples, mas poderoso, centrado em um único conceito: snt.Module . Os módulos podem conter referências a parâmetros, outros módulos e métodos que aplicam alguma função na entrada do usuário. O Sonnet vem com muitos módulos predefinidos (por exemplo, snt.Linear , snt.Conv2D , snt.BatchNorm ) e algumas redes predefinidas de módulos (por exemplo, snt.nets.MLP ), mas os usuários também são incentivados a construir seus próprios módulos.
Ao contrário de muitos frameworks, o Sonnet é extremamente pouco opinativo sobre como você usará seus módulos. Os módulos são projetados para serem independentes e totalmente desacoplados uns dos outros. O Sonnet não é fornecido com uma estrutura de treinamento e os usuários são incentivados a construir os seus próprios ou a adotar aqueles criados por outros.
O Sonnet também foi projetado para ser simples de entender, nosso código é (espero!) claro e focado. Onde escolhemos padrões (por exemplo, padrões para valores iniciais de parâmetros), tentamos apontar o porquê.
A maneira mais fácil de experimentar o Sonnet é usar o Google Colab, que oferece um notebook Python gratuito conectado a uma GPU ou TPU.
snt.distribute Para começar, instale o TensorFlow 2.0 e o Sonnet 2:
$ pip install tensorflow tensorflow-probability
$ pip install dm-sonnetVocê pode executar o seguinte para verificar se tudo foi instalado corretamente:
import tensorflow as tf
import sonnet as snt
print ( "TensorFlow version {}" . format ( tf . __version__ ))
print ( "Sonnet version {}" . format ( snt . __version__ )) O Sonnet vem com vários módulos integrados que você pode usar trivialmente. Por exemplo, para definir um MLP podemos usar o módulo snt.Sequential para chamar uma sequência de módulos, passando a saída de um determinado módulo como entrada para o próximo módulo. Podemos usar snt.Linear e tf.nn.relu para realmente definir nosso cálculo:
mlp = snt . Sequential ([
snt . Linear ( 1024 ),
tf . nn . relu ,
snt . Linear ( 10 ),
]) Para usar nosso módulo precisamos “chamá-lo”. O módulo Sequential (e a maioria dos módulos) define um método __call__ que significa que você pode chamá-los pelo nome:
logits = mlp ( tf . random . normal ([ batch_size , input_size ]))Também é muito comum solicitar todos os parâmetros do seu módulo. A maioria dos módulos no Sonnet criam seus parâmetros na primeira vez que são chamados com alguma entrada (já que na maioria dos casos o formato dos parâmetros é uma função da entrada). Os módulos Sonnet fornecem duas propriedades para acessar parâmetros.
A propriedade variables retorna todos os tf.Variable s referenciados pelo módulo fornecido:
all_variables = mlp . variables É importante notar que tf.Variable s não são usados apenas para parâmetros do seu modelo. Por exemplo, eles são usados para manter o estado nas métricas usadas em snt.BatchNorm . Na maioria dos casos, os usuários recuperam as variáveis do módulo para passá-las a um otimizador para serem atualizadas. Neste caso, as variáveis não treináveis normalmente não devem estar nessa lista, pois são atualizadas através de um mecanismo diferente. O TensorFlow possui um mecanismo integrado para marcar variáveis como "treináveis" (parâmetros do seu modelo) versus não treináveis (outras variáveis). O Sonnet fornece um mecanismo para reunir todas as variáveis treináveis do seu módulo, que provavelmente é o que você deseja passar para um otimizador:
model_parameters = mlp . trainable_variables Sonnet incentiva fortemente os usuários a subclassificar snt.Module para definir seus próprios módulos. Vamos começar criando uma camada Linear simples chamada 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 é trivial:
mod = MyLinear ( 32 )
mod ( tf . ones ([ batch_size , input_size ])) Ao subclassificar snt.Module você obtém muitas propriedades interessantes gratuitamente. Por exemplo, uma implementação padrão de __repr__ que mostra argumentos de construtor (muito útil para depuração e introspecção):
>> > print ( repr ( mod ))
MyLinear ( output_size = 10 ) Você também obtém as variables e as propriedades trainable_variables :
>> > mod . variables
( < tf . Variable 'my_linear/b:0' shape = ( 10 ,) ...) > ,
< tf . Variable 'my_linear/w:0' shape = ( 1 , 10 ) ...) > ) Você pode notar o prefixo my_linear nas variáveis acima. Isso ocorre porque os módulos Sonnet também entram no escopo do nome do módulo sempre que métodos são chamados. Ao inserir o escopo do nome do módulo, fornecemos um gráfico muito mais útil para ferramentas como TensorBoard consumirem (por exemplo, todas as operações que ocorrem dentro de my_linear estarão em um grupo chamado my_linear).
Além disso, seu módulo agora oferecerá suporte ao checkpoint do TensorFlow e ao modelo salvo, que são recursos avançados abordados posteriormente.
Sonnet oferece suporte a vários formatos de serialização. O formato mais simples que suportamos é o pickle do Python, e todos os módulos integrados são testados para garantir que possam ser salvos/carregados via pickle no mesmo processo do Python. Em geral, desencorajamos o uso de pickle, pois ele não é bem suportado por muitas partes do TensorFlow e, em nossa experiência, pode ser bastante frágil.
Referência: https://www.tensorflow.org/alpha/guide/checkpoints
O checkpoint do TensorFlow pode ser usado para salvar o valor dos parâmetros periodicamente durante o treinamento. Isso pode ser útil para salvar o andamento do treinamento caso seu programa trave ou seja interrompido. O Sonnet foi projetado para funcionar de forma limpa com o checkpoint do 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 )Referência: https://www.tensorflow.org/alpha/guide/saved_model
Os modelos salvos do TensorFlow podem ser usados para salvar uma cópia da sua rede que está desacoplada da fonte Python para ela. Isso é habilitado salvando um gráfico do TensorFlow que descreve o cálculo e um ponto de verificação contendo o valor dos pesos.
A primeira coisa a fazer para criar um modelo salvo é criar um snt.Module que você deseja salvar:
my_module = snt . nets . MLP ([ 1024 , 1024 , 10 ])
my_module ( tf . ones ([ 1 , input_size ]))A seguir, precisamos criar outro módulo descrevendo as partes específicas do nosso modelo que queremos exportar. Aconselhamos fazer isso (em vez de modificar o modelo original no local) para que você tenha um controle refinado sobre o que realmente é exportado. Isso normalmente é importante para evitar a criação de modelos salvos muito grandes e para que você compartilhe apenas as partes do seu modelo que deseja (por exemplo, você deseja compartilhar apenas o gerador para um GAN, mas manter o discriminador privado).
@ 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" ) Agora temos um modelo salvo na pasta /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 variablesCarregar este modelo é simples e pode ser feito em uma máquina diferente sem nenhum código Python que construiu o modelo salvo:
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 Observe que o objeto carregado não é um módulo Sonnet, é um objeto contêiner que possui os métodos específicos (por exemplo, inference ) e propriedades (por exemplo, all_variables ) que adicionamos no bloco anterior.
Exemplo: https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb
Sonnet oferece suporte para treinamento distribuído usando estratégias de distribuição personalizadas do TensorFlow.
Uma diferença importante entre o Sonnet e o treinamento distribuído usando tf.keras é que os módulos e otimizadores do Sonnet não se comportam de maneira diferente quando executados sob estratégias de distribuição (por exemplo, não calculamos a média de seus gradientes nem sincronizamos suas estatísticas de norma de lote). Acreditamos que os usuários devem ter controle total sobre esses aspectos do seu treinamento e não devem ficar presos à biblioteca. A desvantagem aqui é que você precisa implementar esses recursos em seu script de treinamento (normalmente são apenas 2 linhas de código para reduzir seus gradientes antes de aplicar seu otimizador) ou trocar módulos que sejam explicitamente conscientes da distribuição (por exemplo, snt.distribute.CrossReplicaBatchNorm ).
Nosso exemplo distribuído do Cifar-10 mostra como fazer o treinamento multi-GPU com o Sonnet.
