MXNet Tutorialを試す（1）: NDArray - CPU/GPU上での命令的なテンソル操作 #Python

MXNet Tutorialを順番にやっていくメモ(頑張って最後まで進む……)

今回は一番最初のNDArray - Imperative tensor operations on CPU/GPU。

NDArrayとは

numpy.ndarrayライクなモジュール
これにより命令的に書けるのがMXNetの大きな特徴

CPU, GPU, multi-GPU環境のサポート
遅延で実行されるので並列で計算実行可能

という点がnumpy.ndarrayと異なっている

重要なAPIs

ndarray.shape: アレイの次元（ここでいう次元とは軸の数のこと）
ndarray.dtype: numpyタイプのオブジェクト、要素の型を示す
ndarray.size: 要素の数
ndarray.context: これが特徴的、アレイが保持される場所(例:cpu() gpu(1))

アレイの作成

いくつかの作成方法がある

Pythonのリストやタプルから生成

import mxnet as mx
# Pythonのリストから1次元のアレイを生成
a = mx.nd.array([1,2,3])
# ネストしたリストから2次元アレイの生成
b = mx.nd.array([[1,2,3], [2,3,4]])

{'a.shape':a.shape, 'b.shape':b.shape}

OUT

    {'a.shape': (3,), 'b.shape': (2, 3)}

numpy.ndarrayオブジェクトから生成

import numpy as np
import math
c = np.arange(15).reshape(3,5)
# numpy.ndarrayオブジェクトから2次元のアレイを生成
a = mx.nd.array(c)
{'a.shape':a.shape}

OUT

    {'a.shape': (3, 5)}

また要素の型をdtypeで指定可能（numpyの型で指定）、デフォルトではfloat32

# デフォルトではfloat32
a = mx.nd.array([1,2,3])
# int32
b = mx.nd.array([1,2,3], dtype=np.int32)
# 16-bit float
c = mx.nd.array([1.2, 2.3], dtype=np.float16)
(a.dtype, b.dtype, c.dtype)

OUT

    (numpy.float32, numpy.int32, numpy.float16)

アレイのサイズはは分かっているが要素の値が未定の時、プレースホルダを作成する関数もある

numpyと一緒の初期化の関数たち、ということらしい

# (2,3)で0で埋められたアレイ
a = mx.nd.zeros((2,3))
# 同じ形で1で埋めたアレイ
b = mx.nd.ones((2,3))
# 7で埋めたアレイ
c = mx.nd.full((2,3), 7)
# ランダムで埋めたアレイ
# 値はメモリの状態依存
d = mx.nd.empty((2,3))

print(a.asnumpy())  # 出力追加, 次の節で説明あり
print(b.asnumpy())
print(c.asnumpy())
print(d.asnumpy())

OUT

    [[ 0.  0.  0.]
     [ 0.  0.  0.]]
    [[ 1.  1.  1.]
     [ 1.  1.  1.]]
    [[ 7.  7.  7.]
     [ 7.  7.  7.]]
    [[  0.00000000e+00   0.00000000e+00   1.23216977e-36]
     [  0.00000000e+00   6.24967988e+34   4.56529027e-41]]

アレイの出力

アレイの中身を見たいときはasnumpy()を使ってnumpy.ndarray型にする

Numpyでの出力レイアウトは

最後の軸（2次元だと列）は左から右へ出力
最後から2番目の軸（2次元だと行）は上から下へ出力
残りも上から下へ出力、それぞれのスライスは空行でセパレートされる

b = mx.nd.arange(24).reshape((3,2,4))  # ※列数変更
b.asnumpy()  # 左から右に4、上から下へ2、上から下へ3(空行あり)

OUT

    array([[[  0.,   1.,   2.,   3.],
            [  4.,   5.,   6.,   7.]],

           [[  8.,   9.,  10.,  11.],
            [ 12.,  13.,  14.,  15.]],

           [[ 16.,  17.,  18.,  19.],
            [ 20.,  21.,  22.,  23.]]], dtype=float32)

基本的な演算

NDArrayでは基本的な演算子は要素ごとに計算される

a = mx.nd.ones((2,3))
b = mx.nd.ones((2,3))  # a, b共に2行3列の1で埋められた行列
# 要素ごとに足し算
c = a + b
print('c:')
print(c.asnumpy())
# 要素ごと符号反転
d = - c
print('d:')
print(d.asnumpy())
# 要素ごとに自乗しSin関数適用、その後転置
e = mx.nd.sin(c**2).T
print('e:')
print(e.asnumpy())
# 要素ごとに最大値比較
f = mx.nd.maximum(a, c)
print('f:')
print(f.asnumpy())

OUT

    c:
    [[ 2.  2.  2.]
     [ 2.  2.  2.]]
    d:
    [[-2. -2. -2.]
     [-2. -2. -2.]]
    e:
    [[-0.7568025 -0.7568025]
     [-0.7568025 -0.7568025]
     [-0.7568025 -0.7568025]]
    f:
    [[ 2.  2.  2.]
     [ 2.  2.  2.]]

Numpyと同様、*は要素ごとのかけ算を示すので、行列同士のかけ算にはdotを使う

a = mx.nd.arange(4).reshape((2,2))
b = a * a
c = mx.nd.dot(a,a)
print("b: \n%s, \nc: \n%s" % (b.asnumpy(), c.asnumpy()))

OUT

    b: 
    [[ 0.  1.]
     [ 4.  9.]], 
    c: 
    [[  2.   3.]
     [  6.  11.]]

+=, *= はその場でアレイを修正するので新しいアレイを生成するためのメモリを割り当てない

a = mx.nd.ones((2,2))
b = mx.nd.ones(a.shape)
b += a
b.asnumpy()

OUT

    array([[ 2.,  2.],
           [ 2.,  2.]], dtype=float32)

インデクシングとスライシング

[]演算子は0軸方向へ適用される

a = mx.nd.array(np.arange(6).reshape(3,2))
a[1:2] = 6  # 2行目（インデックス1以上2未満）を6に
a[:].asnumpy()

OUT

    array([[ 0.,  1.],
           [ 6.,  6.],
           [ 4.,  5.]], dtype=float32)

slice_axisメソッドも使える

d = mx.nd.slice_axis(a, axis=1, begin=1, end=2)  # 2列目(インデックス1)選択
d.asnumpy()

OUT

    array([[ 1.],
           [ 6.],
           [ 5.]], dtype=float32)

形の操作

reshape で全体の要素数が保たれる限り、アレイの形を変更できる

a = mx.nd.array(np.arange(24))
b = a.reshape((2,3,4))
b.asnumpy()

OUT

    array([[[  0.,   1.,   2.,   3.],
            [  4.,   5.,   6.,   7.],
            [  8.,   9.,  10.,  11.]],

           [[ 12.,  13.,  14.,  15.],
            [ 16.,  17.,  18.,  19.],
            [ 20.,  21.,  22.,  23.]]], dtype=float32)

concat メソッドは軸1方向（行方向）にそって重ねる（他の軸の要素数は同じである必要あり）

最初の軸の数を保つ形で、という意図の模様

a = mx.nd.ones((2,3))
b = mx.nd.ones((2,3))*2
c = mx.nd.concat(a,b)   # bを右から結合するイメージ
c.asnumpy()

OUT

    array([[ 1.,  1.,  1.,  2.,  2.,  2.],
           [ 1.,  1.,  1.,  2.,  2.,  2.]], dtype=float32)

要素同士の演算

原題reduce

sumやmeanのような関数でアレイからスカラーへの計算

a = mx.nd.ones((2,3))
b = mx.nd.sum(a)
b.asnumpy()

OUT

    array([ 6.], dtype=float32)

ブロードキャスト

長さが1の軸を複製

a = mx.nd.array(np.arange(6).reshape(6,1))  # 列方向のベクトル
b = a.broadcast_to((6,4))  # 行方向に沿って４つ複製
b.asnumpy()

OUT

    array([[ 0.,  0.,  0.,  0.],
           [ 1.,  1.,  1.,  1.],
           [ 2.,  2.,  2.,  2.],
           [ 3.,  3.,  3.,  3.],
           [ 4.,  4.,  4.,  4.],
           [ 5.,  5.,  5.,  5.]], dtype=float32)

複数の軸方向へのブロードキャストも可能

少し分かりづらいが2,3番目の要素が複製されている

c = a.reshape((2,1,1,3))
d = c.broadcast_to((2,2,2,3))  # 行方向に複製+行列を複製
print('c:')
print(c.asnumpy())
print('d:')
print(d.asnumpy())

OUT

    c:
    [[[[ 0.  1.  2.]]]


     [[[ 3.  4.  5.]]]]
    d:
    [[[[ 0.  1.  2.]
       [ 0.  1.  2.]]

      [[ 0.  1.  2.]
       [ 0.  1.  2.]]]


     [[[ 3.  4.  5.]
       [ 3.  4.  5.]]

      [[ 3.  4.  5.]
       [ 3.  4.  5.]]]]

コピー

NDArrayを他のPythonの変数へ渡すと、同じNDArrayへの参照が渡される

元のアレイを上書きせずに操作したい場合はcopyを使う

a = mx.nd.ones((2,2))
b = a
b is a # will beTrue

OUT

    True

b = a.copy()
b is a  # will be False

OUT

    False

この場合新しいNDArrayをbへと割り当てる。追加でメモリを割り当てたくない場合はcopytoメソッドか[]を利用する

b = mx.nd.zeros(a.shape)
c = b
c[:] = a
d = b
a.copyto(d)
print(c is b, d is b)  # Both will be True
print(c is a, d is a)  # 追加、[]もcopytoも値のコピーだけ行う

OUT

    True True
    False False

print(a.asnumpy(), b.asnumpy(), c.asnumpy(), d.asnumpy())  # 出力結果確認

OUT

    [[ 1.  1.]
     [ 1.  1.]] [[ 1.  1.]
     [ 1.  1.]] [[ 1.  1.]
     [ 1.  1.]] [[ 1.  1.]
     [ 1.  1.]]

発展トピック

GPUサポート

NDArrayはndarray.contextに実行環境の情報が格納されている

USE_CUDA=1の状態でMXNetがコンパイルされかつ1つでもNVIDIAのGPUが備わっていれば、コンテクストをmx.gpu(0)のようにしているすることでGPU上で実行できるようになる

gpu_device=mx.gpu() # GPUがないとエラー、mx.cpu() へと変えること

def f():
    a = mx.nd.ones((100,100))
    b = mx.nd.ones((100,100))
    c = a + b
    print(c)
# デフォルトではmx.cpu() が使用される
f()

# GPU0上へ実行環境を変更
with mx.Context(gpu_device):
    f()

OUT

    <NDArray 100x100 @cpu(0)>
    <NDArray 100x100 @gpu(0)>

アレイ作成の時に明示的に宣言も可

a = mx.nd.ones((100, 100), gpu_device)
a

OUT

    <NDArray 100x100 @gpu(0)>

現在のところMXNetでは、2つのアレイ同士の計算の際これらは同じデバイス上にある必要がある

そのためいくつかコピーする関数がある

a = mx.nd.ones((100,100), mx.cpu())
b = mx.nd.ones((100,100), gpu_device)
c = mx.nd.ones((100,100), gpu_device)
a.copyto(c)  # copyto() でCPUからGPUへコピー
d = b + c
e = b.as_in_context(c.context) + c  #　cと同じコンテクストで実行（上と同じ）
{'d':d, 'e':e}

OUT

    {'d': <NDArray 100x100 @gpu(0)>, 'e': <NDArray 100x100 @gpu(0)>}

（分散）ファイルシステムからの/へのシリアライズ

ディスクへの読み書きには2つのシンプルな方法があり

1つ目はpickle、NDArrayも他のPythonオブジェクトと同じようにpickleで処理可能

import pickle as pkl
a = mx.nd.ones((2, 3))
# ダンプして保存
data = pkl.dumps(a)
pkl.dump(data, open('tmp.pickle', 'wb'))
# ロード
data = pkl.load(open('tmp.pickle', 'rb'))
b = pkl.loads(data)
b.asnumpy()

OUT

    array([[ 1.,  1.,  1.],
           [ 1.,  1.,  1.]], dtype=float32)

2番目の方法はsave,loadメソッドを使って直接読み書き、NDArrayのリストも可能

a = mx.nd.ones((2,3), gpu_device)  # コンテクストも保存される模様
b = mx.nd.ones((5,6))
mx.nd.save("temp.ndarray", [a,b])  # NDArrayのリストを保存
c = mx.nd.load("temp.ndarray")
c

OUT

    [<NDArray 2x3 @gpu(0)>, <NDArray 5x6 @cpu(0)>]

同じ要領で辞書も可能

d = {'a':a, 'b':b}
mx.nd.save("temp.ndarray", d)
c = mx.nd.load("temp.ndarray")
c

OUT

    {'a': <NDArray 2x3 @gpu(0)>, 'b': <NDArray 5x6 @cpu(0)>}

save, load メソッドを使う方が好ましい
1. Pyhon以外とのインターフェースのやりとりに利用できる

# Python
a = mx.nd.ones((2, 3))
mx.nd.save("temp.ndarray", [a,])

Rでこれを読める

a <- mx.nd.load("temp.ndarray")
as.array(a[[1]])
##      [,1] [,2] [,3]
## [1,]    1    1    1
## [2,]    1    1    1

2.S3やHDFSへ直で読み書きできる

mx.nd.save('s3://mybucket/mydata.ndarray', [a,])  # コンパイル時に USE_S3=1 指定の必要
mx.nd.save('hdfs///users/myname/mydata.bin', [a,])  # コンパイル時に USE_HDFS=1 指定の必要

遅延評価と自動並列化

パフォーマンス向上のために遅延で評価が行われる
例えばPythonで a = b + 1を呼んだ時、Python側はバックエンドにこの操作を渡すのみ

Pythonが他の計算を続けられる、オーバーヘッドの大きいフロントエンドの言語（ユーザーが触るPython部分ということだろう）には効果的
自動並列化などの効率化が容易

という利点があるため

バックエンドはデータの依存関係を把握し正確に計算を実行する。
フロントエンド側のユーザーからも触ることができて、明示的にwait_to_read()を呼ぶことで計算終了まで待つことができる

asnumpy()のような他のパッケージのアレイへ変換するようなメソッドは暗黙的にこれを呼んでいる

import time
def do(x, n):
    """バックエンドへ計算をプッシュ"""
    return [mx.nd.dot(x,x) for i in range(n)]
def wait(x):
    """結果が出るまで待つ"""
    for y in x:
        y.wait_to_read()

tic = time.time()
a = mx.nd.ones((1000,1000))
b = do(a, 50)
print('time for all computations are pushed into the backend engine:\n %f sec' % (time.time() - tic))
wait(b)
print('time for all computations are finished:\n %f sec' % (time.time() - tic))

OUT

    time for all computations are pushed into the backend engine:
     0.003978 sec
    time for all computations are finished:
     0.646166 sec

データの読み書きの依存関係以外にも、依存関係のない計算を並列に走らせることもバックエンドはできる

a = mx.nd.ones((2,3))
b = a + 1  # これと
c = a + 2  # これは並列実行可能
d = b * c

# CPUとGPUを直列に実行
n = 10
a = mx.nd.ones((1000,1000))
b = mx.nd.ones((6000,6000), gpu_device)
tic = time.time()
c = do(a, n)
wait(c)
print('Time to finish the CPU workload: %f sec' % (time.time() - tic))
d = do(b, n)
wait(d)
print('Time to finish both CPU/GPU workloads: %f sec' % (time.time() - tic))

OUT

    Time to finish the CPU workload: 0.128905 sec
    Time to finish both CPU/GPU workloads: 3.721957 sec

# 並列に実行
tic = time.time()
c = do(a, n)
d = do(b, n)
wait(c)
wait(d)
print('Both as finished in: %f sec' % (time.time() - tic))

OUT

    Both as finished in: 3.593669 sec

直列の時にGPUでかかった分短縮されてる

Python 3 + Ubuntu, GPU環境で学習中
全文和訳しているわけではないです、あくまでメモ程度に
Jupyterからの出力に手を入れただけなのでレイアウト崩れるやも…

次はSymbolモジュールの予定。