sonnet
v2.0.2
文檔|範例
Sonnet 是一個建構在 TensorFlow 2 之上的函式庫，旨在為機器學習研究提供簡單、可組合的抽象。
Sonnet 是由 DeepMind 的研究人員設計和建造的。它可用於建構用於許多不同目的的神經網路（無監督學習、強化學習…）。我們發現它對我們的組織來說是一個成功的抽象，您可能也是如此！
更具體地說，Sonnet 提供了一個簡單但功能強大的程式設計模型，該模型圍繞著一個概念： snt.Module 。模組可以保存對參數、其他模組和對使用者輸入應用某些功能的方法的參考。 Sonnet 附帶了許多預先定義的模組（例如snt.Linear 、 snt.Conv2D 、 snt.BatchNorm ）和一些預先定義的模組網路（例如snt.nets.MLP ），但也鼓勵使用者建立自己的模組。
與許多框架不同，Sonnet 對於如何使用模組沒有任何意見。模組被設計為獨立的並且彼此完全解耦。 Sonnet 不附帶培訓框架，鼓勵使用者建立自己的框架或採用其他人建立的框架。
Sonnet 的設計也易於理解，我們的程式碼（希望如此！）清晰且重點突出。在我們選擇預設值（例如初始參數值的預設值）的地方，我們嘗試指出原因。
嘗試 Sonnet 最簡單的方法是使用 Google Colab，它提供了一個連接到 GPU 或 TPU 的免費 Python 筆記型電腦。
snt.distribute進行分散式訓練若要開始安裝 TensorFlow 2.0 和 Sonnet 2：
$ pip install tensorflow tensorflow-probability
$ pip install dm-sonnet您可以執行以下命令來驗證安裝是否正確：
import tensorflow as tf
import sonnet as snt
print ( "TensorFlow version {}" . format ( tf . __version__ ))
print ( "Sonnet version {}" . format ( snt . __version__ ))Sonnet 附帶了許多您可以輕鬆使用的內建模組。例如，要定義 MLP，我們可以使用snt.Sequential模組來呼叫一系列模組，將給定模組的輸出作為下一個模組的輸入。我們可以使用snt.Linear和tf.nn.relu來實際定義我們的計算：
mlp = snt . Sequential ([
snt . Linear ( 1024 ),
tf . nn . relu ,
snt . Linear ( 10 ),
])要使用我們的模組，我們需要「呼叫」它。 Sequential模組（以及大多數模組）定義了一個__call__方法，這意味著您可以透過名稱呼叫它們：
logits = mlp ( tf . random . normal ([ batch_size , input_size ]))請求模組的所有參數也很常見。 Sonnet 中的大多數模組在第一次使用某些輸入呼叫時都會建立其參數（因為在大多數情況下，參數的形狀是輸入的函數）。 Sonnet 模組提供了兩個用於存取參數的屬性。
variables屬性傳回給定模組引用的所有tf.Variable ：
all_variables = mlp . variables值得注意的是tf.Variable不僅僅用於模型的參數。例如，它們用於保存snt.BatchNorm中使用的指標的狀態。在大多數情況下，使用者會檢索模組變數以將它們傳遞給優化器進行更新。在這種情況下，不可訓練的變數通常不應出現在該列表中，因為它們是透過不同的機制進行更新的。 TensorFlow 有一個內建機制，將變數標記為「可訓練」（模型的參數）與不可訓練（其他變數）。 Sonnet 提供了一種從模組中收集所有可訓練變數的機制，這可能是您想要傳遞給優化器的內容：
model_parameters = mlp . trainable_variablesSonnet 強烈鼓勵使用者子類化snt.Module來定義自己的模組。讓我們先建立一個名為MyLinear的簡單Linear圖層：
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 . b使用這個模組很簡單：
mod = MyLinear ( 32 )
mod ( tf . ones ([ batch_size , input_size ]))透過子類化snt.Module您可以免費獲得許多不錯的屬性。例如__repr__的預設實現，它顯示建構函數參數（對於調試和內省非常有用）：
>> > print ( repr ( mod ))
MyLinear ( output_size = 10 )您還可以獲得variables和trainable_variables屬性：
>> > mod . variables
( < tf . Variable 'my_linear/b:0' shape = ( 10 ,) ...) > ,
< tf . Variable 'my_linear/w:0' shape = ( 1 , 10 ) ...) > )您可能會注意到上面變數上的my_linear前綴。這是因為每當呼叫方法時，Sonnet 模組也會進入模組名稱範圍。透過輸入模組名稱範圍，我們為 TensorBoard 之類的工具提供了一個更有用的圖表來使用（例如，my_linear 內發生的所有操作都將位於名為 my_linear 的群組中）。
此外，您的模組現在將支援 TensorFlow 檢查點和儲存模型，這些是稍後介紹的進階功能。
Sonnet 支援多種序列化格式。我們支援的最簡單的格式是 Python 的pickle ，並且所有內建模組都經過測試，以確保它們可以在同一 Python 進程中透過 pickle 保存/載入。一般來說，我們不鼓勵使用 pickle，它沒有得到 TensorFlow 許多部分的良好支持，而且根據我們的經驗，它可能非常脆弱。
參考： https://www.tensorflow.org/alpha/guide/checkpoints
TensorFlow 檢查點可用於在訓練期間定期儲存參數值。這對於保存訓練進度很有用，以防程式崩潰或停止。 Sonnet 設計與 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 )參考： https://www.tensorflow.org/alpha/guide/saved_model
TensorFlow 保存的模型可用於保存與 Python 來源分離的網路副本。這是透過保存描述計算的 TensorFlow 圖和包含權重值的檢查點來實現的。
為了創建保存的模型，要做的第一件事是創建一個要保存的snt.Module ：
my_module = snt . nets . MLP ([ 1024 , 1024 , 10 ])
my_module ( tf . ones ([ 1 , input_size ]))接下來，我們需要建立另一個模組來描述我們想要導出的模型的特定部分。我們建議這樣做（而不是就地修改原始模型），以便您可以對實際導出的內容進行細粒度控制。這通常很重要，可以避免創建非常大的保存模型，這樣您就可以隻共享模型中您想要的部分（例如，您只想共享 GAN 的生成器，但保持鑑別器私有）。
@ 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" )現在，我們在/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 variables載入此模型非常簡單，可以在不同的電腦上完成，無需任何建立已儲存模型的 Python 程式碼：
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請注意，載入的物件不是 Sonnet 模組，而是一個容器對象，具有我們在上一個區塊中添加的特定方法（例如inference ）和屬性（例如all_variables ）。
範例： https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb
Sonnet 支援使用自訂 TensorFlow 分發策略進行分散式訓練。
Sonnet 和使用tf.keras的分散式訓練之間的一個關鍵區別是，Sonnet 模組和優化器在分佈策略下運行時的行為沒有不同（例如，我們不會平均您的梯度或同步您的批量標準統計數據）。我們認為使用者應該完全控制他們的培訓的這些方面，並且他們不應該被納入庫中。這裡的權衡是，您需要在訓練腳本中實現這些功能（通常這只是 2 行程式碼，以便在應用優化器之前減少梯度）或交換明確分佈感知的模組（例如snt.distribute.CrossReplicaBatchNorm 。
我們的分散式 Cifar-10 範例示範如何使用 Sonnet 進行多 GPU 訓練。
