カルマンフィルタのコード比較【numpy, pytorch, eager】

はじめに

最近わたしはTensorFlowにボコボコにされています。 計算グラフを書いて、sess.runするまでは何が起こるかわからない！何かが起こっても何が起こっているかは分からない！ そんな状態から、TensorFlow eager executionの登場で解放された！……かに思えました。

今回は下記のTensorFlowとnumpyによるカルマンフィルタの実装例に見て、 これがTensorFlow eager executionでどれくらい楽になるかを試してみました（やる前からそれほど楽ではないことは分かっていた）。

カルマンフィルタといえば理論は非常に難解ですが、実装は意外と楽チンで、実用上非常に優れた予測モデル（と言っていいかは微妙ですが）だと言えます。 今回は理論には立ち入らず、numpyの実装、TensorFlow eagerの実装、そしてpytorchの実装を見比べてみましょう。

ちなみに、下記の記事のコードを全面的に参考にしています。 実際はこれを各ライブラリのAPIで置き換えただけです。

qiita.com

コード比較

numpy

def KalmanFilter(T, x_obs, mu0, V0, A, b, Q, R):
    """
    numpy実装．
    """

    mu = mu0
    V = V0

    ## Tはシミュレーションのタイムステップ数
    ## メモリのアロケーションを行い、後で結果を逐次代入していく
    x_predict = np.zeros((T, 2))
    x_filter = np.zeros((T, 2))
    V_predict = np.zeros((T, 2, 2))
    V_filter = np.zeros((T, 2, 2))

    ## 初期値
    x_predict[0] = mu.transpose()
    x_filter[0] = mu
    V_predict[0] = V
    V_filter[0] = V


    start = time.time()

    ## シミュレーションのループ
    ## 初期値t=0は既に埋めてあるので、1:Tまで
    for t in range(1, T):

        mu_ = A @ mu + b
        V_ = A @ V @ A.transpose() + Q

        x_predict[t] = mu_
        V_predict[t] = V_

        S = V_ + R
        K = V_.dot(np.linalg.inv(S))

        mu = mu_ + K @ (x_obs[t] - mu_)
        V = V_ - K @ V_

        x_filter[t] = mu
        V_filter[t] = V

    elapsed_time = time.time() - start
    print("numpy: ", elapsed_time)

    return x_predict, V_predict, x_filter, V_filter

ある意味最も慣れ親しんだコードですね。 numpyでは行列計算を"@"で実行できます。

numpyのコードを基準に見ていきましょう。

pytorch

def KalmanFilter_th(T, x_obs, mu0, V0, A, b, Q, R):
    """
    pytorch実装
    """

    ## numpyの型をtorch.Tensorの型へ変換
    ## numpyがfloat64なのでこちらもdoubleで
    mu = torch.from_numpy(mu0).double()
    V = torch.from_numpy(V0).double()
    A = torch.from_numpy(A).double()
    b = torch.from_numpy(b).double()
    Q = torch.from_numpy(Q).double()
    R = torch.from_numpy(R).double()
    x_obs = torch.from_numpy(x_obs).double()

    ## メモリのアロケーション
    x_predict = torch.zeros((T, 2))
    x_filter = torch.zeros((T, 2))
    V_predict = torch.zeros((T, 2, 2))
    V_filter = torch.zeros((T, 2, 2))

    ## 初期値
    x_predict[0] = mu
    x_filter[0] = mu
    V_predict[0] = V
    V_filter[0] = V


    start = time.time()

    for t in range(1, T):

        ## pytorchでは転置は "Tensor.t()" で実行可能
        ## ちょっとだけ楽？
        mu_ = A @ mu + b
        V_ = A @ V @ A.t() + Q

        x_predict[t] = mu_
        V_predict[t] = V_

        ## torch.dotはブロードキャストが無い（ベクトル同士などに使う）
        ## 行列をベクトルに作用させたりする場合はtorch.mm
        S = V_ + R
        K = V_.mm(torch.inverse(S))

        mu = mu_ + K @ (x_obs[t] - mu_)
        V = V_ - K @ V_

        x_filter[t] = mu
        V_filter[t] = V

    elapsed_time = time.time() - start
    print("torch: ", elapsed_time)

    return x_predict, V_predict, x_filter, V_filter

基本的にnumpyとほとんど変わりませんでした。

メソッドの差はあれど、ちょこっと調べるだけですぐに書き換えられるようです。 cuda()メソッドで直ちにgpuへ転送できることも考えれば、numpyに代わる数値計算ライブラリになってもおかしくありませんね。

TensorFlow eager execution

def KalmanFilter_tfe(T, x_obs, mu0, V0, A, b, Q, R):
    """
    eager実装
    """
    
    ## numpyをtf.Tensorへ変換。
    ## 型はtf.double (=tf.float64)
    mu = tf.convert_to_tensor(mu0, dtype=tf.double)
    V = tf.convert_to_tensor(V0, dtype=tf.double)
    A = tf.convert_to_tensor(A, dtype=tf.double)
    b = tf.convert_to_tensor(b, dtype=tf.double)
    Q = tf.convert_to_tensor(Q, dtype=tf.double)
    R = tf.convert_to_tensor(R, dtype=tf.double)
    x_obs = tf.convert_to_tensor(x_obs, dtype=tf.double)


    ## 実はメモリのアロケーションを行う方法は
    ## eager では利用できません。
#     x_predict = tf.zeros((T, 2))
#     x_filter = tf.zeros((T, 2))
#     V_predict = tf.zeros((T, 2, 2))
#     V_filter = tf.zeros((T, 2, 2))

    ## 代わりに結果を格納するリストを使います。
    x_predict = []
    x_filter = []
    V_predict = []
    V_filter = []

    ## メモリのアロケーションをして準備した配列(tf.Tensor)に
    ## 値を代入することが今のところ許されていません。
    ## ここでエラーを吐きます。
#     x_predict[0] = mu
#     x_filter[0] = mu
#     V_predict[0] = V
#     V_filter[0] = V

    ## 代わりに結果を逐次リストで追加していく
    x_predict.append(mu)
    x_filter.append(mu)
    V_predict.append(V)
    V_filter.append(V)
    
    start = time.time()

    for t in range(1, T):

        ## 行列とベクトルの演算はtf.tensordotで実行axesには注意
        ## 機械学習ではバッチ処理が基本なので
        ## そもそも１次元Tensorを扱うことは少ないためこのような実装に？
        mu_ = tf.tensordot(A, mu, axes=1) + b
        V_ = A @ V @ tf.transpose(A) + Q

#         x_predict[t] = mu_
#         V_predict[t] = V_
        
        ## リストで追加
        x_predict.append(mu_)
        V_predict.append(V_)

        ## tensorflow はガシガシ関数で囲んでやる
        ## tfe.Tensor自体にメソッドで備わったりはしていない。
        S = V_ + R
        K = tf.matmul(V_, tf.matrix_inverse(S))

        mu = mu_ + tf.tensordot(K, (x_obs[t] - mu_), axes=1)
        V = V_ - K @ V_

#         x_filter[t] = mu
#         V_filter[t] = V
        x_filter.append(mu)
        V_filter.append(V)

    elapsed_time = time.time() - start
    print("eager: ", elapsed_time)

#     return x_predict, V_predict, x_filter, V_filter
    ## 最後にリストをtf.Tensorに組みなおして終わり
    return tf.stack(x_predict), tf.stack(V_predict), tf.stack(x_filter), tf.stack(V_filter)

さて私が大いに期待を寄せているeagerの登場です。実は既にサンプリングの勉強をしていた時に、TensorFlow eagerでは実装がやりにくそうだという感覚を既に抱いていました。

カルマンフィルタという、比較的単純なコードで書けるアルゴリズムならば、その面倒な部分が比較しやすいのではと思い記事にしたというのが本当の動機です。 既にプログラムのコメントにぶつかりうる困難を書き記しました。 もっと膨大なプログラムになればもっと面倒なことが起こりそうです。

今回は各タイムステップの結果を格納していくだけだったので、List.append()で何とかなりました。 これがもっと入り組んだ値の更新を必要とする場合、たとえば３次元Tensorの特定の２次元Tensorのみを更新したいとか、 そのようなときにどうすれば良いのか私には今のところ分かりません（そもそもeagerを使う強い動機が見つかりませんし）。

ここらへんがもっとフレキシブルに動けば、eagerモードが数値計算ライブラリとして用いられることもあり得ると思います。

速度と結果の比較

まずCPU飲みを使った計算を比較します。 numpyはcpuのみしか使うことができないので、これがフェアな比較です。

# データ作成
T = 20
mu0 = np.array([100, 100])
V0 = np.array([[10, 0],
               [0, 10]])

A = np.array([[1.001, 0.001],
              [0, 0.99]])
b = np.array([5, 10])
Q = np.array([[20, 0],
              [0, 20]])
R = np.array([[20, 0],
              [0, 20]])

rvq = np.random.multivariate_normal(np.zeros(2), Q, T)
rvr = np.random.multivariate_normal(np.zeros(2), R, T)
obs = np.zeros((T, 2))
obs[0] = mu0
for i in range(1, T):
    obs[i] = A @ obs[i-1] + b + rvq[i] + rvr[i]
# 作成終わり



x_predict, V_predict, x_filter, V_filter = \
    KalmanFilter(T, obs, mu0, V0, A, b, Q, R)
print(x_filter)

x_predict_th, V_predict_th, x_filter_th, V_filter_th = \
    KalmanFilter_th(T, obs, mu0, V0, A, b, Q, R)
print(x_filter_th)

x_predict_tfe, V_predict_tfe, x_filter_tfe, V_filter_tfe = \
    KalmanFilter_tfe(T, obs, mu0, V0, A, b, Q, R)
print(x_filter_tfe)

fig = plt.figure(figsize=(16, 9))
ax = fig.gca()

ax.scatter(obs[:, 0], obs[:, 1], s=10, alpha=1, marker="o", color='w', edgecolor='k')
ax.plot(obs[:, 0], obs[:, 1], alpha=0.5, lw=1, color='k')

ax.scatter(x_filter[:, 0], x_filter[:, 1], s=10, alpha=1, marker="o", color='r')
ax.plot(x_filter[:, 0], x_filter[:, 1], alpha=0.5, lw=1, color='r')

ax.scatter(x_filter_tfe.numpy()[:, 0], x_filter_tfe.numpy()[:, 1], s=10, alpha=1, marker="o", color='g')
ax.plot(x_filter_tfe.numpy()[:, 0], x_filter_tfe.numpy()[:, 1], alpha=0.5, lw=1, color='g')

ax.scatter(x_filter_th.numpy()[:, 0], x_filter_th.numpy()[:, 1], s=10, alpha=1, marker="o", color='m')
ax.plot(x_filter_th.numpy()[:, 0], x_filter_th.numpy()[:, 1], alpha=0.5, lw=1, color='m')

こちらのコードは冒頭でも紹介したコードそのままです (tensorflowでカルマンフィルタ - Qiita)

結果

numpy : 0.0015537738800048828

pytorch : 0.0021310848236083984

eager : 0.03675341606140137

おえーー、eager遅い！！！なんでや！！？

灰色が観測データ、紫色がカルマンフィルタによって推定された結果（どの結果も同じになったので重なっております）。

計算がちゃんと出来ていることは確認できましたが、やはり遅すぎないでしょうかeagerさん…。 ということで、一旦、numpyもtorchもリストにぶち込んでいく方針で実装してみます。

第二回戦

def KalmanFilter(T, x_obs, mu0, V0, A, b, Q, R):
    """
    numpy実装．
    """

    mu = mu0
    V = V0

    ## Tはシミュレーションのタイムステップ数
    ## メモリのアロケーションを行い、後で結果を逐次代入していく
#     x_predict = np.zeros((T, 2))
#     x_filter = np.zeros((T, 2))
#     V_predict = np.zeros((T, 2, 2))
#     V_filter = np.zeros((T, 2, 2))

    x_predict = []
    x_filter = []
    V_predict = []
    V_filter = []

#     x_predict[0] = mu.transpose()
#     x_filter[0] = mu
#     V_predict[0] = V
#     V_filter[0] = V

    x_predict.append(mu.transpose())
    x_filter.append(mu)
    V_predict.append(V)
    V_filter.append(V)


    start = time.time()

    ## シミュレーションのループ
    ## 初期値t=0は既に埋めてあるので、1:Tまで
    for t in range(1, T):

        mu_ = A @ mu + b
        V_ = A @ V @ A.transpose() + Q

#         x_predict[t] = mu_
#         V_predict[t] = V_
        
        x_predict.append(mu_)
        V_predict.append(V_)

        S = V_ + R
        K = V_.dot(np.linalg.inv(S))

        mu = mu_ + K @ (x_obs[t] - mu_)
        V = V_ - K @ V_

#         x_filter[t] = mu
#         V_filter[t] = V

        x_filter.append(mu)
        V_filter.append(V)

    elapsed_time = time.time() - start
    print("numpy: ", elapsed_time)

    return np.stack(x_predict), np.stack(V_predict), np.stack(x_filter), np.stack(V_filter)

numpyに関しては上記のように変更。ただただリストに突っ込んでいく形で実装しただけです。同様にしてpytorchも以下のように実装。

def KalmanFilter_th(T, x_obs, mu0, V0, A, b, Q, R):
    """
    pytorch実装
    """
    mu = torch.from_numpy(mu0).double()
    V = torch.from_numpy(V0).double()
    A = torch.from_numpy(A).double()
    b = torch.from_numpy(b).double()
    Q = torch.from_numpy(Q).double()
    R = torch.from_numpy(R).double()
    x_obs = torch.from_numpy(x_obs).double()

#     x_predict = torch.zeros((T, 2))
#     x_filter = torch.zeros((T, 2))
#     V_predict = torch.zeros((T, 2, 2))
#     V_filter = torch.zeros((T, 2, 2))

    x_predict = []
    x_filter = []
    V_predict = []
    V_filter = []

#     x_predict[0] = mu
#     x_filter[0] = mu
#     V_predict[0] = V
#     V_filter[0] = V

    x_predict.append(mu)
    x_filter.append(mu)
    V_predict.append(V)
    V_filter.append(V)

    start = time.time()

    for t in range(1, T):


        mu_ = A @ mu + b
        V_ = A @ V @ A.t() + Q

#         x_predict[t] = mu_
#         V_predict[t] = V_

        x_predict.append(mu_)
        V_predict.append(V_)

        S = V_ + R
        K = V_.mm(torch.inverse(S))

        mu = mu_ + K @ (x_obs[t] - mu_)
        V = V_ - K @ V_

#         x_filter[t] = mu
#         V_filter[t] = V

        x_filter.append(mu)
        V_filter.append(V)

    elapsed_time = time.time() - start
    print("torch: ", elapsed_time)

    return torch.stack(x_predict), torch.stack(V_predict), torch.stack(x_filter), torch.stack(V_filter)

これで、純粋に計算処理の速さ勝負以外はフェアになったかなと思います。

結果

numpy: 0.0015443035125732422

torch: 0.001976583480834961

eager: 0.03855392265319824

チーン…。 きっと、サイズの小さな計算に対してはTensorFlowは元々強くないのでしょう（適当）。 numpyとpytorchは、まあnumpyの方が速いのかもしれませんしたまたまかもしれませんが、 eagerはちょっと流石に遅すぎるような…。時間図っているの行列計算のところだけだし…。

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com