Jetson xavierをCUDAで動かして、さしあたって10倍早くなる

MNISTの手書き数字データをアソシアトロンで認識させ、評価したいと思っている。6万個の学習データがあるので、想起ベクトルを75264次元にしている。1文字のピクセル数は文字28ビット×28ビット=784で、文字データのラベルは、1バイトなので本当は、1ピクセル256次元使いたいのだが、そうすると、私のJetson xavierでは、メモリを食い尽くしてしまうので、1文字の1バイトデータを96bit(CUDAは32の倍数が好みだから)に縮小して、結局、784×96=75264で、結局上の次元になる。

だから、アソシアトロンの想起マトリクスは75264×75264=5664669696要素を持つことになって、巨大である。でもこれならば、なんとかGPUのグローバルメモリを取れて、データのコミュニケーションができる。ただし、1データをunsigned shortの2バイトで、マトリクスの右上三角と左下三角は対称的なので、片側の三角だけをデータで取るなどの操作をしなければならない。まともに作ると、グローバルメモリに乗らないのだ。

6万個のデータをすべて想起行列に載せることはまだやっていないが、約50個のデータを載せてみた。XavierのGPUを使わずCPU(8コア乗っている)だけで1スレッドだけ使って動かすと、5分かかる。timeコマンドを使って計算した結果は次のようなものである。

real 5m47.163s
user 5m46.680s
sys 0m0.456s

まともに6万個やると、100時間、つまり4日ほどまるまるコンピュータを動かさないといけない。

これをJetson XavierのGPU(512コア)で並列計算させると、50個30秒で終わる。timeコマンドの結果は次のようになっている。

real 0m30.672s
user 0m28.764s
sys 0m1.656s

10倍早い。それでも6万個やると10時間かかる。1個のデータの計算に用いている総スレッド数は、2次元で、75264×75264=5664669696個にもなるが、いろいろなオーバーヘッドから、10倍にしかなっていない。まだ、ナイーブなプログラム(普通のCPUのプログラムをそのままGPU用にかえただけのもの)なので、パフォーマンスが出ていないのだろうと思う。

ちなみに、6年ほど前に購入したNVIDIAのTelsa k20cで、試してみよう。こちらは2千数百個のコアを持っている。同じように実行した結果は次のようなものだった。

real 0m48.809s
user 0m43.691s
sys 0m4.450s

最新のXavierよりも、遅い。

プログラムは、調整しおえたらGitHUに上げると思う。