Tensorflow を強化学習でversion2.1を使うとversion1.12よりも遅くなるので原因を確認しました。

結論から述べますと、version2.1は、バッチで学習させるときにはversion1.12よりも若干速いのですが、データを1つづつに分けて学習させる場合ではかなり遅くなるということです。今回のテストだと約6倍も時間がかかりました。

普通の教師あり学習で、後者の方法を使うことはないと思いますが、強化学習では使う場合があります。

modelをsequential で作るか、functional で作るかの違いでも試してみましたが、そこはあまり違いはありませんでした。

以下、私のPCで測定した結果です。batchが100個のデータをバッチで学習させたとき、eachが1つずつ100回学習させたときです。

python 3.7, Tensorflow 2.1
model = sequential, mode = batch: time 0.37 sec
model = sequential, mode = each: time 3.63 sec
model = functional, mode = batch: time 0.35 sec
model = functional, mode = each: time 3.81 sec

python 3.6, Tensorflow 1.12
model = sequential, mode = batch: time 0.48 sec
model = sequential, mode = each: time 0.59 sec
model = functional, mode = batch: time 0.51 sec
model = functional, mode = each: time 0.63 sec

以下、プログラムです。

import numpy as np 
import tensorflow as tf 
import time
import pdb


class NeuralNet:
    def __init__(self,
        n_output=2,
        n_dense=16,
        n_dense2=16,
        input_size=(10,)
    ):
        self.n_output = n_output
        self.n_dense = n_dense
        self.n_dense2 = n_dense2
        self.input_size = input_size

    def built(self, mode='sequential'):
        if mode == 'sequential':
            self.model = tf.keras.models.Sequential([
                tf.keras.layers.Flatten(input_shape=self.input_size),
                tf.keras.layers.Dense(self.n_dense, activation='relu'),
                tf.keras.layers.Dense(self.n_dense2, activation='relu'),
                tf.keras.layers.Dense(self.n_output, activation='linear'),
            ])
        elif mode == 'functional':
            inputs = tf.keras.Input(shape=(self.input_size))
            x = tf.keras.layers.Flatten()(inputs) # 2次元の入力にも対応できるように入れている
            x = tf.keras.layers.Dense(self.n_dense, activation='relu')(x)
            x = tf.keras.layers.Dense(self.n_dense2, activation='relu')(x)
            outputs = tf.keras.layers.Dense(self.n_output, activation='linear')(x)
            self.model = tf.keras.Model(inputs=inputs, outputs=outputs)
        else:
            raise ValueError('modeが間違っています')
        self.model.compile(
            optimizer='adam',
            loss='mean_squared_error',
            metrics=['mse']
        )
    
    def learn(self, X, Y, mode='batch'):
        if mode == 'batch':
            # 全てまとめて
            self.model.fit(X, Y, verbose=0, epochs=1)
        elif mode == 'each':
            # 一つずつ
            for xx, yy in zip(X, Y):
                self.model.fit(
                    xx.reshape((1,) + self.input_size), 
                    yy.reshape((1, 2)), 
                    verbose=0, epochs=1)

    def predict(self, X):
        Y = self.model.predict(X)
        return Y
    

if __name__ == '__main__':
    # データ
    N = 100 # データ数
    D = 10  # input の次元
    M = 2   # output の次元
    X = np.random.randn(N, D)
    W = np.random.randn(D, M)
    Y = X.dot(W)

    print('TF version', tf.__version__)
    for model in ['sequential', 'functional']:
        for mode in ['batch', 'each']:
            nn = NeuralNet(n_output=M, input_size=(D,))
            nn.built(mode=model)
            stime = time.time()
            nn.learn(X, Y, mode=mode)
            etime =time.time()
            print('model = %s, mode = %s: time %.2f sec' % (model, mode, etime - stime))