TensorFlow EagerモードとPyTorchの学習コードと速度の比較

はじめに

Eagerモードを触っていこうということで、最近は以下の記事を書きました。

今回はMNISTを使って学習を回していきます。 

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com

題材としては下記の記事のAutoEncoderをPyTorchとEagerでそれぞれ実装して、Pythonのループ部分は全く同じになるように書きました。果たして、どのような差がでるのでしょうか。

aidiary.hatenablog.com

用いるデータと設定

今回は学習率を0.001、バッチサイズを128、epoch数を100に固定しました。

また、sklearnのdatasetを使って以下のようにデータを取り出しています。

from sklearn.datasets import fetch_mldata
from sklearn.cross_validation import train_test_split
from sklearn.utils import shuffle
import numpy as np

mnist = fetch_mldata('MNIST original', data_home=".")
X = mnist.data
y = mnist.target

X = X.astype(np.float64)
X /= X.max()

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)

num_data = X_train.shape[0]
num_epochs = 100
batch_size = 128
learning_rate = 0.001

validationデータは用意していません。ハイパーパラメータのチューニングなどしないので、一発勝負です（そもそも性能を競うつもりは無いです）。

前処理としては各ピクセルが0~1の値を取るように正規化しています。他にも-1~1で正規化したり、平均引いて分散で割る標準化などの方法があります。

また、EagerでもPyTorchでもMSELossを各iteration毎に表示しています。これは特にルールはないと思いますが、iteration毎に記録しておけば、あとからepoch毎のlossを算出することはできますし（必要かは知らん）、確率的勾配降下法の不安定さ（ミニバッチの選ばれ方による学習の進み方の変動）を見るにはiteration毎の方が良いと思います。

PyTorchでのAutoEncoder

モデルと学習

aidiaryさんのほぼ丸パクリで下記のようにAutoEncoderを実装。

MNISTのデータは28×28ピクセルなので、これを4回に分けてLinear層を使って2次元へ圧縮します。活性化関数はReLUです。EncoderとDecoderのノードの数は対象になるようにしています。記事書きながら気づきましたが、EncoderとDecoderのつなぎ目に活性化関数が入っていません。まあ対して問題にならないでしょう（たぶん）。

正則化、ドロップアウト、バッチ正規化などは一切用いません。

class Autoencoder_pytorch(nn.Module):
    
    def __init__(self):
        super(Autoencoder_pytorch, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(True),
            nn.Linear(128, 64),
            nn.ReLU(True),
            nn.Linear(64, 12),
            nn.ReLU(True),
            nn.Linear(12, 2))
        
        self.decoder = nn.Sequential(
            nn.Linear(2, 12),
            nn.ReLU(True),
            nn.Linear(12, 64),
            nn.ReLU(True),
            nn.Linear(64, 128),
            nn.ReLU(True),
            nn.Linear(128, 28 * 28),
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

model_pytorch = Autoencoder_pytorch()
if cuda:
    model_pytorch.cuda()

AutoEncoderの出力が入力自身に近づくようにMSEを最小化する学習を行います。コードは以下のようになっています。今回はPyTorchのデータローダは使いません。実直にnumpyのデータをシャッフルして、バッチサイズ分取り出し、Variableに変換し…というコードを書きました。

これは学習以外の部分をなるべくTensorFlowと揃えたかったためです。

学習開始と終了時に時間を記録して、学習に要した時間も記録しておきます。

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model_pytorch.parameters(),
                             lr=learning_rate)

loss_list = []
start_pytorch = time.time()
for i in range(100):
    ## データをシャッフル
    X, y = shuffle(X_train, y_train)

    for idx in range(0, num_data, batch_size):
        ## データをバッチサイズ分取り出す
        batch_x = X[idx: idx + batch_size 
            if idx + batch_size < num_data else num_data]
               
        batch_x = Variable(torch.from_numpy(batch_x).float()).cuda()
        xhat = model_pytorch(batch_x)
        loss = criterion(xhat, batch_x)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # logging
        loss_list.append(loss.data[0])
        print('loss : {}'.format(loss.data.cpu().numpy()))
end_pytorch = time.time()
learning_time_pytorch = end_pytorch - start_pytorch

結果

結果は以下のコードで表示します。

print(learning_time_pytorch)
z = model_pytorch.encoder(Variable(torch.from_numpy(X_train).float()).cuda())
Z = z.data.cpu().numpy()
plt.figure(figsize=(12, 10))
plt.scatter(Z[:,0], Z[:,1], c=y_train, cmap='Accent')
plt.colorbar()
plt.figure(figsize=(12, 5))
plt.plot(loss_list)

私の環境ではPyTorch利用時、学習終了までに197秒要しました。

うーん、GTX1060だとこんなものなのでしょう。

PyTorchのEncoderは28×28の訓練データを以下のように2次元に圧縮しました。教師なしで実行している割にはある程度分離ができているように見えます。0と6や、9と7が手書きでは似たような見た目になるので、再構成可能なように2次元上に配置した場合には近くなることも想定されますが、ちゃんとそれっぽい気がします（4の扱いは謎）。

学習の経過としては以下のような感じ。100epochもいらなさそうです。

さて、テストデータの方をどのように2次元に圧縮するのかを見てみると以下のようになっていました。訓練データとほとんど変わらない配置になっていますね。これが最も良い圧縮であるかは分かりませんが、訓練データだけしか圧縮できないというわけではなく、ちゃんとテストデータも同じように圧縮してくれるようです。

TensorFlow EagerのAutoEncoder

モデルと学習

PyTorchと同じノード数で同じ活性化関数を利用。torch.nn.Sequentialに相当するクラスはEagerにはないので、仕方なくそれっぽいもので代用します（これをやっておくと実装は楽になると思われます）。

class Sequential(tfe.Network):
    def __init__(self, *args):
        super(Sequential, self).__init__()        
        self.layer_list = [self.track_layer(layer) \
                           for layer in args]
        
    def call(self, x):
        for layer in self.layer_list:
            x = layer(x)
        return x

    
class Autoencoder(tfe.Network):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = Sequential(
            tf.layers.Dense(128, activation=tf.nn.relu),
            tf.layers.Dense(64, activation=tf.nn.relu),
            tf.layers.Dense(12, activation=tf.nn.relu),
            tf.layers.Dense(2))
        
        self.decoder = Sequential(
            tf.layers.Dense(12, activation=tf.nn.relu),
            tf.layers.Dense(64, activation=tf.nn.relu),
            tf.layers.Dense(128, activation=tf.nn.relu),
            tf.layers.Dense(28*28))
        
    def call(self, x):
        z = self.encoder(x)
        y = self.decoder(z)
        return y

model = Autoencoder()

TensorFlowのeagerモードではtf.layersのクラスをself.track_layerで包むことで、計算グラフを記録してくれているようです。Sequentialクラスの方で全部包まれるようにしていますが、普通に実装をする際には各々self.track_layerで包むことを忘れないようにしましょう。

続いて学習についてですが、eagerモードではlossをpython関数で書いておく流儀の模様。tfe.implicit_value_and_gradientsにはPythonで定義した関数を渡しておきます。

学習時に、Pythonで定義した時の引数を実際に渡してやると、ロスの値と勾配と各パラメータを全部ひっくるめて返してくれる便利なやつなようです。optimizerで更新するときには勾配とパラメータの情報が必要になりますがtfe.implicit_value_and_gradientsの戻り値の２つ目がタプルで勾配とパラメータを渡してくれるので、これをそのまま横流しにするだけで大丈夫です。

あと注意事項として、GPUを使うということをwith tf.deviceで明示してあげなければなりません。ここが意外と面倒なところかも。

PyTorchのときと同様に時間を計測しておきます。

def loss(x, model):
    return tf.reduce_mean(tf.square(model(x) - x))
    
value_and_grad_f_w = tfe.implicit_value_and_gradients(loss)
optimizer = tf.train.AdamOptimizer(learning_rate)

loss_log = []
start = time.time()
for i in range(100):
    ## データをシャッフル
    X, y = shuffle(X_train, y_train)

    for idx in range(0, num_data, batch_size):
        ## データをバッチサイズ分取り出す
        batch_x = X[idx: idx + batch_size 
            if idx + batch_size < num_data else num_data]
        
        
        with tf.device('/gpu:0'):
            batch_x = tf.convert_to_tensor(batch_x, dtype=tf.float32)
            loss, grads_and_vars = value_and_grad_f_w(batch_x, model)
            loss_log.append(loss.numpy())
            optimizer.apply_gradients(grads_and_vars)
        print('loss : {}'.format(loss.numpy()))
end = time.time()
learning_time = end - start

結果

コードは以下

with tf.device('/gpu:0'):
    z = model.encoder(tf.convert_to_tensor(X_train, dtype=tf.float32))

print(learning_time)
Z = z.numpy()
plt.figure(figsize=(12, 10))
plt.scatter(Z[:,0], Z[:,1], c=y_train, cmap='tab10')
plt.colorbar()
plt.figure(figsize=(12, 5))
plt.plot(loss_log)

 

学習に掛かった時間は約548秒。おっそ！！9分！！暇すぎてチャーハンを作ってしまいました。

 

学習データは以下のように2次元に埋め込まれました。0が飛び散っています。他は概ね分かれているようですが、0のせいでよく見えん…。

 

 

 

lossの方はPyTorchとほとんど変わらずと言ったところです（これを見て、ひとまずあのSequentialの書き方でもeagerの学習ができていることは確認できました。実際、tensorflowがどのような実装を得てdefine by runを達成しているのかよく分からん）。

 
 
 
テストデータの方も学習データと同じような配置になりました（相変わらず0が飛び散っています）。まあ初期値の問題やミニバッチの選ばれ方次第で学習はどうなるか分かりませんね。 

 

 

 

結論

PyTorch：197秒

TF-eager : 548秒

追試

TF-eager（tf.data.Datasetとtfe.Iterator利用）：542秒

TF-graph(InteractiveSettion, sess.run(feed_dict)) : 151秒 

※TF-graphは従来のDefine and Runのやり方tf.data.Datasetを利用した学習も試しました

遅いのはnumpyとTensorFlowを繋ぐ部分ではなさそう。単純にeagerがまだ未完成だと見るのが良いのかもしれません。あるいはかなりeagerのお気持ちを理解しなければ、PyTorchやTF-graphくらいのパフォーマンスを出せないのだろうか。

学習自体は進んでいるので、根本的に間違っているというわけではなさそうですが、TensorFlowから見て何か非効率なコードになっている可能性は十分にあります。流石に3倍遅いってのはなあ…。

おまけと追試たち

Sequentialをやめて実直に

バッチデータを取ってきたりする部分は同様のコードなので、差がでるとしたらGPUへの転送や動的に計算グラフを作る部分でしょうか。今回作ったSequentialクラスのやり方が、もしかしたらTensorFlowにとって扱いづらいものになっている可能性はあります。

と思って以下のモデルで追試をしました。結果は相変わらず遅かった…546秒。

class Autoencoder_eager(tfe.Network):
    def __init__(self):
        super(Autoencoder_eager, self).__init__()
        
        self.e_d1 = self.track_layer(
            tf.layers.Dense(128, activation=tf.nn.relu))
        self.e_d2 = self.track_layer(
            tf.layers.Dense(64, activation=tf.nn.relu))
        self.e_d3 = self.track_layer(
            tf.layers.Dense(12, activation=tf.nn.relu))
        self.e_d4 = self.track_layer(
            tf.layers.Dense(2))

        self.d_d1 = self.track_layer(
            tf.layers.Dense(12, activation=tf.nn.relu))
        self.d_d2 = self.track_layer(
            tf.layers.Dense(64, activation=tf.nn.relu))
        self.d_d3 = self.track_layer(
            tf.layers.Dense(128, activation=tf.nn.relu))
        self.d_d4 = self.track_layer(
            tf.layers.Dense(28*28))
    
    def encoder(self, x):
        x = self.e_d1(x)
        x = self.e_d2(x)
        x = self.e_d3(x)
        x = self.e_d4(x)
        return x

    def decoder(self, x):
        x = self.d_d1(x)
        x = self.d_d2(x)
        x = self.d_d3(x)
        x = self.d_d4(x)
        return x
    
    def call(self, x):
        z = self.encoder(x)
        y = self.decoder(z)
        return y

 

となると、完全に学習ループの方に問題があるということになります（それか本当にTensorFlowが遅いのか？）。考えられるのはnumpyからtf.convert_to_tensorを使ってtf.Tensorへの変換をミニバッチ毎に行っている点です。githubのexampleでは予めすべてのデータをtf.dataの形式で（中身となるデータはtf.Tensorとなる）保持しているようです。

tf.data.Datasetを利用した学習

tf.Datasetの作り方は以下。教師データと入力データをまとめる場合、それぞれのDatasetを作ってからzipします（今回は、教師データの方は不要です）。訓練ループの際にtfe.Iteratorで取り出すことができます。一々numpyからtf.Tensorへの変換はなくなりましたが、まあ（やっぱり）それほど速度アップには繋がりませんでした。

X_train = tf.cast(tf.convert_to_tensor(X_train), tf.float32)
y_train = tf.cast(tf.convert_to_tensor(y_train), tf.float32)

ds_X_train = tf.data.Dataset.from_tensor_slices(X_train)
ds_y_train = tf.data.Dataset.from_tensor_slices(y_train)
ds_train = tf.data.Dataset.zip((ds_X_train, ds_y_train))
ds_train = ds_train.shuffle(X_train.shape[0]).batch(128)


loss_log = []
global_step = tf.train.get_or_create_global_step()
start = time.time()

with tf.device('/gpu:0'):
    for i in range(100):
        for batch_x, _ in tfe.Iterator(ds_train):                
            loss_value, grads_and_vars = value_and_grad_f_w(batch_x, model)
            loss_log.append(loss_value.numpy())
            optimizer.apply_gradients(grads_and_vars)
            print('loss : {}'.format(loss_value.numpy()))
    end = time.time()
    learning_time = end - start

普通のTensorFlowの学習

静的グラフを作ってfeed_dictする前々からやられている方法。（けどいまはwhile_loopのグラフを作ってTensorFlowの世界で学習ループを回すのが主流なのかな？）まずはグラフ作成。クラスでの実装は特に必要ないですが、eagerのやつコピペしたのであしからず…。

TensorFlowはPythonで関数を呼び出した時には計算などはしていません（sess.runするまでは計算しない！）。ですから下記のコードは計算グラフを定義するコードであって、計算自体は何もしていません。

このとき、encoderの出力（埋め込み）が見たいので、encoderの計算グラフを定義しておきましょう。AutoEncoderで一括でやっちゃうと取り出せなくなります。

sess = tf.InteractiveSession()

class Sequential(object):
    def __init__(self, *args):        
        self.layer_list = [layer \
                           for layer in args]
        
    def build(self, x):
        for layer in self.layer_list:
            x = layer(x)
        return x

    
class Autoencoder(object):
    def __init__(self):
        self.encoder = Sequential(
            tf.layers.Dense(128, activation=tf.nn.relu),
            tf.layers.Dense(64, activation=tf.nn.relu),
            tf.layers.Dense(12, activation=tf.nn.relu),
            tf.layers.Dense(2))
        
        self.decoder = Sequential(
            tf.layers.Dense(12, activation=tf.nn.relu),
            tf.layers.Dense(64, activation=tf.nn.relu),
            tf.layers.Dense(128, activation=tf.nn.relu),
            tf.layers.Dense(28*28))
        
    def encode(self, x):
        z = self.encoder.build(x)
        return z
    
    def decode(self, z):
        y = self.decoder.build(z)
        return y
    
def loss(x_decoded, x_in):
    return tf.reduce_mean(tf.square(x_decoded - x_in))
    

model = Autoencoder()

with tf.name_scope('input'):
    x = tf.placeholder(dtype=tf.float32, shape=[None, 28*28])

with tf.name_scope('AutoEncoder'):
    encoder = model.encode(x)
    decoder = model.decode(encoder)

mseloss = loss(decoder, x)
optimizer = tf.train.AdamOptimizer(1e-3).minimize(mseloss)

init = tf.global_variables_initializer()
sess.run(init)

学習ループで初めて計算を実行します。一度やりたいことを計算グラフで全て定義し終えたら、投げるだけなので楽っちゃ楽。今回は簡単だけど、普通はもっとたくさん計算グラフを書かなければならないので、多分デバッグ大変だと思われます。ただ今は各種metricsもそろっていますし、層の定義も初期化などもかなり簡単になっていました。

loss_list = []
start_graph = time.time()
for i in range(100):
    ## データをシャッフル
    X, y = shuffle(X_train, y_train)

    for idx in range(0, num_data, batch_size):
        ## データをバッチサイズ分取り出す
        batch_x = X[idx: idx + batch_size 
            if idx + batch_size < num_data else num_data]
               
        sess.run(optimizer, feed_dict={x: batch_x})
        loss_value = mseloss.eval(session=sess, feed_dict={x: batch_x})
        
        # logging
        loss_list.append(loss_value)
        print('loss : {}'.format(loss_value))
end_graph = time.time()
learning_time = end_graph - start_graph

結局今回は、この実装が一番早かったなぁ…。

おまけ

主成分分析で28×28から2次元へ線形変換で埋め込んだ場合は以下のようになりました。線形変換にしてはそこそこというべきでしょうか、ある程度同じ数字同士は固まっています（が分離は甘い）。

以下はTF-Graphでの学習結果（学習データの埋め込み）

TensorFlow、0爆発させがち（もちろん偶然です…ｗ）