melhor maneira de preservar matrizes numpy no disco

Estou procurando uma maneira rápida de preservar grandes matrizes numpy. Quero salvá-los no disco em um formato binário e depois lê-los de volta na memória de forma relativamente rápida. Infelizmente, o cPickle não é rápido o suficiente.

Encontrei numpy.savez e numpy.load . Mas o mais estranho é que o numpy.load carrega um arquivo npy no "mapa da memória". Isso significa que a manipulação regular de matrizes é realmente lenta. Por exemplo, algo assim seria muito lento:

#!/usr/bin/python
import numpy as np;
import time; 
from tempfile import TemporaryFile

n = 10000000;

a = np.arange(n)
b = np.arange(n) * 10
c = np.arange(n) * -0.5

file = TemporaryFile()
np.savez(file,a = a, b = b, c = c);

file.seek(0)
t = time.time()
z = np.load(file)
print "loading time = ", time.time() - t

t = time.time()
aa = z['a']
bb = z['b']
cc = z['c']
print "assigning time = ", time.time() - t;

mais precisamente, a primeira linha será realmente rápida, mas as linhas restantes que atribuem as matrizes objsão ridiculamente lentas:

loading time =  0.000220775604248
assining time =  2.72940087318

Existe alguma maneira melhor de preservar matrizes numpy? Idealmente, desejo poder armazenar várias matrizes em um arquivo.

Sou um grande fã do hdf5 por armazenar grandes matrizes numpy. Existem duas opções para lidar com o hdf5 em python:

http://www.pytables.org/

http://www.h5py.org/

Ambos são projetados para trabalhar com matrizes numpy de maneira eficiente.

Comparei o desempenho (espaço e tempo) de várias maneiras de armazenar matrizes numpy. Poucos deles suportam várias matrizes por arquivo, mas talvez seja útil de qualquer maneira.

Arquivos Npy e binários são realmente rápidos e pequenos para dados densos. Se os dados forem escassos ou muito estruturados, convém usar o npz com compactação, o que economizará muito espaço, mas custará algum tempo de carregamento.

Se a portabilidade é um problema, o binário é melhor que o npy. Se a legibilidade humana for importante, você terá que sacrificar muito desempenho, mas isso pode ser alcançado bastante bem usando o csv (que também é muito portátil, é claro).

Mais detalhes e o código estão disponíveis no repositório do github .

Agora existe um clone baseado em HDF5 picklechamado hickle!

https://github.com/telegraphic/hickle

import hickle as hkl 

data = { 'name' : 'test', 'data_arr' : [1, 2, 3, 4] }

# Dump data to file
hkl.dump( data, 'new_data_file.hkl' )

# Load data from file
data2 = hkl.load( 'new_data_file.hkl' )

print( data == data2 )

EDITAR:

Também existe a possibilidade de "selecionar" diretamente em um arquivo compactado, fazendo:

import pickle, gzip, lzma, bz2

pickle.dump( data, gzip.open( 'data.pkl.gz',   'wb' ) )
pickle.dump( data, lzma.open( 'data.pkl.lzma', 'wb' ) )
pickle.dump( data,  bz2.open( 'data.pkl.bz2',  'wb' ) )

Apêndice

import numpy as np
import matplotlib.pyplot as plt
import pickle, os, time
import gzip, lzma, bz2, h5py

compressions = [ 'pickle', 'h5py', 'gzip', 'lzma', 'bz2' ]
labels = [ 'pickle', 'h5py', 'pickle+gzip', 'pickle+lzma', 'pickle+bz2' ]
size = 1000

data = {}

# Random data
data['random'] = np.random.random((size, size))

# Not that random data
data['semi-random'] = np.zeros((size, size))
for i in range(size):
    for j in range(size):
        data['semi-random'][i,j] = np.sum(data['random'][i,:]) + np.sum(data['random'][:,j])

# Not random data
data['not-random'] = np.arange( size*size, dtype=np.float64 ).reshape( (size, size) )

sizes = {}

for key in data:

    sizes[key] = {}

    for compression in compressions:

        if compression == 'pickle':
            time_start = time.time()
            pickle.dump( data[key], open( 'data.pkl', 'wb' ) )
            time_tot = time.time() - time_start
            sizes[key]['pickle'] = ( os.path.getsize( 'data.pkl' ) * 10**(-6), time_tot )
            os.remove( 'data.pkl' )

        elif compression == 'h5py':
            time_start = time.time()
            with h5py.File( 'data.pkl.{}'.format(compression), 'w' ) as h5f:
                h5f.create_dataset('data', data=data[key])
            time_tot = time.time() - time_start
            sizes[key][compression] = ( os.path.getsize( 'data.pkl.{}'.format(compression) ) * 10**(-6), time_tot)
            os.remove( 'data.pkl.{}'.format(compression) )

        else:
            time_start = time.time()
            pickle.dump( data[key], eval(compression).open( 'data.pkl.{}'.format(compression), 'wb' ) )
            time_tot = time.time() - time_start
            sizes[key][ labels[ compressions.index(compression) ] ] = ( os.path.getsize( 'data.pkl.{}'.format(compression) ) * 10**(-6), time_tot )
            os.remove( 'data.pkl.{}'.format(compression) )


f, ax_size = plt.subplots()
ax_time = ax_size.twinx()

x_ticks = labels
x = np.arange( len(x_ticks) )

y_size = {}
y_time = {}
for key in data:
    y_size[key] = [ sizes[key][ x_ticks[i] ][0] for i in x ]
    y_time[key] = [ sizes[key][ x_ticks[i] ][1] for i in x ]

width = .2
viridis = plt.cm.viridis

p1 = ax_size.bar( x-width, y_size['random']       , width, color = viridis(0)  )
p2 = ax_size.bar( x      , y_size['semi-random']  , width, color = viridis(.45))
p3 = ax_size.bar( x+width, y_size['not-random']   , width, color = viridis(.9) )

p4 = ax_time.bar( x-width, y_time['random']  , .02, color = 'red')
ax_time.bar( x      , y_time['semi-random']  , .02, color = 'red')
ax_time.bar( x+width, y_time['not-random']   , .02, color = 'red')

ax_size.legend( (p1, p2, p3, p4), ('random', 'semi-random', 'not-random', 'saving time'), loc='upper center',bbox_to_anchor=(.5, -.1), ncol=4 )
ax_size.set_xticks( x )
ax_size.set_xticklabels( x_ticks )

f.suptitle( 'Pickle Compression Comparison' )
ax_size.set_ylabel( 'Size [MB]' )
ax_time.set_ylabel( 'Time [s]' )

f.savefig( 'sizes.pdf', bbox_inches='tight' )

savez () salva dados em um arquivo zip. Pode levar algum tempo para compactar e descompactar o arquivo. Você pode usar a função save () & load ():

f = file("tmp.bin","wb")
np.save(f,a)
np.save(f,b)
np.save(f,c)
f.close()

f = file("tmp.bin","rb")
aa = np.load(f)
bb = np.load(f)
cc = np.load(f)
f.close()

Para salvar várias matrizes em um arquivo, basta abrir o arquivo primeiro e depois salvar ou carregar as matrizes em sequência.

Outra possibilidade de armazenar matrizes numpy com eficiência é o Bloscpack :

#!/usr/bin/python
import numpy as np
import bloscpack as bp
import time

n = 10000000

a = np.arange(n)
b = np.arange(n) * 10
c = np.arange(n) * -0.5
tsizeMB = sum(i.size*i.itemsize for i in (a,b,c)) / 2**20.

blosc_args = bp.DEFAULT_BLOSC_ARGS
blosc_args['clevel'] = 6
t = time.time()
bp.pack_ndarray_file(a, 'a.blp', blosc_args=blosc_args)
bp.pack_ndarray_file(b, 'b.blp', blosc_args=blosc_args)
bp.pack_ndarray_file(c, 'c.blp', blosc_args=blosc_args)
t1 = time.time() - t
print "store time = %.2f (%.2f MB/s)" % (t1, tsizeMB / t1)

t = time.time()
a1 = bp.unpack_ndarray_file('a.blp')
b1 = bp.unpack_ndarray_file('b.blp')
c1 = bp.unpack_ndarray_file('c.blp')
t1 = time.time() - t
print "loading time = %.2f (%.2f MB/s)" % (t1, tsizeMB / t1)

e a saída do meu laptop (um MacBook Air relativamente antigo com um processador Core2):

$ python store-blpk.py
store time = 0.19 (1216.45 MB/s)
loading time = 0.25 (898.08 MB/s)

isso significa que ele pode armazenar muito rápido, ou seja, o gargalo é tipicamente o disco. No entanto, como as taxas de compressão são muito boas aqui, a velocidade efetiva é multiplicada pelas taxas de compressão. Aqui estão os tamanhos para essas matrizes de 76 MB:

$ ll -h *.blp
-rw-r--r--  1 faltet  staff   921K Mar  6 13:50 a.blp
-rw-r--r--  1 faltet  staff   2.2M Mar  6 13:50 b.blp
-rw-r--r--  1 faltet  staff   1.4M Mar  6 13:50 c.blp

Observe que o uso do compressor Blosc é fundamental para conseguir isso. O mesmo script, mas usando 'clevel' = 0 (ou seja, desativando a compactação):

$ python bench/store-blpk.py
store time = 3.36 (68.04 MB/s)
loading time = 2.61 (87.80 MB/s)

está claramente afunilado pelo desempenho do disco.

O tempo de pesquisa é lento porque, quando você usa mmap, não carrega o conteúdo da matriz na memória quando invoca o loadmétodo. Os dados são carregados com preguiça quando dados específicos são necessários. E isso acontece na pesquisa no seu caso. Mas a segunda pesquisa não será tão lenta.

Esse é um recurso interessante mmapquando você tem uma grande matriz e não precisa carregar dados inteiros na memória.

Para resolver o seu pode usar joblib, você pode despejar qualquer objeto que desejar usando joblib.dumpdois ou mais numpy arrays, veja o exemplo

firstArray = np.arange(100)
secondArray = np.arange(50)
# I will put two arrays in dictionary and save to one file
my_dict = {'first' : firstArray, 'second' : secondArray}
joblib.dump(my_dict, 'file_name.dat')