Scheme で MNIST からの学び

久しぶりに Mosh Scheme に触っている。良い機会なのでゼロから作るDeep Learning ―Pythonで学ぶディープラーニングの理論と実装を参考に Scheme で動く MNIST デモを実装しようと思い立つ。MNIST は 0 から 9 までの手書き数字の認識デモ。Neural Network フレームワークのデモやチュートリアルでよく使われるあれだ。なお Scheme には Python でいうところの numpy は存在しないので必要な行列演算はすべて Scheme で実装することとする。

データサイズ

train data 画像(28 x 28 x 1) と正解ラベル（0-9）
train data: 60000件
test data: 10000件

実装前の期待と予想

他の言語・フレームワークでデモ可能だったのだから CPU 上で十分速く動作する。
比較的小さなモデル・train data なので pure Scheme 実装で大丈夫だろう。
matrix-mul がボトルネックになるかもしれない。

技術的な選択

Matrix は Vector の Vector として表現する。Uniform Vector （格納する型が決まっているもの）もありだが Mosh では未実装なので見送り。
行列演算の高速実装は特に目指さない。ナイーブな実装で良い。

matrix-mul の実装

特に工夫のないナイーブな実装。３重ループ。

(define (matrix-mul a b)
  (unless (= (matrix-shape a 1) (matrix-shape b 0))
    (error "matrix-mul shapes don't match" (matrix-shape a) (matrix-shape b)))
  (let* ([nrows (matrix-shape a 0)]
         [ncols (matrix-shape b 1)]
         [m     (matrix-shape a 1)]
         [mat (matrix nrows ncols)])
    (define (mul row col)
      (let loop ([k 0]
                 [ret 0])
        (if (= k m)
            ret
            (loop (+ k 1) (+ ret (* (mat-at a row k) (mat-at b k col)))))))
    (do ((i 0 (+ i 1)))
        ((= i nrows) mat)
      (do ((j 0 (+ j 1)))
          ((= j ncols))
        (mat-at mat i j (mul i j))))))

結果1

2層の NN を hidden_size = 10 batch_size=3 で動かしてみたが 1 batch 回したプロファイラの結果。matrix-mul は 32165 回呼ばれ 95%（98 秒）を占めている。

Time%        msec      calls   name                              location
  95        98500      32165   (matrix-mul a b)                  <transcoded-textual-input-port <binary-input-port tests/mnist.scm>>:175
   1         1140      64344   (matrix-map proc a)               <transcoded-textual-input-port <binary-input-port tests/mnist.scm>>:96

結果2

1 は明らかに遅いので C++ 側で実装してみる。98秒->75秒だが思ったよりも速くない。逆に言えば VM ループは思ったよりも速い。

  92        74590      32165   matrix-mul                        
   1          680      64543   (matrix-element-wise op a...)     <transcoded-textual-input-port <binary-input-port tests/mnist.scm>>:186

ここで C++ 側の実装を注意して見る必要がある。途中の計算結果の保持に Scheme Object を使っていること。行列に Fixnum もしくは Flonum が存在するので型チェックをしてくれる Arithmetic::add と mul を利用して計算していること。これだと most inner loop で heap allocation が発生してしまう。

    Object ret = Object::makeVector(numRowsA);
    for (size_t i = 0; i < numRowsA; i++) {
        ret.toVector()->set(i, Object::makeVector(numColsB));
    }

    for (size_t i = 0; i < numRowsA; i++) {
        for (size_t j = 0; j < numColsB; j++) {
            Object value = Object::makeFlonum(0.0);
            for (size_t k = 0; k < numColsA; k++) {
                Object aik = a->ref(i).toVector()->ref(k);
                Object bkj = b->ref(k).toVector()->ref(j);
                value = Arithmetic::add(value, Arithmetic::mul(aik, bkj));
            }
            ret.toVector()->ref(i).toVector()->set(j, value);
        }
    }

結果3

もしも Uniform Vector なら内部データの型チェックは不要になり、すべての計算を double のまま行うことが可能なはず。これを模した適当な実装をやってみる。不正なデータを入力されたらインタプリタが落ちてしまうがデモなのでよしとする。 98秒が60m秒になり実用的な速度に落ち着いた。heap allocation は避けたいと思わせる結果。

    1           60      32165   matrix-mul

            double value = 0.0;
            for (size_t k = 0; k < numColsA; k++) {
                Object aa = a->ref(i).toVector()->ref(k);
                Object bb = b->ref(k).toVector()->ref(j);
                double aik = aa.isFixnum() ? aa.toFixnum() : aa.toFlonum()->value();
                double bkj = bb.isFixnum() ? bb.toFixnum() : bb.toFlonum()->value();
                value += aik * bkj;
            }
            ret.toVec

学びと改善のアイデア

Matrix を Vector で表現すると型チェックのオーバーヘッドがある
Matrix の内部表現は double/int の2次元配列とし C++ でアクセサと mul を実装すれば良い。
matrix-mul 以外の行列操作は Scheme 側で書いてもほぼ問題なさそう。多少遅くても呼ばれる回数が matrix-mul と比べて極端に少ない。例）add, element-wise multiply, transpose, shape。
matrix-mul は SIMD などでもっと速くなるだろう。

余談

numpy は歴史が長いだけあってよく出来ている。
オブジェクト指向の操作の方がコードが読みやすい場合がある a.t は a の転置。(t a) だと若干意味合いが違うし広すぎる。
numpy の broadcasting の仕組みは便利だし必須部品であることがよく分かった。スカラ lr と行列A の掛け算とか。
numpy の subset が SRFI で提案されたら面白いと思う。ただ上記の理由により Scheme だけで実装すると遅いと思うので悩ましい。各処理系が独自に実装することを期待はできなそうな。
現実的には Tensorflow フロントエンドで Scheme が使えれば良いのかもしれないが、それなら Python で十分。
Scheme で気軽に機械学習が楽しめるという Path は意外と遠い。Python 1強である理由も分かった。