目录

机器学习用lmdb给数据读取加速

读写LMDB数据库

感觉最近训练模型的时候,GPU利用率经常是间断出现0%,试了很多方法比如dataloader的多个worker,prefetch,感觉速度没达到预期,而且依然会出现0%的情况,所以使用LMDB试试能不能提升效率。

最后的效果如下:

image-20210428105853927

比直接一个个文件读取快了至少3倍,并且GPU利用率再也没出现0%的情况(看到有issue说读取多个大文件的时候可能没有提升,所以主要应用场景应该是大量读取小文件,我这里的文件大小基本上每个在50到400KB)。

Step 1. 统计数据大小

使用du -sh命令统计当前文件夹的总大小,统计结果有343G,将344(343+1)转化成B,有2954937499648B。

Step 2. LMDB数据的添加、修改、删除

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
# -*- coding: utf-8 -*-
import lmdb
 
# map_size定义最大储存容量,单位是B,以下定义344G容量
env = lmdb.open("./train", map_size=2954937499648)

txn = env.begin(write=True)

# 添加数据和键值
txn.put(key='1'.encode(), value=np.array([1,2,3]))
txn.put(key='2'.encode(), value=np.array([1,2,3]))
txn.put(key='3'.encode(), value=np.array([1,2,3]))

# 通过键值删除数据
txn.delete(key='1'.encode())

# 修改数据
txn.put(key='3'.encode(), value='ddd'.encode())

# 通过commit()函数提交更改
txn.commit()
env.close()

此时会创建目录,里面有2个文件data.mdblock.mdb.

其中map_size是指定能够存放的最大的大小,当数据超过map_size时,会报lmdb.MapFullError: mdb_put: MDB_MAP_FULL: Environment mapsize limit reached

Step 3. LMDB数据的查询

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
#查询
env = lmdb.Environment('./train') #env = lmdb.open("./train")
txn = env.begin()  #write=False
# get函数通过键值查询数据
print('打印2的值:',txn.get('2'.encode()))
print('打印3的值:',txn.get('3'.encode()))


# 通过cursor()遍历所有数据和键值
for key, value in txn.cursor():
    print (key, value)

print('样本数量:', txn.stat()['entries'])  #读取LMDB文件的样本数量

# close
env.close()

输出:

1
2
3
4
5
打印2的值: b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00'
打印3的值: b'33333'
b'2' b'\x01\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00'
b'3' b'33333'
样本数量: 2

可以看到保存的格式都是byte,对于numpy可以用np.frombuffer(image_bin, dtype=int),需要注意dtype类型和保存的numpy类型有关。

参考资料