はじめに

Jetson Nanoを入手しました。
とりあえず何かということで、機械学習の練習にMNISTの手書き数字データ使ったいわゆるHello WorldをGPU上で実施して、どれだけCPUより早いのか体感してやろうと思いました。
以前、numpy(CPU)で実施したときのコードが残っていたのでcupy(GPU)に書き換えて動かせばきっと早くなるだろうと思い試したところ、なんとnumpyが約1,000秒に対して、cupyは約5,000秒。すごい！遅くなってる・・・

ということで原因を調査したところ、分析の時間は1/4になっていて前処理の時間が大幅に増えていることがわかりました。
本記事は主に画像や正解ラベルを読み込むための前処理を高速化(というかまとも化)するための記録です。

前提として、今回はsklearnなど、各種ML系ライブラリに含まれているMNISTデータではなく、自分でダウンロードしてイチから読み込んでみています。
また、本論から外れるため一切記述していませんが、MNIST文字列を学習させるためのコードは図書ゼロから作るDeep Learning
――Pythonで学ぶディープラーニングの理論と実装をベースに実装しています。

やったこと

データの整理

MNISTには以下4つのデータがあります。

#	データセット名	簡単な説明
1	train-images-idx3-ubyte.gz	トレーニング用の手書き画像(28x28)の60,000枚セット
2	train-labels-idx1-ubyte.gz	トレーニング画像に対する正解ラベル(60,000個)
3	t10k-images-idx3-ubyte.gz	テスト用の手書き画像(28x28)の10,000枚セット
4	t10k-labels-idx1-ubyte.gz	テスト画像に対する正解ラベル(10,000個)

画像ファイルの構造

トレーニング/テスト用画像ファイルは以下のような構造になっています。

#	開始バイト数	終了バイト数	簡単な説明
1	0	3	magicナンバー(処理にはおそらく不要
2	4	7	イメージ数(トレーニング用は60,000 / テスト用は10,000)
3	8	11	画像一枚あたりの行数(28)
4	12	15	画像一枚あたりの列数(28)
5	16	16	画像一枚目の1行1列目の画素の値
6	17	17	画像一枚目の1行2列目の画素の値
...	...	...	...

同様に正解ラベルのファイルは以下のような構造です。

#	開始バイト数	終了バイト数	簡単な説明
1	0	3	magicナンバー(処理にはおそらく不要
2	4	7	イメージ数(トレーニング用は60,000 / テスト用は10,000)
3	8	8	1枚目の画像の正解の値(0~9)
4	9	9	2枚めの画像の正解の値(0~9)
...	...	...	...

これらをそれぞれ、cupyの配列に格納していきます。

以前numpyでやった方法

それなりの速度で動いてしまっていたので気にしていなかったのですが、numpyを使ったときは以下のようなコードで実現していました。

ここでは上記データのうち、1,2の部分のコードのみ示します。(3,4は同じなので)
ちなみにこれをJetson Nano上で実行すると、私の環境の場合、108[sec]で完了しました。
60,000枚の28x28の画像と考えると、きっとそんなところだろうと思ってあまり考えていませんでした。

トレーニング用の正解ラベルの読み込み

import gzip
import numpy as np

# ファイルのオープン
with gzip.open("mnist/train-labels-idx1-ubyte.gz", "rb") as file:
    label_data = file.read()

# meta的なデータの読み込み
magic = int.from_bytes(label_data[0:4], byteorder='big')
num_of_data = int.from_bytes(label_data[4:8], byteorder='big')
offset = 8

# one-hot表現として格納するため、前もって60,000 * 10の配列を用意
t_train  = np.ndarray((num_of_data,10), dtype='uint8')

# 1つずつ変換
for i in range(0, num_of_data):
  t_train [i] = np.zeros(10,dtype="int8")
  t_train [i][label_data[offset+i]] = 1

print ("shape: {}".format(t_train.shape))

トレーニング用画像の読み込み

# ファイルのオープン
with gzip.open("mnist/train-images-idx3-ubyte.gz", "rb") as file:
    dataset = file.read()

# meta的なデータの読み込み
magic = int.from_bytes(dataset[0:4], byteorder='big')
num_of_data = int.from_bytes(dataset[4:8], byteorder='big')
num_of_rows = int.from_bytes(dataset[8:12], byteorder='big')
num_of_cols = int.from_bytes(dataset[12:16], byteorder='big')
offset = 16

# 画像サイズの計算
data_size = num_of_rows * num_of_cols

# 画像の形に合う配列を事前に用意
x_train = np.ndarray((num_of_data, data_size), dtype='uint8')

# 一画素ずつ格納
for i in range(0, num_of_data):
    for j in range(0, data_size):
      x_train[i][j] = dataset[offset+i*data_size+j]

print ("shape: {}".format(x_train.shape))

1つ1つ手作業で家内制手工業みたいなぬくもり溢れるやり方ですが、numpyだとそれなりの時間で動いてくれます。

(こんなことやってるとnumpy使ってる意味ないかもですね)

これをそのままcupyに変換してみる

とりあえずcupyはnumpyと違いはあるものの概ね互換性があるということだったので、動作確認がてらなのでまずはざっくりと、以下のように書き換えてみました。

変更前

import gzip
import numpy as np

変更後

import gzip
import cupy as np

これで実行した結果、なんと2,881[sec]もかかってしまいました。
想定外過ぎて、一瞬jetson nanoを叩き割ることが頭をよぎりました。

修正

もっとうまいやり方はあるかもしれませんが、やれる範囲で修正してみました。

正解ラベルの読み込み

当初こういった形で一旦、要素数10のすべての要素がゼロの配列を用意して、正解のところだけ1を立てるようにしていました。

変更前

# one-hot表現として格納するため、前もって60,000 * 10の配列を用意
t_train  = np.ndarray((num_of_data,10), dtype='uint8')
offset = 8

# 1つずつ変換
for i in range(0, num_of_data):
  t_train [i] = np.zeros(10,dtype="int8")
  t_train [i][label_data[offset+i]] = 1

cupy配列はforでぶん回したりすると遅いということだったのでやり方を変えました。
これについては参考2を参考にしています。
下から見ていくと、cp.indentity(10)で10行10列の単位行列を用意してくれます。
これの添字に配列を渡すと配列のそれぞれの値に対しての行を選択してくれて、結果的にone-hot表現ができあがります。
2行目では元のデータを1文字ずつ整数値に変換しているだけです。

変更後

offset = 8
label = [int(s) for s in label_data[offset:]]
t_train = cp.identity(10)[label]

画像の読み込み

元々は以下のように大きさの合うnumpy配列を用意して、1つの画像の個数784(=28*28)個分のデータを1つずつ読み込ませていました。

変更前

# 画像の形に合う配列を事前に用意
x_train = np.ndarray((num_of_data, data_size), dtype='uint8')

# 一画素ずつ格納
for i in range(0, num_of_data):
    for j in range(0, data_size):
      x_train[i][j] = dataset[offset+i*data_size+j]

cupyが遅いならlistで処理しちゃえばいいじゃない、というところで、x_train_tmpというpythonネイティブのリストを用意してそこに画像を1枚ごとに切り分けて格納しています。

元々は1画素ずつやっていましたが、それも辛いので1画像単位でまとめてもってこれるようにしているので、その点でも多少は高速化できているかもしれません。

変更後

x_train_tmp = []
dataset_tmp = [int(s) for s in dataset[offset:offset+data_size*num_of_data]]
for i in range(0, num_of_data):
    x_train_tmp.append(dataset_tmp[i*data_size:(i+1)*data_size])
x_train = cp.array(x_train_tmp, dtype=cp.uint16)

結果

結果的に全体としては約3,045[sec] -> 223[sec]となり、当初の7%の速度まで全体を短縮できました。
前処理だけでいうと、約2,881[sec] -> 53[sec]です。

結論

GPUは文句なしに早い
cupyはnumpyよりも扱いに気をつけないととてつもなく損をすることがある
numpyのお作法をちゃんと勉強しよう

MNISTの手書き文字をcupy配列として読み込ませてみた