javascriptでchainerモデルを利用するためのKeras.js

  • 32
    Like
  • 0
    Comment

 この記事では、chainerの学習済みモデルを、Keras.jsを使ってフロント側で利用する方法を紹介します。

 学習済みモデルを用いたサービスを作るためには、サーバー側でモデルを利用するAPIを用意するか、フロント側で直接モデルを利用するコードを書く必要があります。 フロント側で学習済みモデルを使う大きな利点は計算機コストがかからない点です。ファイルサーバだけ用意すれば良いので、小さなデモを公開する際は便利です。逆に欠点としてjavascriptで再実装する必要があります。しかしKeras.jsを使えばネットワーク構造をJSで書き直す必要がなくなるため、非常に簡単になります。

 ちなみにGirl Friend Factory - 機械学習で彼女を創る -では、この方法を用いたデモを作成しました。

記事の流れ

  1. KerasとKeras.js
  2. chainerモデルからKerasモデルへ
  3. KerasモデルからKeras.jsモデルへ
  4. javascriptで順伝播
  5. (発展)学習済みモデルの容量を小さくする試行錯誤
  6. 感想

1. Keras.jsとは

 Keras.jsは、Kerasモデルをブラウザで、GPUを使って動かすライブラリです。ちなみにKerasは、TensorFlow・Theanoなどのラッパーライブラリです。
 chainerモデルをブラウザで使う方法は他にも、ConvNetJSを使う手がありますが、こちらはjavascriptで同等のネットワーク構造を定義する必要があったり、利用したいNN層が無かったりと、少し使いづらいです。

 chainerモデルをKeras.jsで利用するには、まずchainerモデルをKerasモデルに変換し、KerasモデルをKeras.jsモデルに変換する必要があります。後者は専用のコードが用意されていますが、前者は自分で実装する必要があります。

2. chainerモデルからKerasモデルへ

 今回は試しに、DeconvolutionしてBatchNormalizationしてConvolutionするだけのモデルを変換します。

import chainer

batchsize = 1
size = 2
n_in = 3
n_middle = 2
n_out = 1
ksize = 2


class ChainerModel(chainer.Chain):
    def __init__(self):
        super().__init__(
            deconv=chainer.links.Deconvolution2D(n_in, n_middle, ksize=ksize, stride=ksize),
            bn=chainer.links.BatchNormalization(n_middle),
            conv=chainer.links.Convolution2D(n_middle, n_out, ksize=ksize, stride=ksize),
        )

    def __call__(self, x, test):
        return self.conv(self.bn(self.deconv(x), test=test))


chainer_model = ChainerModel()

 モデルを変換するライブラリがあるだろうと探しましたが、見当たらなかったので自分で実装します。

 まず、等価なKerasモデルを作ります。コツとしまして、畳み込み層の次元の順序dim_orderingは、TensorFlowモードtfを利用した方が良いです。(TheanoモードthだとKeras.jsでエラーが発生したため。)

 ConvolutionやDeconvolutionは、使うライブラリによって初期値や引数の仕様が異なるため、同じ形のネットワークを作るのに結構時間がかかります。

 Kerasでネットワークを定義する方法は複数ありますが、層をaddしていくよりも、順伝搬を自分で書くほうがchainerを使う人には直感的だと思います。

import keras

input_keras = keras.layers.Input(shape=(size, size, n_in), name='input')

h = input_keras
h = keras.layers.convolutional.Deconvolution2D(
    n_middle, ksize, ksize, name='deconv',
    output_shape=(batchsize, size * ksize, size * ksize, n_middle),
    dim_ordering='tf', subsample=(ksize, ksize),
)(h)
h = keras.layers.normalization.BatchNormalization(
    epsilon=2e-5, momentum=0.9, axis=3, name='bn',
)(h)
h = keras.layers.convolutional.Convolution2D(
    n_out, ksize, ksize, name='conv',
    dim_ordering='tf', subsample=(ksize, ksize),
)(h)

keras_model = keras.models.Model(input_keras, h)

 一旦このまま適当なベクトルを順伝搬させてみます。chainerのBatchNormalizationは何か学習させないとtest=Trueで順伝搬できないので、最初に適当な学習をさせています。

import numpy
x = numpy.random.randn(batchsize, n_in, size, size).astype(numpy.float32)

# 適当に学習させる
optimize = chainer.optimizers.Adam()
optimize.setup(chainer_model)
y = numpy.zeros(shape=(batchsize, n_out, size, size), dtype=numpy.float32)
for _ in range(10):
    optimize.zero_grads()
    loss = chainer.functions.mean_squared_error(chainer_model(x, test=False), y)
    loss.backward()
    optimize.update()

# chainer 順伝搬
output_chainer = chainer_model(x, test=True).data

# Keras 順伝搬
# `(samples, rows, cols, channels)` if dim_ordering='tf'
output_keras = keras_model.predict(x.transpose(0, 2, 3, 1)).transpose(0, 3, 1, 2)

# 予測の差
print('diff: ', numpy.abs(output_chainer - output_keras).sum())  # diff: 3.54881

 重みをコピーしていないので予測値に大きな差があります。では最後に、重みをコピーしてみます。(コードは適当にこのGistに上げましたので、他のレイヤーも実装した際はぜひ追加してください。)

def copy_weights_deconvolution(chainer_model, keras_model, layer_name):
    deconv_chainer = chainer_model[layer_name]
    W, b = (deconv_chainer.W.data, deconv_chainer.b.data)
    keras_model.get_layer(layer_name).set_weights([numpy.transpose(W, (2, 3, 0, 1)), b])


def copy_weights_convolution(chainer_model, keras_model, layer_name):
    conv_chainer = chainer_model[layer_name]
    W, b = (conv_chainer.W.data, conv_chainer.b.data)
    keras_model.get_layer(layer_name).set_weights([numpy.transpose(W, (2, 3, 1, 0)), b])


def copy_weights_bn(chainer_model, keras_model, layer_name):
    bn_chainer = chainer_model[layer_name]
    w = [bn_chainer.gamma.data, bn_chainer.beta.data, bn_chainer.avg_mean, bn_chainer.avg_var]
    keras_model.get_layer(layer_name).set_weights(w)

copy_weights_deconvolution(chainer_model, keras_model, 'deconv')
copy_weights_convolution(chainer_model, keras_model, 'conv')
copy_weights_bn(chainer_model, keras_model, 'bn')

output_keras = keras_model.predict(x.transpose(0, 2, 3, 1)).transpose(0, 3, 1, 2)
print('diff: ', numpy.abs(output_chainer - output_keras).sum())  # diff: 0.000445604

 差が3.54881から0.00044まで縮まりました。なんかまだ微妙に誤差がありますが、まぁ・・・

3. KerasモデルからKeras.jsモデルへ

 ここからはとても簡単で、Kerasモデルを保存したあと、Keras.jsが用意してくれているコードを実行すれば必要なファイルが生成されます。

# save Keras model
with open('keras_model_arch.json', 'w') as f:
    f.write(keras_model.to_json())

keras_model.save('keras_model.h5')
# https://github.com/transcranial/keras-js/blob/master/encoder.py
python3 encoder.py keras_model.h5

 これで

  • keras_model.h5
  • keras_model_arch.json
  • keras_model_weights.buf
  • keras_model_metadata.json

の4ファイルができたと思います。このうち下3つのファイルをKeras.jsで利用します。

4. javascriptで順伝播

 あとはKeras.jsの使い方に従って使うだけです。

const model = new KerasJS.Model({
  filepaths: {
    model: 'keras_model_arch.json',
    weights: 'keras_model_metadata.json',
    metadata: 'keras_model_weights.buf'
  },
  gpu: true
})

try {
  await model.ready()
  const inputData = {
    'input': new Float32Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1])
  }
  const outputData = await model.predict(inputData)
} catch (err) {
  // handle error
}

 あとはoutputDataに予測値が入っているので、canvasで画像を表示したりできます。

5. (発展)学習済みモデルの容量を小さくする試行錯誤

 最初に紹介したデモはDCGANで学習済みの画像生成モデルを利用しており、そのbufferファイルが70MBを超えました。画像の生成にfloat32の精度は不要なので、float16にしてモデルの軽量化を試みました。

 Kerasモデルの方はencoder.pyを書き換えることでnumpy.float16に対応させることは容易でしたが、どうやらjavascriptの方はFloat32Arrayはあるもののfloat16(半精度浮動小数点数)に対応するものが無いらしく、断念しました。

 もし頑張るなら、半精度浮動小数点数のバイナリから実数値を得るコードを使うようにKeras.jsを修正する必要がありそうです。

6. 感想

Keras.jsとConvNetJSどちらが楽か

 ConvNetJSの方がえらく昔の物なのと、ネットワーク構造をjavascriptで再定義するのがめんどくさそうなので、まだ同じPythonを使えばなんとかなるKeras経由のKeras.js利用の方が楽だと思います。

フロント側とサーバー側どちらで順伝搬させるのが良いか

 断然サーバーを利用することをおすすめします。まず第一に、結果が出るまでの時間が圧倒的に短いです。サーバーなら画像1枚生成につき0.5秒ほどでしたが、Keras.js利用ならダウンロードに数十秒、画像1枚に1秒かかりました。Kerasで等価なモデルを再定義するのはあまり意味が無い上、意外と実装に時間がかかります。また、ネットワーク構造が変わるたびに作り変える必要があり、より一層つらいです。
 もちろん計算資源コストの問題や、サーバー利用関連のめんどくささがあるので、どうしても、どうしてもフロント側で順伝搬させる必要があるなら、Keras.jsを使うと良さそうです。