Sonnet 2.0 : イントロダクション, Getting Started & 直列化 (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 06/10/2020
* 本ページは、Sonnet の以下のドキュメントを翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
イントロダクション
Sonnet は機械学習研究のための単純で、構成可能な抽象を提供するために設計された、TensorFlow 2 上に構築されたライブラリです。
Sonnet は DeepMind の研究者により設計されて構築されました。それは多くの様々な目的のためにニューラルネットワークを構築するために利用可能です (教師なし/あり学習、強化学習, …)。それが私達の組織のためには成功的な抽象であると見出しています、you might too !
より具体的には、Sonnet は単一の概念: snt.Module を中心とする単純でしかしパワフルなプログラミングモデルを提供します。Module はパラメータ、他のモジュールそしてユーザ入力にある関数を適用するメソッドへの参照を保持できます。Sonnet は多くの事前定義されたモジュール (e.g. snt.Linear, snt.Conv2D, snt.BatchNorm) と幾つかの事前定義されたネットワーク・モジュール (e.g. snt.nets.MLP) とともに出荷されますが、ユーザはまた自身のモジュールを構築することも奨励されます。
多くのフレームワークとは違い、Sonnet は貴方のモジュールをどのように使用するかについて非常に意固地ではありません。Module は自己充足的で他の一つから完全に分離されるように設計されています。Sonnet は訓練フレームワークとともに出荷されませんので、ユーザは自身のものを構築するか他の人により構築されたものを採用することが推奨されます。
Sonnet はまた理解するために単純であるようにも設計されていて、私達のコードは (願わくば！) 明瞭で焦点が合っています。デフォルト (e.g. 初期パラメータ値のためのデフォルト) を選択したとこでは何故かを指摘することを試みます。
Getting Started
サンプル
Sonnet を試す最も容易な方法は Google Colab を利用することです、これは GPU か TPU に装着された free な Python ノートブックを供給します。
- Predicting MNIST with an MLP
- Training a Little GAN on MNIST
- Distributed training with snt.distribute
インストール
始めるには TensorFlow 2.0 と Sonnet 2 をインストールします :
To get started install TensorFlow 2.0 and Sonnet 2:
$ pip install tensorflow-gpu 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 の殆どのモジュールはそれらのパラメータをある入力で最初に呼び出されるとき作成します (何故ならば殆どの場合パラメータの shape は入力の関数であるからです)。Sonnet モジュールはパラメータにアクセスするために 2 つのプロパティを提供します。
variables プロパティは与えられたモジュールにより参照される 総ての tf.Variables を返します :
all_variables = mlp.variables
注目すべき点は tf.Variables は単に貴方のモデルのパラメータのために使用されるだけではないことです。例えばそれらは snt.BatchNorm で使用されるメトリクスの状態を保持するために使用されます。殆どの場合ユーザはモジュール変数を取得してそれらを更新されるために optimizer に渡します。この場合非訓練可能変数は典型的にはそのリストにあるべきではありません、何故ならばそれらは異なるメカニズムを通して更新されるからです。TensorFlow は変数を「訓練可能」 (モデルのパラメータ) vs. 非訓練可能 (他の変数) として印をつける組込みメカニズムを持ちます。Sonnet はモジュールから総ての訓練可能な変数を集めるメカニズムを提供します、これは多分貴方が optimizer に渡すことを望むものです :
model_parameters = mlp.trainable_variables
貴方自身のモジュールを構築する
Sonnet はユーザに自身のモジュールを定義するために snt.Module をサブクラス化することを強く奨励します。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.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) ...)>)
上の variables 上の my_linear prefix に気付くかもしれません。これは Sonnet モジュールもメソッドが呼び出されるときはいつでもモジュール名前空間に入るためです。モジュール名前空間に入ることにより消費する TensorBoard のようなツールのための遥かにより有用なグラフを提供します (e.g. my_linear 内で発生する総ての演算は my_linear と呼ばれるグループにあります)。更にモジュールは今では TensorFlow チェックポイントと saved モデルをサポートします、これは後でカバーされる進んだ特徴です。
シリアライゼーション
Sonnet は複数のシリアライゼーション形式をサポートします。サポートする最も単純な形式は Python の pickle です、そして総ての組込みモジュールは同じ Python プロセスで pickle を通してセーブ/ロードできることを確実にするためにテストされています。一般には pickle の利用は推奨されません、それは TensorFlow の多くのパートで上手くサポートされません、そして経験的に非常に不安定である可能性があります。
TensorFlow チェックポイント
参照: https://www.tensorflow.org/guide/checkpoint
訓練の間に定期的にパラメータ値をセーブするために 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)
TensorFlow Saved モデル
参照: https://www.tensorflow.org/guide/saved_model
TensorFlow saved モデルはネットワークのコピーをセーブするために利用できます、これはそのための Python ソースから切り離されています。これは計算を記述する TensorFlow グラフと重みの値を含むチェックポイントをセーブすることにより可能になります。saved モデルを作成するために行なう最初のことはセーブすることを望む snt.Module を作成することです :
my_module = snt.nets.MLP([1024, 1024, 10]) my_module(tf.ones([1, input_size]))
次に、エクスポートすることを望むモデルの特定のパーツを記述するもう一つのモジュールを作成する必要があります。(元のモデルを in-place で変更するよりも) これを行なうことを勧めます、そうすれば実際にエクスポートされるものに渡る極め細かい制御を持ちます。これは非常に大きい saved モデルを作成することを回避するために典型的には重要です、そしてそのようなものとして貴方が望むモデルのパーツを共有するだけです (e.g. GAN のために generator を共有することを望むだけで discriminator は private に保持します)。
@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 フォルダで saved モデルを持ちます :
$ 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
このモデルのロードは単純で、saved モデルを構築したどのような 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 モジュールではないことに注意してください、それは前のブロックで追加した特定のメソッド (e.g. inference) とプロパティ (e.g. all_variables) を持つコンテナ・オブジェクトです。
分散訓練
サンプル: https://github.com/deepmind/sonnet/blob/v2/examples/distributed_cifar10.ipynb
Sonnet は カスタム TensorFlow 分散ストラテジー を使用して分散訓練のためのサポートを持ちます。
Sonnet と tf.keras を使用する分散訓練の間の主要な違いは Sonnet モジュールと optimizer は分散ストラテジーで動作するとき異なる動作をしないことです (e.g. 勾配を平均したりバッチ norm スタッツを同期しません)。ユーザは訓練のこれらの様相の完全な制御にあるべきでライブラリに焼き固められるべきではないと信じます。ここでのトレードオフはこれらの特徴を貴方の訓練スクリプトで実装するか (optimizer を適用する前に勾配を総て減じるために典型的にはこれは単に 2 行のコードです) 明示的に分散 aware なモジュール (e.g. snt.distribute.CrossReplicaBatchNorm) と交換する必要があります。
分散 Cifar-10 サンプルは Sonnet でマルチ GPU 訓練を行なうことをウォークスルーします。
以上