python - 使用 numpy 加载和存储稀疏数据的性能

Question

我考虑过自定义稀疏数据格式。它的目的是为了节省空间的存储和加载，而不是对其进行计算。本质是存储非零条目的索引和值。我想知道是否有一些调整可以提高性能。

处理这样的数据就产生了这种需求:N 个“图像”(32x32)，每个图像有四个 channel 。这些图像平均包含约 5% 的非零值。随着 N 变得非常大，将所有图像存储在 RAM 中效率很低。因此，仅存储非零条目的数量、它们的索引和值以及原始形状。

以下是如何实现此功能的示例:

import numpy as np


def disassemble_data(data):
    # get some dense data and make it sparse
    lengths = np.count_nonzero(data, axis=(1, 2, 3))
    idxs = np.flatnonzero(data)
    vals = data.ravel()[idxs]

    return lengths, idxs, vals, data.shape


def assemble_data(lengths, idxs, vals, shape):
    # get some sparse data and make it dense
    data = np.zeros(shape)

    lower_idx = 0
    for length in lengths:
        upper_idx = lower_idx + length
        data.ravel()[idxs[lower_idx:upper_idx]] = vals[lower_idx:upper_idx]
        lower_idx = upper_idx
    return data


# Create some dummy data 
my_data = np.random.uniform(0, 1, (10, 4, 32, 32))
my_data[my_data > 0.05] = 0

# make data sparse and then dense again
my_new_data = assemble_data(*disassemble_data(my_data))

# assert that this actually works
assert np.allclose(my_data, my_new_data)

现在，我们可以直接看到优点:数据是逐幅致密化的。这使我们能够将整个数据集加载到 RAM 中，并通过生成器按需生成密集图像:

def image_generator(lengths, idxs, vals, shape):          
        idxs %= np.prod(shape[1:])

        lower_idx = 0
        for length in lengths:
            upper_idx = lower_idx + length
            data = np.zeros(shape[1:])
            data.ravel()[idxs[lower_idx:upper_idx]] = vals[lower_idx:upper_idx]
            lower_idx = upper_idx
            yield data

此外，还可以生成批量图像:

def image_batch_generator(lengths, idxs, vals, shape, batch_size):
    idxs %= np.prod((batch_size, *shape[1:]))
    lengths = np.sum(lengths.reshape(-1, batch_size), axis=1)

    lower_idx = 0
    for length in lengths:
        upper_idx = lower_idx + length
        data = np.zeros((batch_size, *shape[1:]))
        data.ravel()[idxs[lower_idx:upper_idx]] = vals[lower_idx:upper_idx]
        lower_idx = upper_idx
        yield data

对于我的需求来说，这是一个相当方便的方法。但我想知道是否有可能加快速度。

例如我看到 numpys 项集比直接赋值更快(根据 docs )。但它仅适用于单个项目，不适用于索引数组。

还有其他方法吗？我对 cython 等一点也不熟悉，所以我很乐意得到一些提示!

Answer 1

我测试了一些如何更有效地完成此操作，并发现对于 np.random.uniform 生成的高度不相关的数据，您的方法一点也不差。根据实际数据，这可能有点不同。

我稍微提高了你的函数的速度，压缩速度为 1.4GB/s，解压缩速度为 1.2GB/s，这还不错。使用 h5py (blosclz)，我只能达到大约 450MB/s，但也将数据写入磁盘。

改进的稀疏算法

import numpy as np
import numba as nb

#We can use uint16 on (4,32,32), since max. idx<2**16
@nb.jit()
def to_sparse_data_uint16(data):
  data_flat=data.reshape(-1)

  idx=np.empty(data.size,dtype=np.uint16)
  data_out=np.empty(data.size,dtype=data.dtype)

  ii=0
  for i in range(data_flat.shape[0]):
    if (data_flat[i]!=0):
      idx[ii]=i
      data_out[ii]=data_flat[i]
      ii+=1

  return idx[0:ii], data_out[0:ii], data.shape


def to_dense_data(idx,data,shape):
  length=np.prod(shape)
  data_out=np.zeros(length,dtype=data.dtype)

  data_out[idx]=data

  return data_out.reshape(shape)


########################
#do you really need float64 here?
images = np.array(np.random.uniform(0, 1, (100000, 4, 32, 32)),dtype=np.float32)
images[images > 0.05] = 0.

res=[]
t1=time.time()
for i in range(100000):
  res.append(to_sparse_data_uint16(images[i,:,:,:]))

print(time.time()-t1)

t1=time.time()
for i in range(100000):
  data=to_dense_data(res[i][0],res[i][1],res[i][2])

print(time.time()-t1)

HDF5 示例

import numpy as np
import tables #register blosc
import h5py as h5
import h5py_cache as h5c
import time

# I assume here that you don't need float64 for images..
# 1650MB Testdata

images = np.array(np.random.uniform(0, 1, (100000, 4, 32, 32)),dtype=np.float32)
images[images > 0.05] = 0.

#Write data (32,7 GB uncompressed)
hdf5_path='Test.h5'
f = h5c.File(hdf5_path, 'w',chunk_cache_mem_size=1024**2*100) #200 MB cache size
dset_images = f.create_dataset("images", shape=(20*100000, 4, 32, 32),dtype=np.float32,chunks=(1000, 4, 32, 32),compression=32001,compression_opts=(0, 0, 0, 0, 9, 1, 1), shuffle=False)

t1=time.time()
#Don't call h5py to often, this will lead to bad performance
for i in range(20):
    dset_images[i*100000:(i+1)*100000,:,:,:]=images

f.close()
print(time.time()-t1)

print("MB/s: " + str(32700/(time.time()-t1)))