PyTorch下训练数据小文件转大文件读写(附有各种存储格式对比)

引言

• Tensorflow有着专门的数据读取模块tfrecord，可以高效地读取训练神经网络模型所用的数据，充分喂饱GPU
• Caffe用lmdb来读取数据，也可以很高效地去读取
• PyTorch有DataLoader读取数据，但是速度比较慢，尤其是小文件较多情况下
• 如何基于PyTorch，高效读取数据，充分利用GPU性能，成为一个关键问题？

TFRecord

• 是否可以将tensorflow下的tfrecord借来一用？未尝不可
• 目前已经有伙伴实现了，详情参见：tfrecord
• 同时，在Kaggle上，也有大神手动实现，详情参见：PyTorch TFRecord-Loader

tfrecord写入代码：

	import cv2
	import numpy as np
	import tensorflow as tf
	from tqdm import tqdm
	 
	from data_loader import TFRecordDataLoader
	 
	 
	def read_txt(txt_path):
	    with open(txt_path, 'r', encoding='utf-8') as f:
	        data = f.readlines()
	    data = list(map(lambda x: x.rstrip('\n'), data))
	    return data
	 
	 
	def bytes_to_numpy(image_bytes):
	    image_np = np.frombuffer(image_bytes, dtype=np.uint8)
	    image_np2 = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
	    return image_np2  
	 
	def list_record_features(tfrecords_path):
	    """查看tfrecords结构
	    https://stackoverflow.com/questions/63562691/reading-a-tfrecord-file-where-features-that-were-used-to-encode-is-not-known
	 
	    Args:
	        tfrecords_path (str): tfrecords路径
	 
	    Returns:
	        dict: 结构信息
	    """
	    features = {}
	    dataset = tf.data.TFRecordDataset([str(tfrecords_path)])
	    data = next(iter(dataset))
	 
	    example = tf.train.Example()
	    example_bytes = data.numpy()
	    example.ParseFromString(example_bytes)
	 
	    for key, value in example.features.feature.items():
	        kind = value.WhichOneof('kind')
	        size = len(getattr(value, kind).value)
	        if key in features:
	            kind2, size2 = features[key]
	            if kind != kind2:
	                kind = None
	 
	            if size != size2:
	                size = None
	        features[key] = (kind, size)
	    return features
	 
	class TFRecorder(object):
	    def __init__(self) -> None:
	        super().__init__()
	        self.feature_dict = {
	            'height': None,
	            'width': None,
	            'depth': None,
	            'label': None,
	            'image_raw': None
	        }
	        self.AUTO = tf.data.experimental.AUTOTUNE
	 
	    def image_to_feature(self, image_string, label):
	        height, width, channel = tf.image.decode_image(image_string).shape
	        self.feature_dict = {
	            'height': self._int64_feature(height),
	            'width': self._int64_feature(width),
	            'depth': self._int64_feature(channel),
	            'label': self._int64_feature(label),
	            'image_raw': self._bytes_feature(image_string)
	        }
	        return tf.train.Example(features=tf.train.Features(feature=self.feature_dict))
	 
	    def write(self, save_path, img_label_dict):
	        with tf.io.TFRecordWriter(save_path) as writer:
	            for file_name, label in tqdm(img_label_dict.items()):
	                img_string = open(file_name, 'rb').read()
	                feature = self.image_to_feature(img_string, label)
	                writer.write(feature.SerializeToString())
	 
	    def read(self, tfrecord_path):
	        reader = tf.data.TFRecordDataset(tfrecord_path)
	        dataset = reader.map(self._parse_image_function,
	                             num_parallel_calls=self.AUTO)
	        return dataset
	 
	    def _parse_image_function(self, example_proto):
	        self.feature_dict = {
	            'height': tf.io.FixedLenFeature([], tf.int64),
	            'width': tf.io.FixedLenFeature([], tf.int64),
	            'depth': tf.io.FixedLenFeature([], tf.int64),
	            'label': tf.io.FixedLenFeature([], tf.int64),
	            'image_raw': tf.io.FixedLenFeature([], tf.string)
	        }
	        example = tf.io.parse_single_example(example_proto,
	                                             self.feature_dict)
	        return example
	 
	    @staticmethod
	    def _bytes_feature(value):
	        """Returns a bytes_list from a string / byte."""
	        if isinstance(value, type(tf.constant(0))):
	            # BytesList won't unpack a string from an EagerTensor.
	            value = value.numpy()
	        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
	 
	    @staticmethod
	    def _float_feature(value):
	        """Returns a float_list from a float / double."""
	        return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
	 
	    @staticmethod
	    def _int64_feature(value):
	        """Returns an int64_list from a bool / enum / int / uint."""
	        return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
	 
	 
	if __name__ == '__main__':
	    tfrecorder = TFRecorder()
	 
	    # val.txt中存放的是图像的相对路径
	    img_path = read_txt('dataset/val.txt')
	 
	    # Path(v).parent.name: 图像的标签
	    img_label_dict = {v: int(Path(v).parent.name) for v in img_path}
	 
	    save_path = 'temp/val.tfrecords'
	    tfrecorder.write(save_path, img_label_dict)
	 
	    dataset = tfrecorder.read('dataset/val.tfrecords')
	    for v in dataset:
	        img, label = v
	        print('ok')
	 
	    # 查看未知tfrecords结构信息
	   list_record_features('xxxx.tfrecords')

基于PyTorch下tfrecord读取代码

import cv2
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds

AUTO = tf.data.experimental.AUTOTUNE


def bytes_to_numpy(image_bytes):
    image_np = np.frombuffer(image_bytes, dtype=np.uint8)
    image_np2 = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    return image_np2


def read_labeled_tfrecord(example_proto):
    feature_dict = {
        'height': tf.io.FixedLenFeature([], tf.int64),
        'width': tf.io.FixedLenFeature([], tf.int64),
        'depth': tf.io.FixedLenFeature([], tf.int64),
        'label': tf.io.FixedLenFeature([], tf.int64),
        'image_raw': tf.io.FixedLenFeature([], tf.string)
    }
    example = tf.io.parse_single_example(example_proto,
                                         feature_dict)
    img = tf.io.decode_image(example['image_raw'], channels=3,
                             expand_animations=False)
    img = tf.image.resize_with_crop_or_pad(img,
                                           target_height=388,
                                           target_width=270)
    return img, example['label']


def get_dataset(files, batch_size=16, repeat=False,
                cache=False, shuffle=False):
    ds = tf.data.TFRecordDataset(files, num_parallel_reads=AUTO)
    if cache:
        ds = ds.cache()

    if repeat:
        ds = ds.repeat()

    if shuffle:
        ds = ds.shuffle(1024 * 2)
        opt = tf.data.Options()
        opt.experimental_deterministic = False
        ds = ds.with_options(opt)

    ds = ds.map(read_labeled_tfrecord, num_parallel_calls=AUTO)
    ds = ds.batch(batch_size)
    ds = ds.prefetch(AUTO)
    return tfds.as_numpy(ds)


def count_data_items(file):
    num_ds = tf.data.TFRecordDataset(file, num_parallel_reads=AUTO)
    num_ds = num_ds.map(read_labeled_tfrecord, num_parallel_calls=AUTO)
    num_ds = num_ds.repeat(1)
    num_ds = num_ds.batch(1)

    c = 0
    for _ in num_ds:
        c += 1
    del num_ds
    return c


class TFRecordDataLoader:
    def __init__(self, files, batch_size=32, cache=False, train=True,
                 repeat=False, shuffle=False, labeled=True,
                 return_image_ids=True):
        self.ds = get_dataset(
            files,
            batch_size=batch_size,
            cache=cache,
            repeat=repeat,
            shuffle=shuffle,)

        if train:
            self.num_examples = count_data_items(files)

        self.batch_size = batch_size
        self.labeled = labeled
        self.return_image_ids = return_image_ids
        self._iterator = None

    def __iter__(self):
        if self._iterator is None:
            self._iterator = iter(self.ds)
        else:
            self._reset()
        return self._iterator

    def _reset(self):
        self._iterator = iter(self.ds)

    def __next__(self):
        batch = next(self._iterator)
        return batch

    def __len__(self):
        n_batches = self.num_examples // self.batch_size
        if self.num_examples % self.batch_size == 0:
            return n_batches
        else:
            return n_batches + 1

# 使用
train_txt_path = 'dataset/minist/train.tfrecords'
train_dataloader = TFRecordDataLoader(train_txt_path,
                                      batch_size=batch_size,
                                      shuffle=True)
for v in train_dataloader:
    pass

LMDB

• 纵观各大论坛，说到基于PyTorch下提高小文件读取速度，不得不说到LMDB(Lightning Memory-Mapped Database)了，我也做了一些尝试，最终结论将在最后给出

写入LMDB

import os
import pickle
from pathlib import Path
 
import cv2
import lmdb
import numpy as np
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from tqdm import tqdm
 
import utils
 
 
class SimpleDataset(Dataset):
    def __init__(self, txt_path, transform=None) -> None:
        self.img_paths = utils.read_txt(txt_path)
        self.transform = transform
 
    def __getitem__(self, index: int):
        img_path = self.img_paths[index]
        label = int(Path(img_path).parent.name)
        try:
            img = Image.open(img_path)
            img = img.convert('RGB')
        except:
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(img)
 
        if self.transform:
            img = self.transform(img)
        img = np.array(img)
        return img, label
 
    def __len__(self) -> int:
        return len(self.img_paths)
 
 
class LMDB_Image:
    def __init__(self, image, label):
        # Dimensions of image for reconstruction - not really necessary
        # for this dataset, but some datasets may include images of
        # varying sizes
        self.channels = image.shape[2]
        self.size = image.shape[:2]
 
        self.image = image.tobytes()
        self.label = label
 
    def get_image(self):
        """ Returns the image as a numpy array. """
        image = np.frombuffer(self.image, dtype=np.uint8)
        return image.reshape(*self.size, self.channels)
 
 
def data2lmdb(dpath, name="train", txt_path=None,
              write_frequency=10, num_workers=4):
    dataset = SimpleDataset(txt_path=txt_path)
    data_loader = DataLoader(dataset, num_workers=num_workers,
                             collate_fn=lambda x: x)
 
    lmdb_path = os.path.join(dpath, "%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,  # 单位byte
                   readonly=False,
                   meminit=False,
                   map_async=True)
 
    txn = db.begin(write=True)
    for idx, data in enumerate(tqdm(data_loader)):
        image, label = data[0]
        temp = LMDB_Image(image, label)
        txn.put(u'{}'.format(idx).encode('ascii'), pickle.dumps(temp))
 
        if idx % write_frequency == 0:
            print("[%d/%d]" % (idx, len(data_loader)))
            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__', pickle.dumps(keys))
        txn.put(b'__len__', pickle.dumps(len(keys)))
 
    print("Flushing database ...")
    db.sync()
    db.close()
 
if __name__ == '__main__':
    save_dir = 'dataset/minist'
    data2lmdb(save_dir, name='val', txt_path='dataset/minist/val.txt')

读取LMDB

class DatasetLMDB(Dataset):
    def __init__(self, db_path, transform=None):
        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() as txn:
            self.length = pickle.loads(txn.get(b'__len__'))
            self.keys = pickle.loads(txn.get(b'__keys__'))
        self.transform = transform
 
    def __getitem__(self, index):
        with self.env.begin() as txn:
            byteflow = txn.get(self.keys[index])
 
        IMAGE = pickle.loads(byteflow)
        img, label = IMAGE.get_image(), IMAGE.label
        return Image.fromarray(img).convert('RGB'), label
 
    def __len__(self):
        return self.length
 
# 使用
train_transforms = transforms.Compose([
    transforms.Resize((388, 270)),
    transforms.RandomChoice([
        transforms.RandomRotation(10),
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomGrayscale(p=0.3),
        transforms.RandomPerspective(distortion_scale=0.6, p=0.5),
        transforms.ColorJitter(brightness=.5, hue=.3),
    ]),
    transforms.ToTensor(),
    normalize,
    transforms.RandomErasing(),
    ])
 
train_dataset = DatasetLMDB(train_txt_path, train_transforms)
train_dataloader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=n_worker,
                              pin_memory=True)
# do other things

二进制大文件

• 直接将现有数据集按照二进制读取，存入一个bins的大文件中，也不失为一种选择

写入bins

import cv2
import numpy as np
from tqdm import tqdm
 
 
def write_bin(save_bin_path, save_index_path, data):
    """将现有基于文件的数据集写为bin大文件
    写入到save_index_path中的索引位置和标签，中间以\t分割
 
    Args:
        save_bin_path (str): 保存bin的位置
        save_index_path (str): 保存bin中索引和对应标签
        data (str): 存放图像路径和对应标签的list,
                    e.g. [['xxx/1.jpg', 'cat'], ['xxx/2.jpg', 'dog']]
    """
    with open(save_bin_path, 'wb') as f_w, \
            open(save_index_path, 'w') as f_index:
        start_index = 0
 
        for img_path, label in tqdm(data):
            with open(img_path, 'rb') as f:
                img_bin = f.read()
 
            f_w.write(img_bin)
 
            len_bin = len(img_bin)
            f_index.write(f'{start_index}\t{len_bin}\t{label}\n')
 
            start_index += len_bin
 
 
def read_bin(bin_path, index_path):
    """读取bin大文件和对应的索引标签txt
 
    Args:
        bin_path (str): bin大文件存放路径
        index_path (str): 索引和标签存放txt的路径
    """
    with open(bin_path, 'rb') as f_bin, open(index_path, 'r') as f_index:
        index_lines = list(map(lambda x: x.strip(), f_index.readlines()))
        index_lines = list(map(lambda x: x.split('\t'), index_lines))
 
        for i, (start_index, length) in enumerate(index_lines):
            start_index = int(start_index)
            length = int(length.strip())
             
            # 定位到当前指针位置到start_index
            f_bin.seek(start_index)
 
            # 读取length的字节值
            img_bytes = f_bin.read(length)
 
            img = np.frombuffer(img_bytes, dtype='uint8')
            img = cv2.imdecode(img, -1)  # -1: cv.IMREAD_UNCHANGED
 
        # 转为PIL
            # img = Image.fromarray(img)
            # img = img.convert('RGB')
 
            # 保存图像
            # cv2.imwrite(f'temp/images/{i}.jpg', img)

读取bins

from io import BytesIO
from PIL import Image
import cv2
import numpy as np
 
class SimpleDataset(Dataset):
    def __init__(self, txt_path, bin_path, transform=None) -> None:
        self.index_info = utils.read_txt(txt_path)
        self.index_info = list(map(lambda x: x.split('\t'), self.index_info))
 
        self.f_bin = open(bin_path, 'rb')
        self.transform = transform
 
    def __getitem__(self, index: int):
        start_index, length, label = list(map(int, self.index_info[index]))
        print(start_index)
 
        self.f_bin.seek(start_index)
        img_bytes = self.f_bin.read(length)
         
        # 方案一：
        img = np.frombuffer(img_bytes, dtype='uint8')
        img = cv2.imdecode(img, -1)
        if img is None:
            return self.__getitem__(random.randint(0, self.__len__() - 1))
 
        img = Image.fromarray(img)
        img = img.convert('RGB')
 
        # 方案二：
        try:
            img = Image.open(BytesIO(img_bytes))
            img = img.convert('RGB')
        except:
            return self.__getitem__(random.randint(0, self.__len__() - 1))
 
        if self.transform:
            img = self.transform(img)
        return img, label
 
    def __len__(self) -> int:
        return len(self.index_info)

Sqlite

• 采用python内置的sqlite3作为存储格式，也是一种好的选择

写入到sqlite数据库中

import sqlite3
from pathlib import Path
 
from tqdm import tqdm
 
 
def read_txt(txt_path):
    with open(txt_path, 'r', encoding='utf-8-sig') as f:
        data = list(map(lambda x: x.rstrip('\n'), f))
    return data
 
 
def img_to_bytes(img_path):
    with open(img_path, 'rb') as f:
        img_bytes = f.read()
        return img_bytes
 
 
class SQLiteWriter(object):
    def __init__(self, db_path):
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()
 
    def execute(self, sql, value=None):
        if value:
            self.cursor.execute(sql, value)
        else:
            self.cursor.execute(sql)
 
    def __enter__(self):
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cursor.close()
        self.conn.commit()
        self.conn.close()
 
 
if __name__ == '__main__':
    dataset_dir = Path('datasets/minist')
 
    save_db_dir = dataset_dir / 'sqlite'
    save_db_path = str(save_db_dir / 'val.db')
 
    # val.txt中 每行为：图像路径\t对应文本值 e.g. xxxx.jpg\txxxxxx
    img_paths = read_txt(str(dataset_dir / 'val.txt'))
 
    with SQLiteWriter(save_db_path) as db_writer:
        # 创建表
        table_name = 'minist'
  
        # 注意这里的表中字段，要根据自己数据集来定义
        # 具体数据库类型,可参考：https://docs.python.org/zh-cn/3/library/sqlite3.html#sqlite-and-python-types
        # demo中示例所涉及到的数据集为文本识别数据集，样本为图像，标签为对应文本，
        # 下面示例字段的数据类型为python下的数据类型，只需转为以下对应数据类型即可写入数据库的表中
        # e.g. img_path: str(xxxx.jpg), img_data: bytes格式的图像数据， img_label: str(xxxxx)
        create_table_sql = f'create table {table_name} (img_path TEXT primary key, img_data BLOB, img_label TEXT)'
        db_writer.execute(create_table_sql)
 
        # 向表中插入数据，value部分采用占位符
        insert_sql = f'insert into {table_name} (img_path, img_data, img_label) values(?, ?, ?)'
        for img_info in tqdm(img_paths):
            img_path, label = img_info.split('\t')
 
            img_full_path = str(dataset_dir / 'images' / img_path)
            img_data = img_to_bytes(img_full_path)
 
            db_writer.execute(insert_sql, (img_path, img_data, label))

读取数据库

class SimpleDataset(Dataset):
    def __init__(self, db_path, transform=None) -> None:
        self.db_path = db_path
        self.conn = None
        self.establish_conn()
 
        # 数据库中表名
        self.table_name = 'Synthetic_chinese_dataset'
 
        self.cursor.execute(f'select max(rowid) from {self.table_name}')
        self.nums = self.cursor.fetchall()[0][0]
        self.transform = transform
 
    def __getitem__(self, index: int):
        self.establish_conn()
 
        # 查询
        search_sql = f'select * from {self.table_name} where rowid=?'
        self.cursor.execute(search_sql, (index+1, ))
        img_path, img_bytes, label = self.cursor.fetchone()
 
        # 还原图像和标签
        img = Image.open(BytesIO(img_bytes))
        img = img.convert('RGB')
        img = scale_resize_pillow(img, (320, 32))
 
        if self.transform:
            img = self.transform(img)
        return img, label
 
    def __len__(self) -> int:
        return self.nums
 
    def establish_conn(self):
        if self.conn is None:
            self.conn = sqlite3.connect(self.db_path,
                                        check_same_thread=False,
                                        cached_statements=1024)
            self.cursor = self.conn.cursor()
        return self
 
    def close_conn(self):
        if self.conn is not None:
            self.cursor.close()
            self.conn.close()
 
            del self.conn
            self.conn = None
        return self  
 
# --------------------------------------------------
train_dataset = SimpleDataset(train_db_path, train_transforms)
# ✧✧使用部分，需要手动关闭数据库连接
train_dataset.close_conn()
train_dataloader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              num_workers=n_worker,
                              pin_memory=True,
                              sampler=train_sampler）

最终结论

• TFRecord

  • 转换前后，数据存储大小不变，可以充分利用GPU
  • tfrecord不能接入到其他数据增强方式（imgaug，opencv），且数据增强方式十分有限

• LMDB

  • 转换前后，数据存储大小会变得很大(原始4.2G→转换后96G)
  • PyTorch多进程读取数据时，会出现图像不能还原为原始图像问题，暂时未找到解决方案
  • 读取效率可以充分利用GPU

• 二进制大文件

  • 转换前后，数据存储大小不变
  • 同样，PyTorch多进程读取，也会出现图像不能正确还原的问题，暂时未找到解决方案

• ✧ sqlite(推荐使用)

  • 转换前后，数据存储大小不变
  • 可以正常多进程读取

参考资料

• pytorch-sqlite
• sqlite_dataset

PS：如果有啥问题，还请多多交流。