numpy.memmapでnumpy arrayメモリに乗らない問題を解消


例題

メモリに乗らないnumpy array np.random.rand(100000, 99999).astype(dtype='float32')を一つのファイルに書き出したい、そして読み込みたい。
こういう事がしたい:

import numpy as np

arr = np.random.rand(100000, 99999).astype(dtype='float32')
np.save('np_save.npy', arr)

しかしarrが約40GBになるので、メモリに乗らず、Killされる:

>>> import numpy as np
>>> arr = np.random.rand(100000, 99999).astype(dtype='float32')
Killed: 9

100000行いっぺんに生成するのではなく、ループで生成してアペンドしていこう。
100行生成のループを1000回することにします。

リストにアペンドして、np.vstackして、np.saveするという方法はメモリに乗る場合うまくいくのですが、こちらも40GBがメモリに展開されるので、Killされる:

>>> arr_list = []
>>> for i in range(1000):
...     arr = np.random.rand(100, 99999).astype(dtype='float32')
...     arr_list.append(arr)
...

>>> np.save('np_save.npy', np.vstack(arr_list))
Killed: 9

解決策

np.memmapでメモリを消費せず、ファイルに書き込める。まずshapeを指定して、memory-mapを生成。ループ内でデータを指定したスライスに書き込む。

import numpy as np
fp = np.memmap('np_save.dat', dtype='float32', mode='w+', shape=(100000, 99999))
del fp
for i in range(1000):
    arr = np.random.rand(100, 99999).astype(dtype='float32')
    fp = np.memmap('np_save.dat', dtype='float32', mode='r+', shape=(100000,99999))
    fp[i*100:(i+1)*100] = arr
    del fp

メモリ使用はこんな感じで抑えられています:

np.memmapのいいところは読み込みも瞬時にでき、ディスク上のデータのセグメントをスライスで読み込めます。使うデータのみメモリ展開できます。

>>> import numpy as np
>>> newfp = np.memmap('np_save.dat', dtype='float32', mode='r', shape=(100000, 99999))
>>> newfp
memmap([[0.9187665 , 0.7200832 , 0.36402512, ..., 0.88297397, 0.90521574,
         0.80509114],
        [0.56751174, 0.06882019, 0.7239285 , ..., 0.04442023, 0.26091048,
         0.07676199],
        [0.94537544, 0.06935687, 0.13554753, ..., 0.3666087 , 0.6137967 ,
         0.1677396 ],
        ...,
        [0.97691774, 0.08142337, 0.18367094, ..., 0.30060798, 0.5744949 ,
         0.3676454 ],
        [0.13025525, 0.36571756, 0.17128325, ..., 0.8568927 , 0.9460005 ,
         0.61096454],
        [0.22953852, 0.00882731, 0.15313177, ..., 0.90694803, 0.17832297,
         0.45886058]], dtype=float32)
>>> newfp[30:40]
memmap([[0.23530068, 0.64027864, 0.7193347 , ..., 0.0991324 , 0.71703744,
         0.0429478 ],
        [0.7835149 , 0.20407963, 0.23541291, ..., 0.12721954, 0.6010988 ,
         0.8794886 ],
        [0.62712234, 0.76977116, 0.4803686 , ..., 0.24469836, 0.7741827 ,
         0.14326976],
        ...,
        [0.4855294 , 0.15554492, 0.6792018 , ..., 0.23985237, 0.59824246,
         0.88751584],
        [0.80865365, 0.73577726, 0.9209202 , ..., 0.9908406 , 0.66778165,
         0.16570805],
        [0.63171065, 0.48431855, 0.57776374, ..., 0.76027304, 0.930301  ,
         0.20145524]], dtype=float32)
>>> arr = np.array( newfp[30:40])
>>> arr
array([[0.23530068, 0.64027864, 0.7193347 , ..., 0.0991324 , 0.71703744,
        0.0429478 ],
       [0.7835149 , 0.20407963, 0.23541291, ..., 0.12721954, 0.6010988 ,
        0.8794886 ],
       [0.62712234, 0.76977116, 0.4803686 , ..., 0.24469836, 0.7741827 ,
        0.14326976],
       ...,
       [0.4855294 , 0.15554492, 0.6792018 , ..., 0.23985237, 0.59824246,
        0.88751584],
       [0.80865365, 0.73577726, 0.9209202 , ..., 0.9908406 , 0.66778165,
        0.16570805],
       [0.63171065, 0.48431855, 0.57776374, ..., 0.76027304, 0.930301  ,
        0.20145524]], dtype=float32)

機械学習でデータを扱う時に使えると思います。

オススメしない方法

直接ファイルにアペンドする方法として、np.savetxtがありますが、オススメしません!とても遅く、ファイルが巨大になります(元のデータサイズ以上):

with open('np_save.dat','ab') as f:
    for i in range(1000):
        arr = np.random.rand(100, 99999).astype(dtype='float32')
        np.savetxt(f,arr)