TensorFlow2.x式の低レベルな書き方

世界中で最も使われていて最も現場で使われているTensorFlow(以下TF)が去年、機能拡張を含んだ大型バージョンアップを行いました。その際にAPIの大掃除が行われTF1.x系列で使っていた低レベルな書き方が通用しなくなりました

この記事は低レベルでゴリゴリ色んな所をカスタマイズして書きたいって人を支援できたら幸いです(自分自身もよくわかってない点もある)

1.前話

・TF2.xで生き残っているTF1.xの書き方

TF1.xとTF2.x両方に通用する書き方はKerasのFunctional APIとSequential APIです。
それ以外の書き方は消えました。というかTensorFlow自体が「Kerasをベースに書いてね」という感じです。tf.layersとかは消されて使えません
また、TF1.xからTF2.xに更新した際、まともに変更せずに動くのはKerasのみ(カスタムレイヤーとか使わずに)を使って書かれたコードだとおもいます(変換プログラムもあるけど使えるのかね)

※TF1.xから更新してTF2.xのKerasを動かす場合、多分Kerasのインポートを以下の様に変更する必要があります。

from tensorfrow import keras

・どうしてそんなひどい事をしたの？

答えは単純です
「PyTorchとかの方がPythonic(Pythonぽくて)ぽくって、書きやすいから」です
あとPyTorchの追い上げがすごいってのもある

(引用:https://trends.google.co.jp/trends/explore?date=today%205-y&q=TensorFlow,PyTorch,Keras,Chainer)

あとPyTorchとかKerasはオブジェクト指向ぽくかけて、TF1.xは関数型(tf.layersとかもろ関数型)ぽいからってのもあります。またはDefine on Run(TensorFlow)とDefine by Run(PyTorch)の違いもあったりします
Define on RunとDefine by Runの解説はこの記事ではしませんが、気になる方は調べていただければ

「Kerasでオブジェクト指向的なモデルの書き方は無いんじゃない？」って感のいい読者なら気づいたかもしれません。確かにTF1.xのKerasにはオブジェクト指向的な書き方は存在しません、TF2.xのKerasにはオブジェクト指向的なモデルの書き方ができます。本記事はその点にも触れます

2.本編

・TensorFlow2.0の書き方

TensorFlow2.0はKerasの作者によるColab Notebookによると以下の書き方があります。
この記事では最も左下の書き方を解説します。

・Fully flexible with TensorFlow2.0(最も低レベルな書き方)

TT2.xで最も低レベルな書き方は以下になります
　・レイヤー :tf.keras.layers.layerを継承したクラスで書く
　・レイヤーブロック :tf.keras.layers.layerかtf.keras.Modelを継承したクラスで書く
　・モデル :tf.keras.Modelを継承したクラスで書く
　・訓練イテレーション:クラスを宣言してその中で回す
一番上以外はちょっとピンとこないと思います。一個ずつ解説します

・カスタムレイヤー :tf.keras.layers.layerを継承したクラスで書く
TF1.xのカスタムレイヤーの書き方がそのまま使えます。
基本的に __iniit(self)__で必要な変数を宣言してbuild(self,input_shape)で使うレイヤーと重み、バイアスを宣言します。以下例
(引用と参考:https://www.tensorflow.org/tutorials/customization/custom_layers)

class MyLayer(tf.keras.layers.layer):
　　　def __iniit__(self, output, **kwargs):
　　　　　super(MyLayer, self).__init__(**kwargs)
　　　　　self.output = output
　　　def build(self, input_shape):
　　　　　self.kernel = self.add_variable("kernel",shape=[　　　
　　　　　　　　　　　　　　　　　　　　int(input_shape[-1]),
　　　　　　　　　　　　　　　　　　　　self.num_outputs])
　　　def __call__(self, inputs):
　　　　　return tf.matmul(inputs, self.kernel)

・レイヤーブロック:tf.keras.layers.layerかtf.keras.Modelを継承したクラスで書く
CNNとかで使われている層を重ねたブロック。上のカスタムレイヤーも入れて書くことも可能
tf.keras.layers.layerで書く場合は重み、バイアスの宣言は行わないようにする
tf.keras.Modelでも書ける(モデルの入れ子ができるため)けど入力の型が取れないから正直微妙である。以下例

class MyLayerBlock_layer(tf.keras.layers.layer):
　　　def __iniit__(self, filter, kernel_size, channel_axes=-1 **kwargs):
　　　　　super(MyLayerBlock_layer, self).__init__(**kwargs)
　　　　　self.filter = filter
　　　　　self.kernel_size = kernel_size
　　　　　self.channel_axes = channel_axes　
　　　def build(self, input_shape):
　　　　　input_filter = input_shape[self.channel_axes].value
　　　　　self.conv2_a = tf.keras.layers.Conv2D(input_filter, (1, 1))
　　　　　self.bn2_a = tf.keras.layers.BatchNormalization()
　　　　　self.depthconv2 = tf.keras.layers.DepthwiseConv2D(input_filter,
　　　　　　　　　　　　　　　　　　　　　　　　　　　　　self.kernel_size,
　　　　　　　　　　　　　　　　　　　　　　　　　　　　　padding='same')
　　　　　self.bn2_depth = tf.keras.layers.BatchNormalization()
　　　　　self.conv2_b = tf.keras.layers.Conv2D(self.filter, (1, 1))
　　　　　self.bn2_b = tf.keras.layers.BatchNormalization()
　　　def __call__(self, inputs, training=True):
　　　　　x = self.conv2_a(input_tensor)
　　　　　x = self.bn2a(x, training=training)
　　　　　x = tf.nn.relu(x)
　　　　　x = self.depthconv2(x)
　　　　　x = self.bn2_depth(x, training=training)
　　　　　x = tf.nn.relu(x)
　　　　　x = self.conv2_b(x)
　　　　　x = self.bn2_b(x, training=training)
　　　　　x += inputs
　　　　　return tf.nn.relu(x)

class MyLayerBlock_Model(tf.keras.Model):
　　　def __iniit__(self,
　　　　　　　　　　input_filter,
　　　　　　　　　　filter,
　　　　　　　　　　kernel_size,
　　　　　　　　　　**kwargs):
　　　　　super(MyLayerBlock_Model, self).__init__(**kwargs)
　　　　　self.conv2_a = tf.keras.layers.Conv2D(input_filter, (1, 1))
　　　　　self.bn2_a = tf.keras.layers.BatchNormalization()
　　　　　self.depthconv2 = tf.keras.layers.DepthwiseConv2D(input_filter,
　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　kernel_size,
　　　　　　　　　　　　　　　　　　　　　　　　　　　　　padding='same')
　　　　　self.bn2_depth = tf.keras.layers.BatchNormalization()
　　　　　self.conv2_b = tf.keras.layers.Conv2D(filter, (1, 1))
　　　　　self.bn2_b = tf.keras.layers.BatchNormalization()
　　　def __call__(self, inputs, training=True):
　　　　　x = self.conv2_a(input_tensor)
　　　　　x = self.bn2a(x, training=training)
　　　　　x = tf.nn.relu(x)
　　　　　x = self.depthconv2(x)
　　　　　x = self.bn2_depth(x, training=training)
　　　　　x = tf.nn.relu(x)
　　　　　x = self.conv2_b(x)
　　　　　x = self.bn2_b(x, training=training)
　　　　　x += inputs
　　　　　return tf.nn.relu(x)

・モデル:tf.keras.Modelを継承したクラスで書く

Kerasのベースモデルクラスを使ってモデルを作ります。これでPyTorchっぽく書けてPythonicなモデルになります(オブジェクト指向的)
このモデルに対してcompileやpredictなどのメソッドを使うことができるが、作成したモデルを使ってtf.keras.models系の関数は使えません
このモデルに対して複数の入力を与える場合、それぞれの1軸をバッチ次元にしたリストを渡して__call__の時に分解します

class MyModel(tf.keras.Model):
　　　def __iniit__(self,output,**kwargs):
　　　　　super(MyModel, self).__init__(**kwargs)
　　　　　self.dense = tf.keras.layers.Dense(output)
　　　def __call__(self, inputs):
　　　　　x = self.dense(inputs)
　　　　　return tf.nn.relu(x)

・訓練イテレーション:クラスを宣言してその中で回す
これが一番の肝です。以下のようにタスククラスを宣言しモデルの初期化などを行った後訓練メソッドを呼び出すようにします
以下の場合だとtrainメソッドを呼び出せば訓練が開始されます。
何故、こんな面倒くさい事をやってるかというと@tf.functionの自動グラフ構築を必要な所で使うことによって最適かつ高速に訓練を行う事ができます(GoogleのTransformerでも使われてるテクニック)
Define on RunとDefine by Runを同時に操作できる強みがある

class MyTask(object):
　　　def __iniit__(self,model,optimizer,loss_func,acc_func):
　　　　　self.model = model
　　　　　self.optimizer = optimizer
　　　　　self.loss_func = loss_func
　　　　　self.acc_func = acc_func
　　　　　self.model.compile() #必要ならコンパイルしとく
　　　def train(self, epoch):
　　　　　for iter in range(epoch):
　　　　　　　x, y = (データの取得用の関数とかクラス)
　　　　　　　loss=train_step( x, y )
　　　　　@tf.function
　　　　　def train_step(self, x, y)
　　　　　　　with tf.GradientTape() as tape:
　　　　　　　　　　pred = self.model(x, training=True)
　　　　　　　　　　loss = self.loss_func(y, pred)
　　　　　　　grd = tape.gradient(loss, self.model.trainable_weights)
　　　　　　　self.optimizer.apply_gradients(zip(grd,
　　　　　　　 self.model.trainable_weights))
　　　　　　　self.acc_func.update_state(y, pred)
　　　　　　　return loss
　　　@tf.function
　　　def model_predict(self,x)
　　　　　return self.model.predict(x)

3.おわりに

これらのを利用することでTF1.x系でも、できた細かなカスタムや調整もできるようになったと思います
最後にTensorFlow v2.x系はまだちょっと未完成なところ(ver2.0では勾配計算に失敗したり、今でのカスタムレイヤーで複数出力を計算時にバグがあったり)があるので研究等で使う場合は十分に気を付けてくださいませ