MLMultiArrayを見据えたUnsafeMutableRawPointerの取り回し


MLMultiArrayの情報が少ない

なので、備忘録を兼ねて学んだことをメモしておこうと思います。

そもそもMLMultiArrayとは

coreMLのデータをxcodeに取り込むとモデルのInputとして指定されている型がMLMultiArrayになります。(特殊な例として画像などは別の型で入力できますが、汎用的な数値の入力としてはこちらを利用します。)

大層な名前が付いていますが言ってしまえば型付のただの配列です。裏側では指定された型のサイズを指定されたサイズ分メモリ上に連番で確保してある、という所謂Cでの配列の考え方そのままで問題ないという認識です。

MLMultiArrayの基本

このクラスの取り回しを全く知らない人からすると、例えば numpy で言うところの

variable[0,0,0] = 10

のようなただの代入、もしくは値の取得ですらどうやったら良いのか苦戦します。(しました。
それくらいネット上に情報が少ないんですよね…

MLMultiArrayではこのようにします。

variable[[0,0,0]] = 10

もう少しわかりやすく書くと

let index = [NSNumber(value: 0), NSNumber(value: 0), NSNumber(value: 0)]
variable[index] = 10

こうなります。添字に全次元のIndexをもつ配列を渡す事で参照できるという事ですね。

次元の順番

MLMultiArrayと直接関係はないのですが、coreMLでは画像を扱う際、(チャンネル, 高さ, 幅)という順番で扱います。そのため、Keras(Tensorflow)で (高さ, 幅, チャンネル) の形式で扱っているとcoremltoolsでconvertした時に入れ替わっているのでMLMultiArrayで入力を作るときには注意が必要です。

実際には少し語弊があって、画像を扱う際というよりは3次元の配列は勝手に画像扱いになるので否応なしに次元が入れ替わります。
画像だと高さや幅とチャンネル数が一致することなんてほぼないのでエラーが出て気づけますが、10x10x10の立方体のようなイメージで値を入力するモデルの場合、知らないままだとcoreMLとKeras(Tensorflow)で全然違う入力をしている事になり、結果が全然違うという事になりかねないので注意が必要です。

MLMultiArrayの変形

例えばnumpyでは入力データを工夫しようと思った時に、取り回しがとても簡単です。例として4x5x3(h,w,c)の画像を縦につないで8x5x3としようとすると

a = np.zeros(shape=(4,5,3))
b = np.zeros(shape=(4,5,3))

concatenated = np.concatenate((a, b), axis=0)

と簡単に書けます。横に繋ごうがチャンネル次元で繋ごうがconcatenateのaxisを調整するだけです。

では同じことをSwift上でMLMultiArrayに対して行うにはどうしたらよいか。もし作ったモデルをアプリ上で動かすのであれば同じ処理をSwiftで行うことは避けては通れません。

しかし、現状(おそらく)MLMultiArrayにはnp.concatenateのような便利関数はありません。そのためポインタを理解した上で、concatenate後のサイズの配列をメモリ上に確保した上でそこにコピーしていく、という作業が必要になります。

(ポインタを無視して愚直に1要素ずつ代入をループで回しても可能ですが、サイズが大きくなるとループ数が膨大になるのでおすすめしません。)

上でnumpyでやった事と同じことをしようと思うと

// 1次元目がチャンネル次元になっている事に注意
let a = try! MLMultiArray(shape: [3, 4, 5], dataType: .double)
let b = try! MLMultiArray(shape: [3, 4, 5], dataType: .double)

// メモリを確保
let resultArray = try! MLMultiArray(shape: [3, 8, 5], dataType: .double)

// UnsafeMutableRawPointerだと型が付いておらずassignできないのでbindMemoryで型付け
let resultArrayPointer = resultArray.dataPointer.bindMemory(to: Double.self, capacity: 3*8*5)

// assignで値をコピー
for i in 0..<8 {
    resultArrayPointer.advanced(by: 8*5*i).assign(from: a.dataPointer.bindMemory(to: Double.self, capacity: 3*4*5).advanced(by: 4*5*i), count: 4*5)
    resultArrayPointer.advanced(by: 8*5*i + (4*5)).assign(from: b.dataPointer.bindMemory(to: Double.self, capacity: 3*4*5).advanced(by: 4*5*i), count: 4*5)
}

こんな感じになります。
おそらく最後の assignで値をコピー が何をしているのかさっぱり分からないという人が多いと思いますのでちょっと解説します。(byやcountは敢えて掛け算でそれぞれの次元数を指していることが分かるように記載しているので、Pointerのイメージが付いている人は少し考えれば分かると思います。)

UnsafeMutablePointerを使ったassignのイメージ

※ ここからはわかりやすいように自分の中の「イメージ」で書きます。そのため厳密にハードウェア的な定義などと比べると微妙に異なることを書いているかもしれませんがご了承ください。

Cのポインタを1から説明し始めると相当長くなってしまうので、基本的な概念だけ最初に説明します。ポインタというのはメモリ上のアドレスを指す変数で、例えばメモリ上の100番地を指すのであれば100といった数字が入っていると考えてください。

なので上の例でいう resultArrayPointer には「X番地にデータがありますよ」という情報が入っており、コードが示しているのは「X番地からDoubleのサイズでcapacity個のデータが入るだけの枠がありますよ」という情報が含まれています。

簡単のために
X = 100
Doubleのサイズ = 8
capacity = 20

と考えると、100番地から8のサイズで20個分の枠が確保されています、という事になります。すなわち100〜259番地が今回用意された領域という事ですね。(厳密にはその他の情報を保持するための余白などがあったりしますが、今回は考えません。)
そしてDoubleの配列だと考えて値を取得する場合は100番地、108番地、116番地、・・・を見ていく事になります。

ここで「3x4x5みたいな複数次元の配列はどうやって入っているの?」と疑問を持った方は鋭いです。
実は何次元のデータだとしても、メモリ上には1列で並んでいます。では、どのような順番で値が入っているのかですが、ここからはMLMultiArrayの話になります。CやSwiftの他の似たようなクラスではどうなっているのかはこの限りではありませんので注意してください。(基本は同じだと思いますが。)

numpyの扱いに長けている人には「np.flattenした時に出来上がる順番」と説明すれば一発で終わると思います。これで分からない人は以下の説明を見た上でいろんな3次元の配列を作ってnumpyで実験してみると良いでしょう。

言葉で説明するのは難しいのですが、高次元側から順に番号をつけていく、と考えるとわかりやすいと思います。以下は3次元の場合の例示です。

let a = try! MLMultiArray(shape: [3, 4, 5], dataType: .double)

a[[0,0,0]] ・・・1番
a[[0,0,1]] ・・・2番
a[[0,0,2]] ・・・3番
a[[0,0,3]] ・・・4番
a[[0,0,4]] ・・・5番
a[[0,1,0]] ・・・6番
a[[0,1,1]] ・・・7番
a[[0,1,2]] ・・・8番
a[[0,1,3]] ・・・9番
...
a[[0,3,3]] ・・・19番
a[[0,3,4]] ・・・20番
a[[1,0,0]] ・・・21番
...
a[[2,3,4]] ・・・60番

高次元(後ろの次元)から順に増えていくように順序付けられます。

これを理解した上で、numpyでいう「縦に画像を繋ぐ」を実現しようとすると、MLMultiArrayではアドレスでいう所の前から順に埋めていけば良いわけではない事が分かります。
「縦」の次元が2番目の次元なので、aの21番目が入るべき場所はresultArrayPointerの41番目になります。(Indexでいうとaの20がresultArrayPointerの40です。)

これらを踏まえて行なっている処理を言葉で書くと、

resultArrayPointerの41番目(advanced(by: 8*5)で、ちゃんとDouble型のサイズを踏まえて41番目を指すポインタが返ってきます)を起点として aの21番目〜40番目までをassignする
resultArrayPointerの61番目を起点として bの21番目〜40番目までをassignする

という事になります。もちろんこれはループの2回目の処理なので、実際のコードではiを使ってループごとに適切な位置へのassignを行っています。

これにより最終的に2次元目でabをつなげたMLMultiArrayが完成するわけです。

もちろんチャンネル次元で結合したい場合、3次元目で結合したい場合はそれぞれこの概念を理解した上で処理の仕方を変える必要があります。(numpyのaxis指定どんだけ便利なんだっていう話

まとめ

唐突ですが、これ以上書いても混乱する人は余計混乱するし、理解できた人は理解できたと思うのでまとめます。

  • MLMultiArrayでは画像でいうところのチャンネル次元が1次元目なので気をつけよう!
  • MLMultiArray.dataPointerはUnsafeMutableRawPointer(型が無いポインタ)なので取り回すためにはbindMemoryで型付けしよう!
  • pointerで指定した時の値の1列の並び方を把握しよう!
  • assing元やassign先の計算を間違えると確保していない領域への書き込みになってしまうので気をつけよう!【重要】