kotlinで簡単なテンソル処理がしたい

テンソル処理用にライブラリ（といってもプロジェクトは分かれてないが）を作ったら割と良いものが出来たのでここにページを作っておく。 TF Lite使う人の決定版が作れたんじゃないか？と自分では思っている。自分しか使わないだろうが。

最終的には以下のような使い方をしている。

tensor_ns {
    val len = newStroke.shape[0]
    unnormalize[n(curIndex), r(0, len), r(0, 2)] = newStroke
    unnormalize[n(curIndex), r(0, len), n(2)] = INPUT_TYPE_POS

    val nonzeroMask = unnormalize[all, all, n(2)].scalar_equal(INPUT_TYPE_POS)
    val nonzero = unnormalize[nonzeroMask]
    val xmax = nonzero[all, n(0)].max()
    val xmin = nonzero[all, n(0)].min()
    val ymax = nonzero[all, n(1)].max()
    val ymin = nonzero[all, n(1)].min()
    val xdelta = xmax - xmin + 0.0001f
    val ydelta = ymax - ymin + 0.0001f
    val scale = min(NORMALIZE_MAX.toFloat() / xdelta, NORMALIZE_MAX.toFloat() / ydelta)


    repeat(curIndex + 1) {
        val rowMask = unnormalize[n(it), all, n(2)].scalar_equal(INPUT_TYPE_POS)
        val rowXY = unnormalize[n(it), all, all][rowMask]
        val rowLen = rowXY.shape[0]
        val originTensorX = rowXY[all, n(0)] - xmin
        val originTensorY = rowXY[all, n(1)] - ymin
        outputTensor[n(it), r(0, rowLen), n(0)] = originTensorX * scale
        outputTensor[n(it), r(0, rowLen), n(1)] = originTensorY * scale
    }
    outputTensor[n(curIndex), r(0, len), n(2)] = INPUT_TYPE_POS
}

これは実際に実機でやってるノーマライズの処理。この位の事をやるライブラリが欲しいな、というのがモチベーション。

上記コードの元になった実際の使用例
KdFTensorのUnitTest type safe builder使ってないがだいたいの雰囲気は分かると思う
今の所3ファイルで400行くらい。こういうの簡単に作れてkotlinは(・∀・)ｲｲﾈ!!

numpyとの比較

numpyとステップbyステップで比較してデバッグするのを目的としているので、numpyと同じような振る舞いになるのを重視している。以下はtensor_ns type safe builderの中での記述。

インデックス

indexは三種類

数字単体： n(1)とかn(3)とかnという関数で表す
範囲： r(0, 5)とかrという関数で表す
全体： allで表す

例えばnumpyで、

# numpy
a[3, 5, 2]

と書く所は、

// kotlin
a[n(3), n(5), n(2)]

と書く。以下のようなの

# numpy
a[:, :, 2]

は、

// kotlin
a[all, all, n(2)]

と書く。

範囲はkotlinの常識とは少し違って、numpy同様endがexclusive。つまり、以下のようなnumpyのコードは、

# numpy
a[3:5]

以下のように書くが、この5は含まない。

// kotlin
a[r(3, 5)]

sub tensorの取り出し

次元の指定で、数値が指定されるとその次元はsequeezeされる。

    val target = arange(30).reshape(2, 3, 5)
    val sub = target[n(1), r(0, 2), r(1, 4)]

    assertEquals(Shape(2, 3), sub.shape)

n(1)で指定された次元は落ちている。これはnumpyと同じ挙動。ただし、最終的にスカラーになるケースではnumpyはスカラーになるけれど、こちらは1次元の要素一つのテンソルのままとした。

型が変わるのはkotlin的には気持ち悪いだろうから。

booleanのマスク

booleanのtensorをインデックスに指定する事も出来る。こちらもだいたいnumpyと同じshapeになる。

booleanのテンソルはnumpyだと以下のように==で作るが、

# numpy
mask = (target == 1)

kotlinだと==はboolean以外返すのは一般的では無いので、scalr_equalというのを用意した。

// kotlin
mask = (target.sclar_equal(1f))

maskのshapeはtargetと同じ。

さらに、maskのshapeとtargetのshapeは先頭から一致をしていって、一致した所までのマスクとみなす。だからそれ以外の所はallのようにふるまう。

言葉にするとややこしいが、以下みたいなコードがちゃんと動く、という話。

    val data = listOf(0f, 1f,
        3f, 0f,
        5f, 1f,
        2f, 1f,
        10f, 0f)
    val target = tensor(data).reshape(5, 2)
    val mask = target[all, n(1)].scalar_equal(1f)

    val actual = target[mask]

    assertEquals(Shape(3, 2), actual.shape)

maskがn(1)で一次元落とされているので、maskは(5, )のテンソル。中身は[true, false, true, true, false]となってる。

これをtargetの(5, 2)のインデックスとして使うと、最初の5にだけマッチしてあとはallとしたテンソルが得られる。この場合はtrueが三つなので(3, 2)となっている。

この辺はnumpyと同じ挙動。

代入

sub tensorを取り出すのと同じインデクシングが使えて、スカラーはブロードキャストする。 shapeは先頭から突き合わせていって、次元が合わないでかつ1だったらそこの次元はブロードキャストとみなす。それで解決できない時はエラー。numpyと同じ挙動。

    val target = arange(10).reshape(5, 2)
    val input = tensor((35..40).map { it.toFloat() }.toList()).reshape(3, 2)

    target[r(1, 4), all] = input
    assertEquals(0f, target[0, 0])
    assertEquals(1f, target[0, 1])
    assertEquals(35f, target[1, 0])
    assertEquals(36f, target[1, 1])
    assertEquals(40f, target[3, 1])
    assertEquals(8f, target[4, 0])

モチベーション

このライブラリを作ったモチベーションも書いておく。

numpyでノーマライズの処理をしてモデルに食わせてトレーニングするので、実機の方でも同じような処理を書く必要がある。デバッグの為に、ステップ実行で止めて中のデータを別ウィンドウのブラウザのnumpyで求めたものと比較、とかしたい。また、このノーマライズの処理は結構変更される。

必要な物としては、

デバッガで止めてすぐ中が見れる（効率よりはデバッガでの見やすさ重視）
- データの保持は生の配列くらいにしておきたい
- genericsとかはなるべく使わず、DRYよりもデバッグ時の読みやすさ優先
記述が簡潔でnumpyと同じ処理が書きやすい
reshapeはやりまくる(下は一次元でdupしないで欲しい、-1のguessは頑張って欲しい)
数値型を意識したByteBufferとのやり取り（最後はTF Liteなので）
shape回りはかなりいろいろやるのでちゃんと作りたい

要らない機能としては、

あくまで前処理なので掛け算とかその他の行列計算は早く無くて良い（それはTF Liteでやるから）
使う機能だけで良い（むしろ要る物を気軽に足していけるように外部ライブラリじゃない方がうれしい）

全体的に、numpy上で書いた処理を同じ感じで持って来やすい、という記述上の簡易さが重要。表面上の記述というかテンソル処理のDSL的なのが欲しい、という事でシンタックスが重要という話。ただあまりkotlin的に予想外の振る舞いもしてほしくない。

制作過程

せっかくなので作る時に考えていた事なども残しておく。

やりたい事

やりたいのはストロークのノーマライズとか。そういうのを簡単に書きたいが、そのためにAny型をいろいろdynamicに処理するのは嫌だ。型のための多少の冗長性は受け入れよう。

ストロークはshape的には(strokenum, one_stroke_len, 3)で、最後の3はx, y, 0か1のトリプレット。最後の0か1は、座標がある所は1、無い所は0。

で、ノーマライズはトリプレットの最後が1の所だけのxの最小値をxから引く、とかそういうのが書きたい。

numpy的には、あってるかわからんが、こんな感じ。

mask = strok[:, :, 2] == 1
xmax = stroke[mask][0].max()

こういう感じの処理を書きたい。 maskのshapeは(strokenum, one_stroke_len)で、boolだった。

stroke[mask]のshapeは(trueの数, 3)となる。

と思ったが、モデルの都合でこのノーマライズは要らない事に気づいた。だいたい定数倍でいいか。だから

stroke[:, :, 0]*k
stroke[:, :, 1]*k

あたりが出来れば良いか。inplaceの方が本当は良いんだろうが、代入出来ればまぁ最初はいいかなぁ。

案1 assignerを作る（追記：こちらはボツ）

ようするに、インデックスのシーケンスが得られればいいんだよな。何かしらのtype safe builderがあって、

{all; all; n(0)}

とかでインデックスの列が返ってくるとかならどうだろう？ある範囲ならrangeでrとか？

{all; all; r(0, 10)}

とか。rはもうちょっと考えた方が良さそうだが。

こんな感じのDSLが定義出来れば、部分配列を取り出すのをsと呼べば、

stroke.s{all; all; r(0, 10)} * k

とかで

stroke[:, :, 0:10]*k

的な事が出来るとか？代入は変なトリック使うよりは、assignerみたいなのを返すとして、それがaとすると

stroke.a{all; all; r(0, 10)}(stroke.s{all; all; r(0, 10)}*k)

読みにくいが書きやすさは悪くも無いか。

案2: 全体をtype safe builderにする（追記：現状はこれになった）

indexの所だけをtype safe builderにするより、そもそもnumpy的操作全体をtype safe builderに押し込める方がいいのでは、と言われて、少し考えてみる。具体的にとりあえず今必要になってる操作だけを考えると、

fun onExec(xyposlist: List<List<Float>>) {
    val stroke = build() {
        val res = zeros(MAX_STROKE_NUM, MAX_ONE_STROKE_LEN, 3)
        xyposlist.forEachIndexed { i, xypos ->
            val xy = tensor(xypos).apply{ reshape(-1, 2) }

            val xmax = xy[all, n(0)].max()
            val ymax = xy[all, n(1)].max()
            val scale = NORMALIZE_MAX/max(xmax, ymax)

            val sp = xy.shape
            res.assign(n(i), r(0, sp[0]), r(0, 1))(xy*scale)
            res.assign(n(i), r(0, sp[0]), n(2))(tensor(1))
        }
        res
    }
    val classes = model.predict(stroke)
    Log.d(TAG, "classes = $classes")
}

こんな感じか。ブロードキャストするのにtensor(1)とか書くのはだるいが、こちらの方がスッキリ書けるし実装コストも同じくらいか。

assignが高階関数になるのは仕方ないかなぁ。

rは読みにくくて0..sp[0]と書きたい気もするが、半開区間みたいなのを実装したいのと、..ではケツが含まれてしまうので仕方ないかなぁ。 toとかの方がマシか？

res.assign(n(i), 0 to sp[0], 0 to 1)(xy*scale)
res.assign(n(i), 15 to end, 2)(tensor(0))

これならrの方がマシか。nと揃ってる分。これじゃendの型はなんだよ、という気がするし。

assignerをgetをオーバーライドしてsub assigner返すように作れないか？とkotlinのcallableを軽く調べたが、KCallableはメンバ多くて辛そう。これならメソッドを呼ぶ方がマシか。と思ったがそんな事しないでinvokeをオーバーライドするのか。なるほど。

res.assign(tensor(1))
res.assign[n(i), r(0, sp[0]), 1](xy*scale)

いいかも。

お、良く見たらkotlinはindexerのassignもoperator overloading出来るのか。知らなんだ。じゃあこうか。

res[n(i), r(0, sp[0]), r(0, 1)] = xy*scale
res[all, all, all] = tensor(1)

こっちの方が自然か。

書き直すと、こんな感じか？

    val stroke = build() {
        val res = zeros(MAX_STROKE_NUM, MAX_ONE_STROKE_LEN, 3)
        xyposlist.forEachIndexed { i, xypos ->
            val xy = tensor(xypos).apply{ reshape(-1, 2) }

            val xmax = xy[all, n(0)].max()
            val ymax = xy[all, n(1)].max()
            val scale = NORMALIZE_MAX/max(xmax, ymax)

            val sp = xy.shape
            res[n(i), r(0, sp[0]), r(0, 1)] = xy*scale
            res[n(i), r(0, sp[0], n(2))] = tensor(1)
        }
        res
    }

これでいい気がしてきた。これでいこう。

追記：　結局こんな感じになった。

    tensor_ns {
        val one = tensor(xylist).reshape(-1, 2)
        val len = one.shape[0]
        strokes[n(curIndex), r(0, len), r(0, 2)] = one*scale
        strokes[n(curIndex), r(0, len), n(2)] = INPUT_TYPE_POS
    }

なかなか良く無い？すげー快適。

実際の使用例