深度学习训练前标准的LMDB文件(data.mdb和lock.mdb)制作（附代码）

一、LMDB文件

在训练的时候使用LMDB 存储形式可以加快IO 和CPU 解压缩的速度(测试的时候数据较少, 一般就没有太必要使用LMDB)。其具体的加速要根据机器的配置来，以下几个因素会影响：

有的机器设置了定时清理缓存，而LMDB 依赖于缓存。因此若一直缓存不进去，则需要检查一下。一般free -h 命令下, LMDB 占用的缓存会记录在buff/cache 条目下面。
机器的内存是否足够大，能够把整个LMDB 数据都放进去。如果不是，则它由于需要不断更换缓存，会影响速度。
若是第一次缓存LMDB 数据集，可能会影响训练速度。可以在训练前，进入LMDB 数据集
目录，把数据先缓存进去：cat data.mdb > /dev/null

二、准备训练集

训练集的种类及格式是多样的，这里主要以.png格式的图片进行讲解。如下：

在这里插入图片描述

三、安装basicsr包

进入终端后输入下面命令安装basicsr包：

pip install basicsr -i https://pypi.mirrors.ustc.edu.cn/simple/

四、LMDB文件制作

4.1 参数修改

在这里插入图片描述

4.2 其它格式图片修改

下面是数据集中，图片后缀为其它格式时，怎么修改代码，这里以.bmp格式图片为例：

在这里插入图片描述

4.3 代码

转换的详细代码见下：

import cv2
from basicsr.utils import scandir
import lmdb
import sys
from multiprocessing import Pool
from os import path as osp
from tqdm import tqdm


def create_lmdb_for_div2k():
    folder_path = 'Images/LMDB/Visible_Images'
    lmdb_path = 'Images/LMDB/Visible_Images.lmdb'
    img_path_list, keys = prepare_keys_div2k(folder_path)
    make_lmdb_from_imgs_lmdb(folder_path, lmdb_path, img_path_list, keys)

def prepare_keys_div2k(folder_path):
    print('Reading image path list ...')
    img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=False)))   # 获取文件夹下所有后缀为 png 的文件路径，但不进行递归搜索
    keys = [img_path.split('.png')[0] for img_path in sorted(img_path_list)]        # 从图像路径中分离出键，去掉 '.png' 后缀并排序

    return img_path_list, keys

def prepare_keys_div2k_bmp(folder_path):
    print('Reading image path list ...')
    img_path_list = sorted(list(scandir(folder_path, suffix='bmp', recursive=False)))
    keys = [img_path.split('.bmp')[0] for img_path in sorted(img_path_list)]

    return img_path_list, keys


def read_img_worker(path, key, compress_level):
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    if img.ndim == 2:
        h, w = img.shape
        c = 1
    else:
        h, w, c = img.shape
    _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
    return (key, img_byte, (h, w, c))

def make_lmdb_from_imgs_lmdb(data_path,
                        lmdb_path,
                        img_path_list,
                        keys,
                        batch=5000,
                        compress_level=1,
                        multiprocessing_read=False,
                        n_thread=40,
                        map_size=None):


    assert len(img_path_list) == len(keys), ('img_path_list and keys should have the same length, '
                                             f'but got {len(img_path_list)} and {len(keys)}')
    print(f'Create lmdb for {data_path}, save to {lmdb_path}...')
    print(f'Totoal images: {len(img_path_list)}')
    if not lmdb_path.endswith('.lmdb'):
        raise ValueError("lmdb_path must end with '.lmdb'.")
    if osp.exists(lmdb_path):
        print(f'Folder {lmdb_path} already exists. Exit.')
        sys.exit(1)

    if multiprocessing_read:
        # read all the images to memory (multiprocessing)
        dataset = {}  # use dict to keep the order for multiprocessing
        shapes = {}
        print(f'Read images with multiprocessing, #thread: {n_thread} ...')
        pbar = tqdm(total=len(img_path_list), unit='image')

        def callback(arg):
            """get the image data and update pbar."""
            key, dataset[key], shapes[key] = arg
            pbar.update(1)
            pbar.set_description(f'Read {key}')

        pool = Pool(n_thread)
        for path, key in zip(img_path_list, keys):
            pool.apply_async(read_img_worker, args=(osp.join(data_path, path), key, compress_level), callback=callback)
        pool.close()
        pool.join()
        pbar.close()
        print(f'Finish reading {len(img_path_list)} images.')

    # create lmdb environment
    if map_size is None:
        # obtain data size for one image
        img = cv2.imread(osp.join(data_path, img_path_list[0]), cv2.IMREAD_UNCHANGED)
        _, img_byte = cv2.imencode('.png', img, [cv2.IMWRITE_PNG_COMPRESSION, compress_level])
        data_size_per_img = img_byte.nbytes
        print('Data size per image is: ', data_size_per_img)
        data_size = data_size_per_img * len(img_path_list)
        map_size = data_size * 10

    env = lmdb.open(lmdb_path, map_size=map_size)

    # write data to lmdb
    pbar = tqdm(total=len(img_path_list), unit='chunk')
    txn = env.begin(write=True)
    txt_file = open(osp.join(lmdb_path, 'meta_info.txt'), 'w')
    for idx, (path, key) in enumerate(zip(img_path_list, keys)):
        pbar.update(1)
        pbar.set_description(f'Write {key}')
        key_byte = key.encode('ascii')
        if multiprocessing_read:
            img_byte = dataset[key]
            h, w, c = shapes[key]
        else:
            _, img_byte, img_shape = read_img_worker(osp.join(data_path, path), key, compress_level)
            h, w, c = img_shape

        txn.put(key_byte, img_byte)
        # write meta information
        txt_file.write(f'{key}.png ({h},{w},{c}) {compress_level}\n')
        if idx % batch == 0:
            txn.commit()
            txn = env.begin(write=True)
    pbar.close()
    txn.commit()
    env.close()
    txt_file.close()
    print('\nFinish writing lmdb.')


if __name__ == '__main__':
    create_lmdb_for_div2k()

4.4 转换结果

运行上面脚本，输出如下：

在这里插入图片描述
最终的LMDB文件在代码中设置的路径下：

在这里插入图片描述
生成的各个文件解析加下。

4.4.1 data.mdb文件

存储数据库中的所有实际数据，包括键值对、元数据等。

是一个二进制文件，可以直接使用内存映射的方式进行读写，访问速度非常快。

文件大小取决于数据库中存储的数据量。

4.4.2 lock.mdb文件

一个用于控制数据库访问的锁文件。

确保同一时间只有一个进程可以对数据库进行读写操作，防止数据损坏。

文件大小很小，通常只有几百字节。

4.4.3 meta_info.txt文件

采用txt 来记录，是为了可读性，文件中内容如下：

在这里插入图片描述
上面每一行记录了一张图片，有三个字段，分别表示：

图像名称(带后缀): 0001.png
图像大小：(1404,2040,3) 表示是1404 × 2040 × 3的图像
其他参数(BasicSR 里面使用了cv2 压缩png 程度)：因为在复原任务中，通常使用png 来存储，所以这个1 表示png 的压缩程度，也就是CV_IMWRITE_PNG_COMPRESSION 为1。CV_IMWRITE_PNG_COMPRESSION可以取值为[0, 9] 的整数，更大的值表示更强的压缩，即更小的储存空间和更长的压缩时间。

五、总结

以上就是深度学习训练前标准的LMDB文件(data.mdb和lock.mdb)制作过程，希望能帮我你，有问题欢迎留言。

感谢您阅读到最后！关注公众号「视觉研坊」，获取干货教程、实战案例、技术解答、行业资讯！