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

 

版权声明：本文为CSDN博主「Liekkas Kono」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：

https://blog.csdn.net/shiwanghualuo/article/details/120778553

=======================================================

引言

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)

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

=====================================================

引言Tensorflow有着专门的数据读取模块tfrecord，可以高效地读取训练神经网络模型所用的数据，充分喂饱GPUCaffe用lmdb来读取数据，也可以很高效地去读取PyTorch有DataLoader读取数据，但是速度比较慢，尤其是小文件较多情况下如何基于PyTorch，高效读取数据，充分利用GPU性能，成为一个关键问题？TFRecord是否可以将tensorflow下的tfrecord借来一用？未尝不可目前已经有伙伴实现了，详情参见：tfrecord同时，在Kaggle上，也有大神手动实现，详情参见：PyTorch TFRecord-Loadertfrecord写入代码：import cv2import numpy as npimport tensorflow as tffrom 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')123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137基于PyTorch下tfrecord读取代码import cv2import numpy as npimport tensorflow as tfimport 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:    pass123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113LMDB纵观各大论坛，说到基于PyTorch下提高小文件读取速度，不得不说到LMDB(Lightning Memory-Mapped Database)了，我也做了一些尝试，最终结论将在最后给出写入LMDBimport osimport picklefrom pathlib import Path import cv2import lmdbimport numpy as npfrom PIL import Imagefrom torch.utils.data import DataLoader, Datasetfrom torchvision import transformsfrom 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')123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899读取LMDBclass 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 things123456789101112131415161718192021222324252627282930313233343536373839404142434445二进制大文件直接将现有数据集按照二进制读取，存入一个bins的大文件中，也不失为一种选择写入binsimport cv2import numpy as npfrom 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)12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061读取binsfrom io import BytesIOfrom PIL import Imageimport cv2import 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)123456789101112131415161718192021222324252627282930313233343536373839404142Sqlite采用python内置的sqlite3作为存储格式，也是一种好的选择写入到sqlite数据库中import sqlite3from 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))1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768读取数据库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）1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859最终结论TFRecord
转换前后，数据存储大小不变，可以充分利用GPUtfrecord不能接入到其他数据增强方式（imgaug，opencv），且数据增强方式十分有限LMDB
转换前后，数据存储大小会变得很大(原始4.2G→转换后96G)PyTorch多进程读取数据时，会出现图像不能还原为原始图像问题，暂时未找到解决方案读取效率可以充分利用GPU二进制大文件
转换前后，数据存储大小不变同样，PyTorch多进程读取，也会出现图像不能正确还原的问题，暂时未找到解决方案✧ sqlite(推荐使用)
转换前后，数据存储大小不变可以正常多进程读取参考资料pytorch-sqlitesqlite_dataset————————————————版权声明：本文为CSDN博主「Liekkas Kono」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。原文链接：https://blog.csdn.net/shiwanghualuo/article/details/120778553