この記事はTensorFlow Advent Calendar 2017の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で動作するのでここを見ながら用意します。Eagerの使い方は以下のようなまじないを入れておくだけです。以降もtensorflowはtf、eagerはtfeで統一していきます。
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のドキュメントを調べると、どうやらEagerで使うLayersを管理するために利用しているみたいですね。
余談にはなるのですが、__init__で入力するTensorのチャンネルを場合分けしていますが、これはここに書いてあるようにGPUとCPUで最適な持ち方が違うんですね。勉強になりました。
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をぐるぐる回してます。
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などで利用するかもですね。
# 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)
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
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の実装などすごい勢いで進化してるのでこれからも目が離せないですね。