DeepMind の深層学習ライブラリ Sonnet を早速試してみた

どうも，Ryobot です．夜桜を眺めながら酒を呑みたい季節になって参りました．

追記 4/19
DeepMind から Differentiable Neural Computers の Sonnet 実装 が公開されました．今後も PathNet や Elastic Weight Consolidation 等の実装が公開されることを期待したいですね．

Sonnet は 4月 7日に公開されたばかりの DeepMind 謹製の深層学習ライブラリである．もともと DeepMind の社内で使用されていた TensorFlow のラッパーライブラリだったが，論文の実装を共有しやすくするためにオープンソースとして公開したようだ．Sonnet の最たる特徴として再利用可能なモジュールを複数回接続して計算グラフを構成することが挙げられる．Sonnet は TensorFlow Core の関数や他の高水準ライブラリ (TF Slim 等) と併用して利用できるので，同じくラッパーの Keras に比べて TensorFlow コードの残滓が見てとれる．汎用的だが冗長でもあった TensorFlow をニューラルネット向けに抽象化することによって RNN 等の複雑なネットワークを少ないコード量で可読的に記述できるわけである．

Sonnet は深層学習ライブラリ Touch にインスパイアされているようで，似たようなオブジェクト指向を採用している (標準ライブラリに TensorFlow を採用する前の DeepMind では Touch を採用していた)．より詳細を知りたい方は以下を参照されたい．

• プレスリリース: https://deepmind.com/blog/open-sourcing-sonnet/
• GitHub: https://github.com/deepmind/sonnet
• Qiita: DeepMindのSonnetを触ったので、TensorFlowã‚„Kerasと比較しながら解説してみた (じんべえざめさん)

ここでは，インストールから具体的な使用法・使用例を紹介したい．

インストール

リポジトリの README には Linux / Mac OS X および Python 2.7 と互換性があると明記されているが，プルリク Python3 compatibility? #10 を見ると数カ所の変更で一応 3系で正常に動作するっぽい．僕も Python 3.5 環境でマージ後のものを実行してみたが幾つかエラーが残っていた．まぁ今週中に examples くらいなら 3系で動くのでは？ (他力本願)

Sonnet のインストールには，bazel を使用してライブラリをコンパイルする必要がある．bazel のドキュメント

Mac であれば Homebrew で簡単に bazel をインストールできる．

$ brew install bazel

当然 TensorFlow も事前にインストールしておく．インストール方法はこちらを参照．注意点として最新バージョンの 1.0.1 以外では Sonnet のインストール時にエラーとなる．

# 最新バージョン以外をすでにインストールしている場合，
# TensorFlow をアンインストールする．
$ sudo pip uninstall tensorflow
# Mac の Python2.7 環境に TensorFlow の最新バージョン 1.0.1 をインストールする．
$ sudo pip install --upgrade https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.0.1-py2-none-any.whl
# もし "Cannot remove entries from nonexistent file ほげほげ easy-install.pth" 
# みたいなエラーが表示されたら  --ignore-installed つけたら良いかも．  
$ sudo pip install --upgrade https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.0.1-py2-none-any.whl --ignore-installed

Sonnet のインストール

# Sonnet のソースコードをダウンロード
$ git clone --recursive https://github.com/deepmind/sonnet

# configure の実行
$ cd sonnet/tensorflow
$ ./configure
$ cd ../

# インストールスクリプトを実行して wheel ファイルを作成する
$ mkdir /tmp/sonnet
$ bazel build --config=opt :install
$ ./bazel-bin/install /tmp/sonnet

# 生成した wheel ファイルをインストールする
# ここで最新バージョンの TensorFlow 1.0.1 以外だとエラーになる (1.0.0 はダメでした)
$ pip install /tmp/sonnet/*.whl

これでインストールは完了．

Sonnet の使用法

Sonnet の基本的な使い方は，Torch 風のオブジェクト指向を採用して順伝播を定義するモジュールを作成することである．具体的には sonnet.AbstractModule のサブクラスの Python オブジェクト (これを Modules と呼び，ニューラルネットの一部である) を構築し，これらを個別に複数回接続して計算グラフを構成する．モジュール内で宣言された各変数は接続呼び出しで自動的にラッパーのモジュールや計算グラフに共有される．また，変数の共有を制御する名前スコープや reuse= フラグなどの TensorFlow の低水準な機能はユーザーから抽象化 (隠蔽) される．

基本

import sonnet as snt
import tensorflow as tf

# 訓練データとテストデータ
train_data = get_training_data（）
test_data = get_test_data（）

# Linear モジュールは変数 (線形変換のための重みとバイアス) をもった全結合層．
lin_to_hidden = snt.Linear(output_size=FLAGS.hidden_size,
                                          name='inp_to_hidden')
hidden_to_out = snt.Linear(output_size=FLAGS.output_size,
                                            name='hidden_to_out')

# Sequential モジュールは与えられたデータに一連の内部モジュールや
# 操作 (op) を適応する．この例では Linear モジュール，シグモイド関数 op，
# Linear モジュールを順に適応している．sigmoid のような TF の低水準な
# op は変数を持たないので，Sonnet のモジュールと混ぜて使用できる．
mlp = snt.Sequential([lin_to_hidden, tf.sigmoid, hidden_to_out])

# Sequential は計算グラフに複数回接続できる．
train_predictions = mlp(train_data)
test_predictions = mlp(test_data)

# 重みとバイアスを初期化する．
# tf.truncated_normal(shape, mean=0.0, stddev=1.0,
#                                 dtype=tf.float32, seed=None, name=None)
# tf.truncated_normal は正規分布からのランダムサンプリング．
# 平均値から 2 標準偏差を超える値を削除し再選択する
initializers={"w": tf.truncated_normal_initializer(stddev=1.0),
              "b": tf.truncated_normal_initializer(stddev=1.0)}

# 重みとバイアスを正則化する．
# L1 正則化は重みをスパースにし，L2 正則化は訓練データの過学習を防ぐ．
# scale は正則化項の係数である．0.0 のとき正則化はなくなる．
regularizers = {"w": tf.contrib.layers.l1_regularizer(scale=0.1),
                "b": tf.contrib.layers.l2_regularizer(scale=0.1)}

# 初期化と正則化を追加した Linear モジュール．
linear_regression_module = snt.Linear(output_size=FLAGS.output_size,
                                      initializers=initializers,
                                      regularizers=regularizers)

# 計算された正則化の損失は tf.GraphKeys.REGULARIZATION_LOSSES
# という名前でコレクションに追加される．
# コレクションの正則化の損失を取得し，総和を求める．
graph_regularizers = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
total_regularization_loss = tf.reduce_sum(graph_regularizers)

# 通常の損失と正則化の損失の和を最小化する
train_op = optimizer.minimize(loss + total_regularizer_loss)

独自モジュール

さて，実際に snt.AbstractModule を継承した新しいクラスを作成して独自のモジュールを定義したい．クラス定義は Chainer っぽい書き方である．

このクラス (モジュール) のコンストラクタ __init__ はまずスーパークラスのコンストラクタ super(hoge, self).__init__ を呼び出してモジュール名 name を渡す．また __init__ の可変長引数の辞書型 name キーは常にリストの最後に与え，そのデフォルト値はスネークネーム (my_mlp のようなアンダーバー区切りの名前) のクラス名 (モジュール名) になる．

class MyMLP(snt.AbstractModule):
  def __init__(self, hidden_size, output_size,
               nonlinearity=tf.tanh, name="my_mlp"):
    super(MyMLP, self).__init__(name)
    self._hidden_size = hidden_size
    self._output_size = output_size
    self._nonlinearity = nonlinearity

次に _build() メソッドを実装する．_build() はモジュールが tf.Graph に接続されるたびに呼び出され，入力として複数のテンソル (空や 1つでもよい) を含んだ構造を受け取る．このとき複数のテンソルはタプルもしくは namedtuple で，その要素もまたテンソルやタプル / namedtuple となる．また入力テンソルはバッチとして与え，もしカラーチャネルがあれば最後の階に与える．ちなみに，リストや辞書の可変性はバグの原因になりやすいので利用がサポートされない．

def _build(self, inputs):
    """入力テンソルから出力テンソルを計算する．"""
    lin_x_to_h = snt.Linear(output_size=self._hidden_size, name="x_to_h")
    lin_h_to_o = snt.Linear(output_size=self._output_size, name="h_to_o")
    return lin_h_to_o(self._nonlinearity(lin_x_to_h(inputs)))

_build() メソッドの役割はおおむね以下の 3つある．

1. 内部モジュールの使用
2. コンストラクタに渡された既存のモジュールの使用
3. 変数の作成

変数は必ず tf.get_variable で作成する．tf.Variable コンストラクタを呼び出すとモジュールが初めて接続された時は動作するが，2回目の呼び出しはエラーになる．

上記のコードでは，_build() を呼び出すたびに新しい snt.Linear インスタンスが生成される．このとき作成される変数は互いに異なり値も共有されないというわけでなく，MLP を計算グラフに何度も接続したとしても 4つの変数 (1つの Linear につき重みとバイアスの 2つの変数) のみ作成される．

サブモジュールの宣言

Sequential のように，モジュールは外部で構築された他のモジュールを受け取って使用できる．

サブモジュールは親モジュールと呼ばれる他のモジュールのコード内部で構築されたモジュールである．LSTM を含むほとんどの実装では 1つ以上の Linear モジュールを内部に構築する．サブモジュールは変数スコープが正しいネストになるよう _build() に作成する．

class ParentModule(snt.AbstractModule):
  def __init__(self, hidden_size, name="parent_module"):
    super(ParentModule, self).__init__(name=name)
    self._hidden_size = hidden_size

  def _build(self, inputs):
    lin_mod = snt.Linear(self._hidden_size)  # サブモジュールの構築し
    return tf.relu(lin_mod(inputs))          # 接続する

Linear で作成される変数は parent_module/linear/w のような名前をもつ．

実用的な理由からコンストラクタ内でサブモジュールを構築したいユーザーもいるだろう．この場合は適切に変数をネストするために self._enter_variable_scope の呼び出し内でサブモジュールを構築しなければならない．

class OtherParentModule(snt.AbstractModule):
  def __init__(self, hidden_size, name="other_parent_module"):
    super(OtherParentModule, self).__init__(name=name)
    self._hidden_size = hidden_size
    with self._enter_variable_scope():  # この行が重要
      self._lin_mod = snt.Linear(self._hidden_size)  # ここにサブモジュールを構築

  def _build(self, inputs):
    return tf.relu(self._lin_mod(inputs))  # 事前に構築したモジュールを接続する

リカレントモジュール

Sonnet には TensorFlow のセル (cells) に相当するリカレントコアモジュールが用意され，1タイムステップの計算を実行する．これは TensorFlow の dynamic_rnn() 関数によって時間軸方向に展開できる．LSTM モジュールの場合は次のとおり．

hidden_size = 5
batch_size = 20
# input_sequence は [タイムステップ, バッチサイズ, 入力の特徴]
# サイズのテンソルである．
input_sequence = ...
lstm = snt.LSTM(hidden_size)
initial_state = lstm.initial_state(batch_size)
output_sequence, final_state = tf.nn.dynamic_rnn(
    lstm, input_sequence, initial_state=initial_state, time_major=True)

batch_size を渡した initial_state() メソッドは int32 型のテンソルになる．

最後に，独自のリカレントモジュールを定義したい．

リカレントモジュールは snt.AbstractModule と tf.RNNCell の両方を継承した snt.RNNCore のサブクラスである．この多重継承により Sonnet の変数共有モデルを使用でき，TensorFlow の RNN コンテナも使用できる．ようはいいとこ取りである．

class Add1RNN(snt.RNNCore):
  """リカレントの内部状態に 1 を追加し，出力として 0 を生成するシンプルなコア．
  
  このコアは以下を計算する．

  (`input`, (`state1`, `state2`)) -> (`output`, (`next_state1`, `next_state2`))
  
   (`state1`, `state2`) はバッチサイズの memory cell と hidden vector を表す．
  要素はすべてテンソルで，`next_statei` = `statei` + 1 と `output` = 0 が成り立つ．
  出力 (`output` and `state`) はすべて (`batch_size`, `hidden_size`) サイズである．
  """

  def __init__(self, hidden_size, name="add1_rnn"):
    """モジュールのコンストラクタ．

    Args:
      hidden_size: int型のモジュールの出力サイズ
      name: モジュール名
    """
    super(Add1RNN, self).__init__(name=name)
    self._hidden_size = hidden_size

  def _build(self, inputs, state):
    """1タイムステップの計算を実行する TensorFlow のサブグラフを構築する．"""
    batch_size = tf.TensorShape([inputs.get_shape()[0]])
    outputs = tf.zeros(shape=batch_size.concatenate(self.output_size))
    state1, state2 = state
    next_state = (state1 + 1, state2 + 1)
    return outputs, next_state

  @property
  def state_size(self):
    """バッチの階がない内部状態のサイズを返す．"""
    return (tf.TensorShape([self._hidden_size]),
            tf.TensorShape([self._hidden_size]))

  @property
  def output_size(self):
    """バッチの階がない出力サイズを返す．"""
    return tf.TensorShape([self._hidden_size])

  def initial_state(self, batch_size, dtype):
    """初期値 0 の内部状態を返す．

    NOTE: このメソッドは説明目的で明記している．
                スーパークラスに対応するメソッドがすでに存在している．
    """
    sz1, sz2 = self.state_size
    # 内部状態の shape の先頭にバッチサイズを追加し，zeros を作成する．
    return (tf.zeros([batch_size] + sz1.as_list(), dtype=dtype),
            tf.zeros([batch_size] + sz2.as_list(), dtype=dtype))

examples の rnn_shakespeare (多層 LSTM) を試してみた

rnn_shakespeare.py は多層 LSTM を使用した文字レベル言語モデル (一文字ずつ生成していく言語モデル) のソースコードである．

モデルのコードを見てみる．その他の訓練コードなどは一般的な TensorFlow と同様なのでここでは触れない．

class TextModel(snt.AbstractModule):
  """小さなシェイクスピアデータセットを使用した深層 LSTM モデル．"""

  def __init__(self, num_embedding, num_hidden, lstm_depth, output_size,
               use_dynamic_rnn=True, use_skip_connections=True,
               name="text_model"):
    """`TextModel` を構築する．
    Args:
      num_embedding: one-hot エンコード入力のあとに使用する埋め込み表現のサイズ．
      num_hidden: 各 LSTM 層の隠れユニット数．
      lstm_depth: LSTM 層の数．
      output_size: 深層 RNN の頂上にある出力層のサイズ．
      use_dynamic_rnn: TensorFlow の dynamic_rnn を利用するか否か．
      `False` は static_rnn を使用する．デフォルトは `True`.
      use_skip_connections: `snt.DeepRNN` で skip connections を
      利用するか否か．デフォルトは `True`.
      name: モジュール名．
    """

    super(TextModel, self).__init__(name=name)

    self._num_embedding = num_embedding
    self._num_hidden = num_hidden
    self._lstm_depth = lstm_depth
    self._output_size = output_size
    self._use_dynamic_rnn = use_dynamic_rnn
    self._use_skip_connections = use_skip_connections

    with self._enter_variable_scope():
      self._embed_module = snt.Linear(self._num_embedding,
                            name="linear_embed")
      self._output_module = snt.Linear(self._output_size,
                            name="linear_output")
      self._lstms = [
          snt.LSTM(self._num_hidden, name="lstm_{}".format(i))
          for i in range(self._lstm_depth)
      ] # lstm_depth 個の LSTM モジュールのリスト
      self._core = snt.DeepRNN(self._lstms,
                               skip_connections=True,
                               name="deep_lstm")

  def _build(self, one_hot_input_sequence):
    """深層 LSTM モデルのサブグラフを構築する．
    Args:
      one_hot_input_sequence: one-hot 表現としてエンコードされた入力
      シーケンスのテンソル．階は `[truncation_length, batch_size, output_size]` になる．
    Returns:
      output_sequence_logits: バッチの出力ロジットのテンソルのタプル．
      テンソルの階は `[truncation_length, batch_size, output_size]` になる．
      final_state: 時間展開したコアの final_state．
    """

    input_shape = one_hot_input_sequence.get_shape()
    batch_size = input_shape[1]
    
    # BatchApply モジュールはテンソルの最初の階にバッチの階を追加し，
    # 入力の次元にマッチするようにバッチの階の次元を分割する．  
    batch_embed_module = snt.BatchApply(self._embed_module)
    input_sequence = batch_embed_module(one_hot_input_sequence)
    input_sequence = tf.nn.relu(input_sequence)

    initial_state = self._core.initial_state(batch_size)

    if self._use_dynamic_rnn:
      output_sequence, final_state = tf.nn.dynamic_rnn(
          cell=self._core,
          inputs=input_sequence,
          time_major=True,
          initial_state=initial_state)
    else:
      # tf.unstack はテンソルの最初の階のスタックを解除し，テンソルのリストを返す．
      rnn_input_sequence = tf.unstack(input_sequence)
      output, final_state = tf.contrib.rnn.static_rnn(
          cell=self._core,
          inputs=rnn_input_sequence,
          initial_state=initial_state)
      output_sequence = tf.stack(output) # tf.stack はテンソルのリストをスタックする．

    batch_output_module = snt.BatchApply(self._output_module)
    output_sequence_logits = batch_output_module(output_sequence)

    return output_sequence_logits, final_state
  
  # @snt.experimental.reuse_vars をデコレートすると，
  # build() と同じように変数を再利用できる．
  @snt.experimental.reuse_vars
  def generate_string(self, initial_logits,
                                  initial_state, sequence_length):
    """サブグラフを構築し，モデルからサンプルされた文字列を生成する．
    Args:
      initial_logits: サンプリングした初期のロジット．
      initial_state: RNN コアの最初の内部状態．
      sequence_length: サンプルする文字の数．
    Returns:
      generated_string: 文字のテンソル．
      階は`[sequence_length, batch_size, output_size]` になる．
    """

    current_logits = initial_logits
    current_state = initial_state

    generated_letters = []
    for _ in xrange(sequence_length):
      # 分布から文字のインデックスをサンプルする．
      char_index = tf.squeeze(tf.multinomial(current_logits, 1))
      char_one_hot = tf.one_hot(char_index, self._output_size, 1.0, 0.0)
      generated_letters.append(char_one_hot)

      # deep_lstm に one-hot の文字を与え，ロジットを得る．
      gen_out_seq, current_state = self._core(
          tf.nn.relu(self._embed_module(char_one_hot)),
          current_state)
      current_logits = self._output_module(gen_out_seq)

    generated_string = tf.stack(generated_letters)

    return generated_string

訓練とテストを実行する．

FLAGS はデフォルト引数なので，10000 イテレーション，3 層のLSTM，バッチサイズ 32，埋め込みサイズ 32，隠れ層サイズ 128 である．

$ python rnn_shakespeare.py
INFO:tensorflow:Create CheckpointSaverHook.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.1 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.2 instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX instructions, but these are available on your machine and could speed up CPU computations.
W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX instructions, but these are available on your md could speed up CPU computations.
INFO:tensorflow:Saving checkpoints for 0 into /tmp/tf/rnn_shakespeare/model.ckpt.
INFO:tensorflow:0: Training loss 270.180481.
INFO:tensorflow:1: Training loss 248.135193.
INFO:tensorflow:2: Training loss 249.848511.
...
INFO:tensorflow:498: Training loss 106.377846.
INFO:tensorflow:Saving checkpoints for 500 into /tmp/tf/rnn_shakespeare/model.ckpt.
INFO:tensorflow:499: Training loss 103.404739.
...
INFO:tensorflow:9996: Training loss 76.107834.
INFO:tensorflow:9997: Training loss 79.831573.
INFO:tensorflow:9998: Training loss 79.587143.
INFO:tensorflow:Saving checkpoints for 10000 into /tmp/tf/rnn_shakespeare/model.ckpt.
INFO:tensorflow:9999: Training loss 82.172134. Validation loss 92.072212. Sample = b_0: 
jealouses|Wherevoctmen to rail this night, my earnest lie.|
Not say, fare you wench them: lest not lip fight:|
Then, to old Prince what false as heart:|
But sight God straight your lordship his tongue|
To shameliness smell there. Concey, I'll keep|
the deserves of substigues to me to me from mine:|
As ever aimish thee, to do no mannerers,|
Farewell. Let no strength-houser sorrow!|
Forbight but easy and stay all doing, Marcius,|
Romeo, most smaction upon our thanks|
The petition of princely man, that boldy's soldier;|
RIbe the rambroak and bird makes it have I din|
Wedly seal and meet live.||BENVOLIO:|
What with the tribunes? Why, this is the|
upon speaks of tell thou? I'll six him untutor'd their ease:|
'Tis tear in the boar of his infrict so drinks;|
And he harks out upon this issue.||Nurse:|
Ay, to do him in rise and death.|
Potness me to going and solely secret's|
And my pentice comes the that in, about a word,|
Ready with words, till the poor prevail'd too|
But counterft and happy hath told thee.|Art t
INFO:tensorflow:Reducing learning rate.
INFO:tensorflow:Test loss 117.373634

生成文をみてみると学習が足りないせいかそもそも英語になっていない．

おわりに

深層学習ライブラリは数多あるが最近では TensorFlow, Keras のほかに PyTouch, DyNet に勢いを感じる．PyTouch, DyNet (そして Chainer) はともに Define by Run (データの流れがネットワークの構造を決めるアプローチ) を採用している．Define by Run で作成されるグラフ構造は動的計算グラフ (Dynamic Computation Graphs) や動的ニューラルネットワーク (Dynamic neural networks) と呼ばれ最近やたら流行っている気がする．ちなみに TensorFlow, Caffe, Theano, Touch 等の他のライブラリは Define and Run (ネットワークの構造を決めてからデータを流すアプローチ) を採用している．どちらの設計思想もコーディングのしやすさ・固定グラフの保存や分散システム対応などでメリット・デメリットがあるので一概に甲乙をつけがたい．詳しくはこちらを参照.

TensorFlow ベースの Sonnet は Define and Run の立場であるが，RNN 等の複雑なネットワークをわりと簡単に書くことができる．このライブラリの真価はいまだ未知数だが，DeepMind の最先端の技術 (e.g., AlphaGo, DNC, EWC) が Sonnet 実装で公開される可能性を考えると今後の動向は追っておくのが無難であろう．

話が逸れるが今更 TensorFlow の TensorBoard を使ってみてすごく便利でびっくりした．下の動画は MNIST の埋め込みを t-SNE で可視化してみた様子である．

TensorFlow は最近 Tensor Fold で動的計算グラフにも対応し，Keras, Sonnet, Edward, seq2seq, tf.contrib.learn, TensorFlow-Slim 等の優秀なラッパー&高水準ライブラリを 6つも持っている．NARUTO でいったら長門 + ペイン六道みたいなチート級の強さだろう．