TensorFlow Eager Executionの紹介

この記事は[TensorFlow Advent Calendar 2017](https://qiita.com/advent-calendar/2017/tensorflow)のDay8の記事になります。

# TensorFlow Eager
## 概要
国内で根強い人気を持っているChainerやここ1年くらいで海外で大きくユーザを増やしたPyTorchでお馴染みのDefine-by-Runでモデルを記述することができるTensorFlowのパッケージの一つ。記事投稿時点ではまだPreview releaseとなっています。

## Define-by-Run
これまでDefine-and-RunだったTensorFlowがなぜDefine-by-Runを導入するのか？Define-by-Runには次のようなメリットがあるので、RNN系の実装が多い方に好まれています。

* 動的なネットワークを構築することが出来る (Ex. mini-batch毎に処理が違っても良い)
* モデルのエラーのデバッグが行いやすい

また、その反面以下のようなデメリットも存在します。

* グラフを構成するのが動的なのでオーバーヘッドが発生する
* 計算の最適化が行いにくい

## 基本的な使い方
Eagerは現在nightly packageで動作するので[ここ](https://github.com/tensorflow/tensorflow#installation)を見ながら用意します。Eagerの使い方は以下のようなまじないを入れておくだけです。以降もtensorflowは`tf`、eagerは`tfe`で統一していきます。

```python3

import tensorflow as tf
import tensorflow.contrib.eager as tfe

tfe.enable_eager_execution()
```

# MNIST Example
コードはexamplesを見ればわかるので、これまでと異なる部分だけ解説します。
コードまで出しては説明しないですが、dataの読み出しは1.4.0からcontribが外れた`tf.data`を利用します。これも便利なのでEagerじゃなくても使ってみてください。

## Model Definition
Chainerに近いような`__call__`の中にモデルを書いていくような記述になっていますね。これまでのTensorFlowのモデルと異なるのは、Dropoutのテスト時の対応が直感的にわかりやすい`if`文で書かれていますね。後、もう一点気になるのが`layer`の前に挟まっている`self.track_layer`かと思います。これは`tfe.Network`の[ドキュメント](https://www.tensorflow.org/versions/master/api_docs/python/tf/contrib/eager/Network)を調べると、どうやらEagerで使う`Layers`を管理するために利用しているみたいですね。

余談にはなるのですが、`__init__`で入力するTensorのチャンネルを場合分けしていますが、これは[ここ](https://www.tensorflow.org/performance/performance_guide#data_formats)に書いてあるようにGPUとCPUで最適な持ち方が違うんですね。勉強になりました。

```python3
class MNISTModel(tfe.Network):

  def __init__(self, data_format):
    super(MNISTModel, self).__init__(name='')
    if data_format == 'channels_first':
      self._input_shape = [-1, 1, 28, 28]
    else:
      assert data_format == 'channels_last'
      self._input_shape = [-1, 28, 28, 1]
    self.conv1 = self.track_layer(
        tf.layers.Conv2D(32, 5, data_format=data_format, activation=tf.nn.relu))
    self.conv2 = self.track_layer(
        tf.layers.Conv2D(64, 5, data_format=data_format, activation=tf.nn.relu))
    self.fc1 = self.track_layer(tf.layers.Dense(1024, activation=tf.nn.relu))
    self.fc2 = self.track_layer(tf.layers.Dense(10))
    self.dropout = self.track_layer(tf.layers.Dropout(0.5))
    self.max_pool2d = self.track_layer(
        tf.layers.MaxPooling2D(
            (2, 2), (2, 2), padding='SAME', data_format=data_format))

  def call(self, inputs, training):
    x = tf.reshape(inputs, self._input_shape)
    x = self.conv1(x)
    x = self.max_pool2d(x)
    x = self.conv2(x)
    x = self.max_pool2d(x)
    x = tf.layers.flatten(x)
    x = self.fc1(x)
    if training:
      x = self.dropout(x)
    x = self.fc2(x)
    return x
```

## Mini-batch Learning
これまでと違うのは1epoch内のループ(mini-batchを回す)もfor文で書いてます。tfe.IteratorにDatasetを渡すことでbatchをぐるぐる回してます。

```python3
def loss(predictions, labels):
  return tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(
          logits=predictions, labels=labels))

def train_one_epoch(model, optimizer, dataset, log_interval=None):
  tf.train.get_or_create_global_step()

  def model_loss(labels, images):
    prediction = model(images, training=True)
    loss_value = loss(prediction, labels)
    tf.contrib.summary.scalar('loss', loss_value)
    tf.contrib.summary.scalar('accuracy',
                              compute_accuracy(prediction, labels))
    return loss_value

  for (batch, (images, labels)) in enumerate(tfe.Iterator(dataset)):
    with tf.contrib.summary.record_summaries_every_n_global_steps(10):
      batch_model_loss = functools.partial(model_loss, labels, images)
      optimizer.minimize(
          batch_model_loss, global_step=tf.train.get_global_step())
      if log_interval and batch % log_interval == 0:
        print('Batch #%d\tLoss: %.6f' % (batch, batch_model_loss()))
```

以下のように書くとgradientを出すこともできます。Distributedなどで利用するかもですね。

```python3
# optimizer.minimize(
#     batch_model_loss, global_step=tf.train.get_global_step())

grads = tfe.implicit_gradients(model_loss)(images, labels)
optimizer.apply_gradients(grads)
```

## PyTorchのMNIST Exampleとの比較
すでにあるDefine-by-RunのPyTorchとコードを比較してみます。全体を通して純粋なDefine-by-Runで設計されたPyTorchの方がコード量は少なく書けそうです。しかし、1.0以前のTensorFlowのことを考えるとコード量はかなり少なくなったと思います。

### Model Definition (PyTorch)

```python3
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)
```

### Mini-batch Learning

```python3
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.data[0]))
```
 
# おわりに

PyTorchの比較を見るとDefine-by-Runで書くならPyTorchの方が良さそうに思ってしまいますが、Googleのブログには次のような一文が載っています。

```
In the near future, we will provide utilities to selectively convert portions of　your model to graphs. In
this way, you can fuse parts of your computation (such as internals of a custom RNN cell) for high-
performance, but also keep the flexibility and readability of eager execution.
```

最初に書いた通り、パフォーマンスの観点からDefine-by-Runにはデメリットがありますが、ここに書かれた通りであれば、近いうちにパフォーマンスとEagerの可読性や柔軟性を両立することが可能になるのかもしれません。
実際に`tfe.Network`は`tf.layers.Layer`を実装しているので、examplesのresnet50にあるようなネットワークの一部として別のネットワークを利用することが出来ます。(`tf.keras`も使えますね。)

TensorFlowが最初にリリースされて以降、多くDeep Learningのフレームワークがリリースされました。どれも多種多様ですが、TensorFlowはDistributedの実装、Kerasの取り込み、Define-by-Runの実装などすごい勢いで進化してるのでこれからも目が離せないですね。


# 参考
* [Research Blog: Eager Execution: An imperative, define-by-run interface to TensorFlow](https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html)
* [eager execution examples](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples)
* [深層学習フレームワークChainerの特徴](https://www.slideshare.net/unnonouno/chainer-59664785)