Deep Learning Code Golfやってみた part03 Tensorflowへの移植

前回（PyTorch）で終わりのつもりだったのですがちょっとだけ続編です。

PyTorchと並び著名な深層学習フレームワークとして、Tensorflowがあります。私はPyTorchを使うことがほとんどですが、TensorflowでもDLCGを行うとどんな違いがあるか検証しました。

結論から言えば、PyTorchを利用する場合と解答となるモデルそのものは変わりませんでした。畳み込み層のほうが全結合層より短く書けるというような変化はないようです。

環境構築

環境構築はPyTorchより簡単です。学習ループを自前で書いたりライブラリを入れなくてもmodel.fit()を呼ぶだけで自動的にループを回してくれます。そのため、コードの長さの計測機能を除けば、以下のコードだけで学習・評価ができます。データオーグメンテーションがないなど、PyTorchと若干違うため同じモデル構造でも正解率に若干の差があります。

import tensorflow as tf
tf.random.set_seed(1) # 乱数固定
# データセット読み込み
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
# モデル構築
model = m() # プレイヤーが定義するモデル
# 学習設定
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5, batch_size=256)
# 正解率の評価
test_loss, test_acc = model.evaluate(x_test, y_test)

PyTorchとの相違点

実際のコードを示しながら、PyTorchとの違いを解説します。

最初に、Tensorflowでの通常のモデル定義コードを示します。SequentialやFlattenはPyTorchとほぼ同じです。ただし、Sequentialは可変長引数ではなくレイヤーのリストを引数にとります。PyTorchでのLinearはDenseに変わり、入力チャンネル数は不要となります。すなわちDenseはPyTorchのLazyLinear相当です。もちろん、ResNetのような分岐があるモデルを定義するための、クラスを用いたモデル定義方法もありますが明らかにコード量が増大します（あとで触れます）。

画像データのデータ配置が、PyTorchでは[バッチサイズ, チャンネル, 幅, 高さ]なのに対し、Tensorflowではデフォルトでは[バッチサイズ, 幅, 高さ, チャンネル]となっています。この違いは学習済みモデルを移植したりReshapeを行ったりする場合に影響が出るものの、本稿の範囲では影響ありませんでした。

import tensorflow as tf
def m():
    return tf.keras.models.Sequential([
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.Dense(10) 
    ])

正解率を問わない最短のコード（38.08%, 85文字）を示します。PyTorchではSequentialとFlattenは同じモジュール（torch.nn）からインポートできたのですが、Tensorflowでは、tf.keras.models.Sequentialとtf.keras.layers.Flattenというように別々のモジュールからインポートする必要があります。ここで、tf.keras.models.Sequentialはtf.keras.Sequentialとしてもインポートできるため、from tensorflow.keras import*でSequentialをインポートしつつ、同時にlayersにtf.keras.layersがインポートされるようにするのが最短と考えれます。インポートの煩雑さが影響して、コードの長さはPyTorchの67文字から85文字に増加してしまいました。

from tensorflow.keras import*
L=layers
m=lambda:Sequential([L.Flatten(),L.Dense(10)])

次に、複数の全結合層を利用するコード（43.41%, 103文字）を示します。Tensorflowでの活性化関数は、Denseへのパラメータ引数を用いてDense(99,activation='elu')と記述することが可能ですが、それよりも単体のELUクラスを用いるほうが短くなります。

from tensorflow.keras import*
L=layers
D=L.Dense
m=lambda:Sequential([L.Flatten(),D(99),L.ELU(),D(10)])

畳み込みを用いたコードを示します。順に51.93%, 108文字、60.22%, 110文字です。畳み込みはConv2Dで、引数はConv2D(filters, kernel_size, strides=(1,1), ...)です。入力チャンネル数は不要です。データオーグメンテーションの違い等により（過学習気味）、チャンネル数99では正解率60%に達しなかったためチャンネル数が3桁になっています。

from tensorflow.keras import*
L=layers
m=lambda:Sequential([L.Conv2D(9,3),L.Flatten(),L.ReLU(),L.Dense(10)])

from tensorflow.keras import*
L=layers
m=lambda:Sequential([L.Conv2D(512,3),L.Flatten(),L.ReLU(),L.Dense(10)])

最後に、ループを導入した深いモデル（70.26%, 148文字）を示します。PyTorchと同様、出力チャンネル数が10より大きくてもエラーになりません。チャンネル数68は、正解率70%以上となるようパラメータの探索を行った結果です。Sequentialの引数は可変長引数ではなくリストなので、リストの結合を用いて最後のFlattenを付加しています。

from tensorflow.keras import*
L=layers
S=Sequential
m=lambda:S([S([L.Conv2D(68,3,i),L.BatchNormalization(),L.ReLU()])for i in[1,2]*3]+[L.Flatten()])

Residual構造を記述する

PyTorchで書かれたConvMixerの解説で登場した、Residual構造( $y=f(x)+x$ )をTensorflowでも短く書いてみます。Tensorflowで分岐があるモデルを記述する方法は2つあり、Functional APIによる方法と、サブクラス化による方法と呼ばれます。

まず、Functional APIでのモデル記述例を示します。

from tensorflow.keras import *
from tensorflow.keras.layers import *
def m():
    i=Input([32,32,3])
    j=Flatten()(i)
    k=Dense(128)(j)
    l=ReLU()(k)
    m=Dense(10)(l)
    return Model(i,m)

Functional APIでは、モデルの入力を表す変数をInputで生成し、これをレイヤーオブジェクトに渡してその出力を表す変数を順次生成していくことでモデルの構造を表します。最後に、モデル全体の入出力となる変数をModel関数に渡すことで、モデルが完成します。変数は複数回使用できるので、分岐を表現することができます。l=ReLU(k)+kと記述すれば、非常にシンプルにResidual構造が表現できます。Residualレイヤーを定義するのではなく、モデル定義文の中で2回同じ変数を使うことで表現します。レイヤー同士のつながりを表す変数を明示的に扱えることがシンプルさの要因です。ただし、変数を明示的に代入して取り回さないといけないため代入文が必須で、リスト内包表記（＝式）で深いモデルを表現することが難しくなります。DLCG的には使いづらいでしょう。

最後に最も柔軟性の高いモデル定義方法であるサブクラス化による方法を示します。Residualレイヤーの記述例を示します。PyTorchのときと考え方は同じで、Sequentialクラスを継承し、レイヤー実行時に呼び出されるcallメソッドをオーバーライドしてResidual構造を実装します。PyTorchにおけるforwardをcallに、self[0]をself.layers[0]に置換した形となります。

from tensorflow.keras import *
from tensorflow.keras.layers import *
class Residual(Sequential):
    def call(self, inputs):
        return self.layers[0](inputs) + inputs

# 使用例
r = Residual([ReLU()])
x = tf.constant([1.0, -1.0])
r(x)
# => [ 2., -1.]

DLCGらしく、このコードを短く記述した結果を示します。

from tensorflow.keras import *
from tensorflow.keras.layers import *
R=type('x',(Sequential,),{'call':lambda s,x,**d:s.layers[0](x)+x})

type関数を利用するのはPyTorchと同じなのですが、Tensorflowでは余分な記述が必要でした。第1引数'x'ですが、空文字列にすると内部のコードでエラーが発生します。クラス名の1文字目が'_'かどうかで挙動を変える機構があり、1文字目がないためIndexErrorが生じます。ラムダ式の引数には、キーワード引数を受け取る**dが必要でした。モデルの実行が学習中か推論中かを示すtrainingキーワード引数（Dropout等の挙動が変わる）が渡されるため、これを受け取らないとエラーになります。しかし、素直にclass構文を用いた定義ではキーワード引数を定義しなくてもエラーになりません。この理由は、Tensorflow内部の実装がメソッドが受け取る引数を列挙し、定義に合わせて適切な引数を渡すように実装されているためです。しかしながらラムダ式でメソッドを定義することが想定されていないのか、判定が間違ってしまい{'call':lambda s,x:s.layers[0](x)+xという定義に対してtrainingキーワード引数を与えてしまいエラーになっています。

以上のように、TensorflowでもPyTorchとほぼ同じようにDLCGを遊ぶことができることがわかりました。"tensorflow.keras"が"torch.nn"より長いことや、気を利かせて挙動を変える機構の誤判定等で、PyTorchよりコードが長くなってしまう傾向があることがわかりました。