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

草祭收录于编程-ღ-技术

2021-04-28 约 914 字预计阅读 2 分钟

LMDB数据库

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

最后的效果如下：

比直接一个个文件读取快了至少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.mdb和lock.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类型有关。

参考资料

PyTorch 60.读写LMDB数据库 https://zhuanlan.zhihu.com/p/266136706
python之lmdb https://blog.csdn.net/zhayushui/article/details/103896519