sonnet
v2.0.2
ドキュメント|例
Sonnet は、機械学習研究にシンプルで構成可能な抽象化を提供するように設計された TensorFlow 2 上に構築されたライブラリです。
Sonnet は、DeepMind の研究者によって設計および構築されました。さまざまな目的 (教師なし学習、強化学習など) のニューラル ネットワークを構築するために使用できます。これは私たちの組織にとって成功した抽象化であることがわかりました。あなたもそうかもしれません。
具体的には、Sonnet は、 snt.Moduleという 1 つの概念を中心としたシンプルかつ強力なプログラミング モデルを提供します。モジュールは、ユーザー入力に何らかの機能を適用するパラメーター、他のモジュール、メソッドへの参照を保持できます。 Sonnet には多くの定義済みモジュール (例: snt.Linear 、 snt.Conv2D 、 snt.BatchNorm ) といくつかの定義済みモジュール ネットワーク (例: snt.nets.MLP ) が同梱されていますが、ユーザーは独自のモジュールを構築することも推奨されます。
多くのフレームワークとは異なり、Sonnet はモジュールの使用方法について非常に偏見がありません。モジュールは自己完結型であり、相互に完全に分離されるように設計されています。 Sonnet にはトレーニング フレームワークが同梱されていないため、ユーザーは独自に構築するか、他の人が構築したフレームワークを採用することが推奨されます。
Sonnet は理解しやすいように設計されており、コードは (願わくば!) 明確で焦点が絞られています。デフォルトを選択した場合 (初期パラメータ値のデフォルトなど)、その理由を指摘しようとします。
Sonnet を試す最も簡単な方法は、GPU または TPU に接続された無料の Python ノートブックを提供する Google Colab を使用することです。
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 モジュールは、パラメータにアクセスするための 2 つのプロパティを提供します。
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 の使用はお勧めしません。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 トレーニングの実行について説明します。
