python将矩阵存为lmdb文件

Posted 2021-10-13 白水baishui

由于工作需要，将C++生成的矩阵存入LMDB再用caffe进行处理，输出的矩阵失去了它原本的shape，因此只能记录下来：
矩阵X（n*ell）:

-372302407
1319544887
-223830618
-184109131
-328009648
-182855917
…
749988319
-1387505086
-1967883425
1022949309
1680638116
1541327498

矩阵Y（r*n）:

1
0
0
0
11
-34
-280
109
151
-205
-105
…
156
-152
-151
27
199
-202
-125
-95

# 向lmdb数据库存入两个矩阵X和y，大小分别为n*ell、r*n

import numpy as np
import lmdb

# X:r*n
# Y:n*ell

r = 4
n = 32
ell = 30

X_path = './X.txt'
Y_path = './Y.txt'

def write_lmdb(filename):
    print('Write lmdb')
    lmdb_env = lmdb.open(filename, map_size=int(1e9))
    n_samples= 1
    # 读入X
    with open(X_path, 'r') as f:
        # 读取图像文件的二进制格式数据
        X_temp = f.read()
        X_temp = X_temp.split('\\n')
        del (X_temp[-1])
        X = [int(i) for i in X_temp]
        X = np.array(X).astype(np.int64)
        # print(X)
    # 读入Y
    with open(Y_path, 'r') as f:
        # 读取图像文件的二进制格式数据
        Y_temp = f.read()
        Y_temp = Y_temp.split('\\n')
        del (Y_temp[-1])
        y = [int(i) for i in Y_temp]
        y = np.array(y).astype(np.int64)
        # print(y)

    # 写入LMDB
    for i in range(n_samples):
        with lmdb_env.begin(write=True) as lmdb_txn:
            lmdb_txn.put(str('X_'+str(i)).encode(), X)
            lmdb_txn.put(str('y_'+str(i)).encode(), y)
            print('X:',X)
            print('y:',y)

def read_lmdb(filename):
    print('Read lmdb')

    lmdb_env = lmdb.open(filename)
    lmdb_txn = lmdb_env.begin()
    lmdb_cursor = lmdb_txn.cursor()

    n_samples=1
    read_array = []
    with lmdb_env.begin() as lmdb_txn:
        with lmdb_txn.cursor() as lmdb_cursor:
            for key, value in lmdb_cursor:
                if('X'.encode() in key):
                    read_temp = np.frombuffer(value, dtype=np.int64)
                    read_array.append(read_temp)
                if('y'.encode() in key):
                    read_temp = np.frombuffer(value, dtype=np.int64)
                    read_array.append(read_temp)
                n_samples = n_samples + 1

    x_read_temp = read_array[0]
    x_read = []
    x_read_line = []
    print("X: ")
    for i in range(x_read_temp.size):
        x_read_line.append(x_read_temp[i])
        if len(x_read_line) == (n+r)*ell: # 列数
            x_read.append(x_read_line.copy())
            x_read_line.clear()
    print(x_read)

    y_read_temp = read_array[1]
    y_read = []
    y_read_line = []
    print("Y: ")
    for i in range(y_read_temp.size):
        y_read_line.append(y_read_temp[i])
        if len(y_read_line) == n+r: # 列数
            y_read.append(y_read_line.copy())
            y_read_line.clear()
    print(y_read)

    print('n_samples',n_samples)

write_lmdb('temp.db')
read_lmdb('temp.db')

输出如下：

Write lmdb
X: [-372302407 1319544887 -223830618 ... 1998175905 -616375716 -132621858]
y: [   1    0    0    0   11  -34 -280  109  151 -205 -105    3  156 -152
 -151   27  199 -202 -125  -95  318  -88 -173  273  262 -127 -152 -322
  -18  198  110   65   14 -310  192  103    0    1    0    0 -215 -217
    0  -65 -123 -110  171   39  -78  -21   24  -61   23  -44  113  184
  169 -240 -157  -16  136  182 -331  156   26  112 -188 -175  244 -195
   30 -228    0    0    1    0  -55  -15  -62  -41  -24 -140  -94  198
 -130   -4  -33 -206  256 -230 -406 -170   84  -45 -283 -136   38  -39
   64  -84   85  -85   41 -241  -42 -107  -19 -135    0    0    0    1
  124    5 -235 -121   49 -115  242  -61 -353  106  353 -210 -194  219
 -228  183  -80  110  218 -243   10 -190 -338  -75   11 -177   31 -208
  176  -66   97 -134]

Read lmdb(由于数据太长，省略了后面的)
X: 
[[-372302407, 1319544887, -223830618, -184109131, ...], ...]
Y: 
[[1, 0, 0, 0, 11, -34, -280, 109, 151, -205, ...], ...]