HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

【TensorFlow】EagerExecutionがデフォルトになるぞ!!基本的な書き方を身につけておこう

 

 

follow us in feedly

https://cdn.blog.st-hatena.com/images/theme/og-image-1500.png

はじめに

9月の下旬に下記の記事を書きました。それなりに反響があったので、きっとTensorFlowの発展に興味を持っている人は数多くいるのだと思います。

www.hellocybernetics.tech

TensorFlowメジャーアップデートに伴う代表的な変化として「EagerExecution」がデフォルトになるという点があります。

Define and RunとDefine by Run

これまでTensorFlowと言えば、Chainer・PyTorchを筆頭とするDefine by Runというスタイルと比較され「Define and Run」であると言われてきました。 Define and Runとは、計算グラフを明確に記述してから、その計算グラフの計算処理を実行するというスタイルです(Define by Runは計算処理それ自体が計算グラフを動的に構築するスタイル)。

Define and Runでは、Pythonで計算グラフを書くのだけど、計算グラフを書き終えたらそれを一旦コンパイル(+最適化)して Pythonの外の世界で高速に計算できる環境を整えます。そしてコンパイルされた計算グラフに、適宜Pythonからデータを流してやるという方式を取ります。

Define by Runでは、計算グラフのコンパイルはせずに(ただし個々の行列演算等はNumpy・CupyやTorchなどの低レベル実装を使う)、Pythonの世界で計算処理もグラフの構築も行います。 なのでPythonで書けるようなグラフ構造は好きなように実現できるし、そもそもコードもPythonの書き方に準拠するのでわかりやすいというところがあり、非常に評判が良かったスタイルになります。

今回、世論の流れを組み(?)TensorFlowがDefine by Runとして生まれ変わります(もちろん必要に応じて、従来のDefine and Runでの記述もできるでしょうし、あるいはコンパイル・最適化された計算グラフへの変換やエクスポートもサポートするはずです)。 そこで、今回はTensorFlow流のDefine by Runの書き方についてざっとまとめておきましょう。

Eager Execution

TensorFlowの今までのスタイルと対比して「Eager Execution」と言っているだけで、その実体は要するに「NumpyやTorchと同じように使うことになりますよ」という話です。 簡単に、ただ足し算をするだけの計算グラフを書いて、実行してみましょう。

x = tf.placeholder(tf.float32, shape=[1, 1])
y = tf.placeholder(tf.float32, shape=[1, 1])
 
z = tf.add(x, y)

with tf.Session() as sess:
  print(sess.run(z, feed_dict={x: [[2.]], y: [[4.]]}))

# Will print [[6.]]

いちいち、ややこしいですね。これがEagerでは下記のようになります。

tf.enable_eager_execution()

x = [[2.]]
y = [[4.]]

z = tf.add(x, y)

print(z)

極めて直感的です。直感的というか、普通のPythonの書き方と変わりありません。私達が意識しなければならないのは、足し算をするときにtf.add()という関数を使うのだな、、、ということだけです(実際には x + y でも計算できるが)。 ただし、Eager Executionに入るためには、tf.enable_eager_execution()を宣言する必要があります。これは1回のセッションのただ1回だけ行えばいいので、Jupyter Notebookなどを利用している場合にはこのコードを書いたセルを何度も実行しないようにしましょう(エラーが出るだけです)。

モデルの書き方

線形回帰モデル

入力が多次元で、出力がスカラーであるような線形回帰モデルに使われるモデルを見てみましょう。

class LinearModel(tf.keras.Model):

  def __init__(self):
    super(LinearModel, self).__init__()
    self._hidden_layer = tf.keras.layers.Dense(1)

  def call(self, x):
    return self._hidden_layer(x)

いくつか注目すべき点は、まず、classを書くときにはtf.keras.Modelを継承するということです。 これは、今までKerasのfunctional APIを使う時のクラスと同じなのですが、Eagerに対応するためにtf.keras.Modelがかなり整備されました。 こいつがPyTorchでいうところのnn.ModuleやChainerで言うところのchainer.Chainを担ってくれます。

また、もう1つ注目すべきなのはself._hidden_layer = tf.keras.layers.Dense(1)という部分です。ここでは出力したい次元数1だけを渡しており、入力の次元を明示する必要はありません。 初めてtf.keras.layers.Dense(1)が呼びだされ、計算を実行するというタイミング(ここでは callメソッドに相当する)で入力されたデータの次元に合わせて重みパラメータを準備してくれます(Chainerにもある機能です)。

この機能はあってもなくても良いようなものですが、次の畳み込みニューラルネットワークを利用する際に便利です。

畳み込みニューラルネットワーク

次はMNISTを想定した適当なCNNを見てみます。

class Cnn(tf.keras.Model):

  def __init__(self, data_format):
    super(Cnn, self).__init__()
    
    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]
    
    layers = tf.keras.layers

    self.conv1 = layers.Conv2D(
        64, 5, padding='same', data_format=data_format, activation=tf.tanh)
    self.pool1 = layers.AveragePooling2D(2, 2, data_format=data_format)
    self.conv2 = layers.Conv2D(
        128, 5, data_format=data_format, activation=tf.tanh)
    self.pool2 = layers.AveragePooling2D(2, 2, data_format=data_format)
    self.flatten = layers.Flatten()
    self.fc1 = layers.Dense(1024, activation=tf.tanh)
    self.fc2 = layers.Dense(1, activation=None)

  def call(self, inputs):
    x = tf.reshape(inputs, self._input_shape)
    x = self.conv1(x)
    x = self.pool1(x)
    x = self.conv2(x)
    x = self.pool2(x)
    x = self.flatten(x)
    x = self.fc1(x)
    x = self.fc2(x)
return x

data_formatはTensorFlowではチャンネルのインデックスを最初(ただしバッチの次)にするか最後にするかを設定できるために存在するものです。ここはTFユーザなら知っていることなので流しましょう(確かGPUならチャンネルファーストのほうが良いんだっけか)。layers = tf.keras.layersは、基本的にこのモジュール内で実装された層しか利用しないために表記を省略しているだけです。

かつてはtf.layersや、他の3rdpartyによるlayerが多数存在したため、このような命名をすると混乱を招くところでしたが、今後は割とこの書き方が普通になるかもしれません(個人的にはモジュール名をちゃんと書くのが基本だと思いますが)。

さて、作法的なところは終わりしましょう。通常の畳み込みニューラルネットワークでは、畳み込み層が続いた後、最後のほうで全結合層を利用します。この結合層の入力に対応させるためには、畳込み層に合わせたデータの次元から、全結合層に合わせたデータの次元にしなければなりません。通常、畳み込み層に合わせたデータの形は(N, H, W, C)となっているため、(N, H * W * C)という形にしてやる必要があります。tf.keras.Flatten()クラスを当ててやれば自動で実行されます。

また、畳み込み層の次に位置する線形層は「入力の次元を明示しなくても良いため」、線形層を準備するときに畳み込み層の出口でデータの高さと幅がどうなっているかを予め精密に計算しておく必要はありません。

TensorFlowに従来より実装されていた、padding='same'も健在で、これのお陰で画像サイズを変更しないようなパディングを自動で実行してくれます(少なくとも私がPyTorchを使っていた時は、PyTorchはこのあたりは手計算だった)。

損失関数と勾配の書き方

TensorFlow Eager Executionで癖があるのはこの項目だと思われます。

ここまでの紹介は、要するにTFがPyTorchのような形式で実装が出来ますよというだけの話でした。全結合相が入力の次元を必要としなかったり、畳み込み層のパディングが便利だったりするのは、前々からTFに備わっていた機能そのままでありEager特有ではありません。

しかし、ここで説明する損失関数については明確に特殊な書き方をします。 PyTorchなどが損失関数の「クラス」を準備し、「メソッド」を使って勾配計算を実施するの言わばオブジェクト指向のモジュールになっているのに対して、TensorFlow Eager Executionは損失関数を高階関数で表して、部分適用するなどの形式を取ります(これは関数型の指向だと思われます)。

なぜこのようなモジュールの体系になっているのかは分かりません。TFがそもそも関数型指向(というか宣言的?)で記述を行っていたようにも見えますので、もしかしたら、このような形式を好んでいるだけかもしれませんし、従来の実装体系をEagerに移行する上で仕方のない実装方法だったのかもしれません。あるいは、柔軟な学習コードを書くためには、このような形式のほうが向いているのかもしれません。

理由は定かではありませんが、若干、慣れが必要かなと思われる部分です。

勾配の求め方

まずEagerでの勾配計算のスタンダードな方法を見てみましょう。

下記は非常に単純な$f(x)=x^3$の微分を求める実装方法です。

tfe = tf.contrib.eager

def f_cubed(x):
    return x**3
grad = tfe.gradients_function(f_cubed)
grad(3.)[0].numpy()
## =>27

上記は一体何をやっているかというと、Pythonの関数f_{\rm cubed}(x)で$x$を引数とする$f_{\rm cubed}(x)$を実装しておきます(こういう数学の関数を作ったと認識して構わない)。 この関数を、grad = tfe.gradients_function(f_cubed)と渡してやると、その(偏)導関数を返してくれるというわけです。

grad自身もこれまた関数なので(数学で言う導関数になっている)、この関数にgrad(3.)と渡すのは$\frac{d}{dx} f_{\rm cubed}(3)$と引数を与えていやっていることに相当します。grad(3.)[0][0]は単にIndexを指定しているだけです。numpy()メソッドはtf.Tensor型をnumpy.arrayに変換するメソッドです。

もしも単にEager Executionの実装方法が知りたいだけなら深追いする必要はありません。そういうもんだと思って構いませんが、基本的に関数型のパラダイムで使うので、若干Python利用者にとっては慣れが必要かもしれません。

仮に2変数関数にした場合は下記のように振る舞います。

def f2_cubed(x, y):
    return x**3 + y**2
grad = tfe.gradients_function(f2_cubed)
grad(3., 3.)[0].numpy()
## =>27
grad(3., 2.)[1].numpy()
## => 6

これはgrad(3., 2.)はリストとなっており、最初の要素にはxによる偏導関数にx=3, y=2を代入した値が、2つ目の要素にはyによる偏導関数にx=3, y=2を代入した値が入っていることを示しています。つまりindexが、何番目の変数の偏導関数になっているのかを指定しています。

最後に確認のために下記の例も見ましょう。

def f3_cubed(x, y):
    return x**3 + y**2 + 3*x*y
grad = tfe.gradients_function(f3_cubed)
grad(3., 3.)

さて、この場合grad(3, 3)にはどんな値が格納されているでしょうか(当然、最初の要素にはxによる偏導関数にx=3, y=2を代入した値が、2つ目の要素にはyによる偏導関数にx=3, y=2を代入した値が入っている)。

損失関数

今回はモデルと入力データと教師データを渡して、損失を返すPythonの関数を書きましょう。 損失関数は例えば下記のように地道にコードで書くこともできます(これは出力がスカラーの回帰の例)。

def mean_square_loss(model, xs, ys):
  return tf.reduce_mean(tf.square(tf.subtract(model(xs), ys)))

あるいはtf.keras.losses.mseなどを使って、

def mean_square_loss(model, xs, ys):
  return tf.keras.losses.mse(model(xs), ys)

と実装してしまってもいいでしょう。大抵のものはtf.keras.lossesモジュール内に実装されています。

学習コード

今回はモデルとデータセット、オプティマイザーを渡して1epochをすすめるfit関数を書いておきましょう。 ここでは損失関数とそのパラメータに拠る勾配の求め方、加えてパラメータの更新の仕方に焦点を当てるため、 ひとまずデータセットの準備の詳細は省きます(こういう雛形なのだなくらいに認識していただきたいです)。

def mean_square_loss(model, xs, ys):
  return tf.keras.losses.mse(model(xs), ys)

def fit(model, dataset, optimizer=tf.train.AdamOptimizer()):
  """
  Args:
    model: Model to fit.
    dataset: The tf.data.Dataset to use for training data.
    optimizer: The TensorFlow Optimizer object to be used.
  """
  tfe = tf.contib.eager

  loss_and_grads = tfe.implicit_value_and_gradients(mean_square_loss)

  # Training loop.
  for i, (xs, ys) in enumerate(tfe.Iterator(dataset)):
    loss, grads = loss_and_grads(model, xs, ys)
    optimizer.apply_gradients(grads)

tfe.implicit_value_and_gradientsにも先ほど作ったmean_square_loss関数を渡して関数を作っており、loss_and_gradsは、損失関数のパラメータに拠る勾配(とそのときの損失)を、model, xs, ysを入れたときに返してくれる関数です。ここで何がパラメータであるかは計算に関わっているデータがtfe.Variable型であるかを見て、自動で判断してくれます。 若干トリッキーではありますが、このような実装の仕方になります(おそらく、modelの部分等、場合によっては部分適用した関数をlambda式などで作って、関数を適宜小さくしたり、不要な繰り返しの代入を防いだりする実装ができるように、このような使い方を提供しているのではないかなと思っています…)。

勾配を計算するときにはloss_and_grads(xs, ys)とデータを渡してやって、返ってくるタプルの2つ目(index 1)にgradsが来ます。コイツをoptimizer.apply_gradients()メソッドに与えてやれば勾配を更新してくれます(ココで急に関数型っぽくない書き方になる…。)。

発展的事項

発展的と言いつつ、おそらく最も気になる実用的な部分になります。

Eager Executionでのバッチ正規化の書き方

今のとこPyTorchのmodel.train()model.eval()に相当するメソッドがtf.keras.Modelには存在しません。例えばバッチ正規化は訓練時と学習済のモデル利用時では振る舞いが異なるので、下記のように明示的に指定しなければなりません。

class ResnetIdentityBlock(tf.keras.Model):
  def __init__(self, kernel_size, filters):
    super(ResnetIdentityBlock, self).__init__(name='')
    filters1, filters2, filters3 = filters

    self.conv2a = tf.keras.layers.Conv2D(filters1, (1, 1))
    self.bn2a = tf.keras.layers.BatchNormalization()

    self.conv2b = tf.keras.layers.Conv2D(filters2, kernel_size, padding='same')
    self.bn2b = tf.keras.layers.BatchNormalization()

    self.conv2c = tf.keras.layers.Conv2D(filters3, (1, 1))
    self.bn2c = tf.keras.layers.BatchNormalization()

  def call(self, input_tensor, training=False):
    x = self.conv2a(input_tensor)
    x = self.bn2a(x, training=training)
    x = tf.nn.relu(x)

    x = self.conv2b(x)
    x = self.bn2b(x, training=training)
    x = tf.nn.relu(x)

    x = self.conv2c(x)
    x = self.bn2c(x, training=training)

    x += input_tensor
    return tf.nn.relu(x)

    
block = ResnetIdentityBlock(1, [1, 2, 3])

## 訓練時
block(input_tensor, training=True)

## 評価時
block(input_tensor, training=False)

あるいは暫定的には

class ResnetIdentityBlock(tf.keras.Model):
  def __init__(self, kernel_size, filters):
    super(ResnetIdentityBlock, self).__init__(name='')
    ...同上省略
    self.training = False

  def call(self, input_tensor):
    x = self.conv2a(input_tensor)
    x = self.bn2a(x, training=self.training)
    x = tf.nn.relu(x)

    x = self.conv2b(x)
    x = self.bn2b(x, training=self.training)
    x = tf.nn.relu(x)

    x = self.conv2c(x)
    x = self.bn2c(x, training=self.training)

    x += input_tensor
    return tf.nn.relu(x)

  def to_train(self):
    self.training = True

  def to_eval(self):
    self.training = False


block = ResnetIdentityBlock(1, [1, 2, 3])

## 訓練時
block.to_train()
block(input_tensor)

## 評価時
block.to_eval()
block(input_tensor, training=False)

などと実装しておくことです。しかし、容易に想像がつくようにResBlockはResNetのコンポーネントとして使われることになるので、後々、

class ResNet(tf.keras.Model):
  def __init__(self, kernel_size, filters):
    super(ResNet, self).__init__(name='')
    self.block1 = ResnetIdentityBlock(1, [1, 2, 3])
    self.block2 = ResnetIdentityBlock(1, [1, 2, 3])
    self.training = False

  def call(self, input_tensor):
    x = self.block1(input_tensor)
    x = self.block2(input_tensor)
    return x

  def to_train(self):
    self.training = True

  def to_eval(self):
    self.training = False

としたのでは上手く動作しません(ResNetの中の各self.blockResNetself.trainingを見ているわけではない。なんとか紐付ける方法が必要であるが、PyTorchの場合、nn.Moduleapplyメソッドなるものがあり上手く働く。ここらへんはTFは未整備っぽい)。

tf.keras.Sequentialも健在

my_seq = tf.keras.Sequential([tf.keras.layers.Conv2D(1, (1, 1)),
                               tf.keras.layers.BatchNormalization(),
                               tf.keras.layers.Conv2D(2, 1, 
                                                      padding='same'),
                               tf.keras.layers.BatchNormalization(),
                               tf.keras.layers.Conv2D(3, (1, 1)),
                               tf.keras.layers.BatchNormalization()])

みたいな書き方はEagerでも可能。こちらの書き方でのBatchNormalizationの振る舞いは要調査。

最後に

とりあえず、基本的なモデルの構築方法と学習コードの書き方はこれを参考に作ることは可能でしょう。

必要に応じてloss関数を自分で書き換えたり、モデルを書き換えたり、必要なepoch数だけfit関数をループで包んだりすればすぐに活用できるはずです。ちなみにデータがtf.data.Datasetクラスで包まれている前提に書いていますが、ミニバッチを取り出せる方法があればなんでも構いませんので、自前のイテレータを作っても構いません。が、tf.data.Datasetはジェネレータになっているようで(たぶん)、大規模データでもメモリを殺さず云々カンヌン…ごにょごにょ(勉強不足ですすいません)

TensorFlowには他にも便利なAPIが多数存在するため、高レベルから低レベルまで柔軟な実装が可能になった今、最も便利に使えるフレームワークではないだろうかと私は思っております。   

今後もうちょっと勉強して記事を書けたら書きます。