背景#

在深度学习的时候，如果你的batch size调的很大，或者你每次获取一个batch需要许多的预操作，那么pytorch的Dataloader获取一个batch就会花费较多的时间，那么训练的时候就会出现GPU等CPU的情况，训练的效率就会下降。

为了应对这种情况，Tensorflow有TFrecord，但是Pytorch没有对应的数据格式，在查询各类资料之后，我决定使用LMDB这个数据库

LMDB是一种数据库，可以实现多进程访问，访问简单，而且不需要把全部文件读入内存，总而言之就是速度很快

方法#

我们想用LMDB来读取的话，首先就需要将我们原始的数据集转换为LMDB的数据格式，然后在训练的时候读取这个文件就行了。首先我们来实现将原始数据集转换为LMDB的过程：

转换为LMDB格式#

相必大家在寻找一个更加快速的Dataloader的时候，已经写好了Pytorch常规的Dataloader，我们这里就可以利用上这个已有的Dataloader

首先打开一个lmdb数据库文件，如果之前有用过其他文件数据库的话，会发现这有点相似~

db = lmdb.open(lmdb_path, subdir=isdir,

               map_size=1099511627776 * 2, readonly=False,

               meminit=False, map_async=True)

然后准备向其中写入数据：

txn = db.begin(write=True)

然后就是将图片放到lmdb中

txn.put(u'{}'.format(idx).encode('ascii'), dumps_pyarrow(

    (pic.numpy())

))

txn.commit()

总体的操作与数据库一样，总体就是打开数据库，放入数据，最后commit

最后放出完整的方法：

def data2lmdb(dataloader, db_path='data_lmdb', name="train", write_frequency=50):

    """

    Args:

        dataloader: the general dataloader of the dataset, e.g torch.utils.data.DataLoader

        db_path: the path you want to save the lmdb file

        name: train or test

        write_frequency: Write once every ? rounds

    Returns:

        None

    """

    if not os.path.exists(db_path):

        os.makedirs(db_path)

    lmdb_path = os.path.join(db_path, "%s.lmdb" % name)

    isdir = os.path.isdir(lmdb_path)

    print("Generate LMDB to %s" % lmdb_path)

    db = lmdb.open(lmdb_path, subdir=isdir,

                   map_size=1099511627776 * 2, readonly=False,

                   meminit=False, map_async=True)

    txn = db.begin(write=True)

    for idx, data in enumerate(dataloader):

        # get data from dataloader

        if name == 'train':

            pic = data

        elif name == 'test':

            pic = data

        else:

            raise 'unexpect name :{}'.format(name)

        # put data to lmdb dataset

        # {idx, (in_LDRs, in_HDRs, ref_HDRs)}

        txn.put(u'{}'.format(idx).encode('ascii'), dumps_pyarrow(

            (pic.numpy())

        ))

        if idx % write_frequency == 0:

            print("[%d/%d]" % (idx, len(dataloader)))

            txn.commit()

            txn = db.begin(write=True)

    # finish iterating through dataset

    txn.commit()

    keys = [u'{}'.format(k).encode('ascii') for k in range(idx + 1)]

    with db.begin(write=True) as txn:

        txn.put(b'__keys__', dumps_pyarrow(keys))

        txn.put(b'__len__', dumps_pyarrow(len(keys)))

    print("Flushing database ...")

    db.sync()

    db.close()

上面的函数适合我的使用场景，如果你需要搬去用的话，需要修改dataloader输出的地方，根据自己的dataloader来写

读取LMDB文件#

这里放上IMDB的Dataloader

class ImageFolderLMDBTest(data.Dataset):

    def __init__(self, db_path):

        self.db_path = db_path

        self.env = lmdb.open(db_path, subdir=os.path.isdir(db_path),

                             readonly=True, lock=False,

                             readahead=False, meminit=False)

        with self.env.begin(write=False) as txn:

            # self.length = txn.stat()['entries'] - 1

            self.length = pa.deserialize(txn.get(b'__len__'))

            self.keys = pa.deserialize(txn.get(b'__keys__'))

    def __getitem__(self, index):

        env = self.env

        with env.begin(write=False) as txn:

            byteflow = txn.get(self.keys[index])

        unpacked = pa.deserialize(byteflow)

        # load image

        # 这里写你写入lmdb时的数据，上面的我写入了pic，这里展开就还是pic

        pic = unpacked

        return pic

    def to_tensor(self, img):

        img_t = torch.from_numpy(img.copy())

        if isinstance(img_t, torch.ByteTensor):

            return img_t.float().div(255)

        else:

            return img_t

    def __len__(self):

        return self.length

    def __repr__(self):

        return self.__class__.__name__ + ' (' + self.db_path + ')'

这个Dataloader的结构和一般的Dataloader一样，需要注意的我也写在了注释里面，其实就是根据你写入的东西不同，从LMDB里取出的东西也不一样

后记#

在用上这个LMDB文件格式之后，数据的读取速度也快了很多，GPU也终于不会歇着了