目录
引言
- Tensorflow有着专门的数据读取模块tfrecord，可以高效地读取训练神经网络模型所用的数据，充分喂饱GPU
- Caffe用lmdb来读取数据，也可以很高效地去读取
- PyTorch有DataLoader读取数据，但是速度比较慢，尤其是小文件较多情况下
- 如何基于PyTorch，高效读取数据，充分利用GPU性能，成为一个关键问题？
TFRecord
- 是否可以将tensorflow下的tfrecord借来一用？未尝不可
- 目前已经有伙伴实现了，详情参见：tfrecord
- 同时，在Kaggle上，也有大神手动实现，详情参见：PyTorch TFRecord-Loader
tfrecord写入代码：
```python
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://*.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读取代码：
```python
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(推荐使用)
- 转换前后，数据存储大小不变
- 可以正常多进程读取