DeepMindのTensorFlowライブラリSonnetを使う

(2017-08-06)

AlphaGoを開発したGoogle DeepMind社のTensorFlowライブラリSonnetを使う。 当初はPython2しか対応していないようだったけど、今は3にも対応している。

準備

TensorFlowを使うライブラリはほかにもいくつかあるのだけど、 Kerasと比較してみると、 KerasがTensorFlowの部分を完全にラップしているのに対して、 Sonnetは必要に応じてTensorFlowの関数も呼ぶ、比較的抽象度が低いライブラリのようだ。

SonnetとTensorFlowとPython3入りイメージをDockerHubに上げた。 Dockerfileはここ

内容は基本的にREADME通りだけど、 configureのところで対話的に聞かれないように前もって環境変数で設定を与えている。 あとは、TensorFlowのビルドに使われているGCCのバージョンが古いようで、sonnetをimportするときに以下のエラーが出たため、bazelのオプションに--copt="-D_GLIBCXX_USE_CXX11_ABI=0"を付けている。

tensorflow.python.framework.errors_impl.NotFoundError: /usr/local/lib/python3.5/dist-packages/sonnet/python/ops/_resampler.so: undefined symbol: _ZN10tensorflow7strings6StrCatB5cxx11ERKNS0_8AlphaNumES3_S3_S3_

起動。

$ docker run -itd --name sonnet -p 6006:6006 -p 8888:8888 sambaiz/sonnet
$ docker logs sonnet
...
   Copy/paste this URL into your browser when you connect for the first time,
    to login with a token:
        http://localhost:8888/?token=*****

Jupyter Notebookを開いてSonnetとTensorFlowがimportできることを確認する。

import sonnet as snt
import tensorflow as tf
snt.resampler(tf.constant([0.]), tf.constant([0.]))
# => <tf.Tensor 'resampler/Resampler:0' shape=(1,) dtype=float32>

MNIST

TensorFlowのチュートリアルのデータを使って、畳み込みを行わない簡単なMNISTをやってみる。 このデータはtrain、validation、test用に最初から分かれていて、 それぞれピクセル濃度配列の画像データと、その画像がどの数字なのかを表すone-hot vectorのラベルを含んでいる

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
train, validation, test = mnist
print(train.images[0]) # ピクセルの濃度を[0,1]の値で表した配列: [0, 0, ..., 0.41568631  0.6156863, 0.99607849, ...]
print(len(train.images[0])) # 28 * 28 = 784
print(train.labels[0]) # 正解のみ1のone-hot vector: [ 0.  0.  0.  0.  0.  0.  0.  1.  0.  0.]
images, labels = mnist.train.next_batch(100)

Sonnetではニューラルネットワークの一部をModuleとして表現し、それらをTensorFlowの計算グラフに接続していく。 Moduleはグラフに複数回接続することができ、中の変数は共有される。 素のTensorFlowだと変数のスコープを作って共有するのに reuse=Trueでtf.variable_scopeしてtf.get_variableしたりする必要があるけど、そのあたりは抽象化されているので tf.Variableを含むような処理はModuleで行う。

Linear Moduleは 重みの乗算とバイアスの加算をするもの。 これにtf.Sigmoidのような活性化関数を適用するのを繰り返し、最後に出力層とつなげるとMulti Layer Perceptronを構築できる。

import sonnet as snt
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

FLAGS = tf.flags.FLAGS

tf.flags.DEFINE_integer("hidden_size", 100, "Size of hidden layer.")
tf.flags.DEFINE_integer("output_size", 10, "Size of output layer.")

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

x = tf.placeholder(tf.float32, [None, 784])
y_ = tf.placeholder(tf.float32, [None, 10])
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')
mlp = snt.Sequential([lin_to_hidden, tf.sigmoid, hidden_to_out, tf.nn.softmax])
y = mlp(x)
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(1000):
        images, labels = mnist.train.next_batch(100)
        sess.run(train_step, feed_dict={x: mnist.train.images, y_: mnist.train.labels})
    correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    print(sess.run(accuracy, feed_dict={x: mnist.test.images, y_: mnist.test.labels}))
    # => 0.9307

Moduleを作る

Moduleを作るにはsnt.AbstractModuleを継承し、 スーパークラスのコンストラクタを呼んで、グラフに接続されるたびに呼ばれる_buildメソッドを実装する。

class MyMLP(snt.AbstractModule):
  """test mlp module"""
  def __init__(self, hidden_size, output_size,
               nonlinearity=tf.sigmoid, name="my_mlp"):
    """hidden_size & output_size is required"""
    super(MyMLP, self).__init__(name=name)
    self._hidden_size = hidden_size
    self._output_size = output_size
    self._nonlinearity = nonlinearity
  
  def _build(self, inputs):
    """Compute output Tensor from input Tensor."""
    lin_to_hidden = snt.Linear(output_size=self._hidden_size, name='inp_to_hidden')
    hidden_to_out = snt.Linear(output_size=self._output_size, name='hidden_to_out')
    return snt.Sequential([lin_to_hidden, self._nonlinearity, hidden_to_out, tf.nn.softmax])(inputs)

このModuleを使うとこんな感じ。 exampleのように データセットもModuleにすることができる。

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

mymlp = MyMLP(hidden_size=FLAGS.hidden_size, output_size=FLAGS.output_size)

x = tf.placeholder(tf.float32, [None, 784])
y_ = tf.placeholder(tf.float32, [None, 10])
y = mymlp(x)
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for i in range(1000):
        images, labels = mnist.train.next_batch(100)
        sess.run(train_step, feed_dict={x: mnist.train.images, y_: mnist.train.labels})
    correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    print(sess.run(accuracy, feed_dict={x: mnist.test.images, y_: mnist.test.labels}))
    # => 0.9307