目录
1. 概述
2. 创建 TFRecord
2.1 准备工作
2.2 图片数据集转换为 TFRecord 格式
2.3 图片地址集转换为 TFRecord 格式
2.4 小结
3. 读取 TFRecord 文件
为了训练一个深度神经网络,通常会在本地保存一些数据,这些数据分别被保存在 train,validation,test 文件夹中。这些文件夹散列地保存了成千上万的图片或文本文件,无论是直接读取图片还是通过 csv 文件保存图片路径的间接读取方式都会带来两个问题:
读取速度慢:读取图片内容时需要逐个读取本地磁盘中的散列文件
内存空间:许多大型数据(如ImageNet)无法一次性加载到内存中
TensorFlow 提供一种统一的格式来存储数据,即 TFRecord,允许将任意的数据转换为TensorFlow所支持的格式。TFRecord 内部使用了“Protocol Buffer”二进制数据编码方案,只占用一个内存块,只需要一次性加载一个二进制文件的方式即可,简单,快速,尤其对大型训练数据很友好。因此,在进行深度模型训练时,建议将训练样本存储为 TFRecord 格式,当训练数据量比较大的时候,可以将数据分成多个TFRecord文件,提高处理效率。
TFRecord 文件中的数据通过 tf.train.Example Protocol Buffer 的格式存储,tf.train.Example 的定义如下:
message Example {
Features features = 1;
};
message Features {
map feature = 1;
};
message feature {
oneof kind {
BytesList bytes_list = 1;
FloatList float_list = 2;
Int64List int64_list = 3;
}
};
从定义看,tf.train.Example 的数据结构较为简洁,包含了一组 Features(一个从属性名到取值的字典) ,其中属性名是 string 类型,属性的取值 feature
是可以是 BytesList
(字符串)、FloatList
(实数列表) 或者 Int64List
(整数列表)。
为本地训练数据集创建 TFRecord 格式的训练数据集有两种方式:
对于小规模数据集,将数据数据集集整体转换成 TFRecord 格式
对于大规模数据集,将图片数据集对应的地址集转换为 TFRecord 格式。
本节以猫狗数据集为例,部分训练样本如下图,使用以上两种方法分别创建 tfrecords 文件。
将猫狗数据集中所有的训练样本的路径及其对应的标签写入到本地文件中。
def get_data(data_dir):
'''
:description: 获取数据集中所有训练样本的路径和标签
:param data_dir: 数据集存放目录
:return: 返回训练样本的路径列表和标签列表
'''
images_path = []
images_label = []
for root, _, names in os.walk(data_dir):
for name in names:
images_path.append(os.path.join(root, name))
if name.split('.')[0] == 'cat':
images_label.append(0)
if name.split('.')[0] == 'dog':
images_label.append(1)
return images_path, images_label
def save_data_path(data_dir, save_path):
'''
:description: 保存训练样本的路径和对应标签到指定文件中
:param data_dir: 训练样本所在目录
:param save_path: 指定文件保存路径
:return: 无返回值
'''
if os.path.exists(save_path):
os.system("rm {}".format(save_path))
os.system("touch {}".format(save_path))
images_path, images_label = get_data(data_dir)
temp = np.array([images_path, images_label])
temp = temp.transpose()
np.random.shuffle(temp)
images = list(temp[:, 0])
labels = list(temp[:, 1])
with open(save_path, 'w') as f:
for i in range(len(images)):
content = '{} {} \n'.format(images[i], labels[i])
f.write(content)
if __name__ == 'main':
train_dir = '../data/images/train'
save_path = '../data/images/train.txt'
save_data_path(data_dir, save_path)
整体转换为本地训练数据集创建 TFRecord 格式的训练数据集可以按照以下四个步骤进行:
准备数据并使用 tf.io.TFRecordWriter
创建一个 TFRecordWriter
对象
确定本地数据集的 Features
,并使用以下三个函数:tf.train.Example()
、tf.train.Features()
和 tf.train.Feature()
,将单个样例的 Features
转换成一个 record
使用 writer.write(record)
方法将单个样例的 record
写入 TFRecord 文件
等到所有样例都写入 TFRecord 文件后,关闭 TFRecordWriter
对象
import os
from tqdm import tqdm
from PIL import Image
def read_dataset(filepath):
'''
:description: 从本地文件中读取训练集中所有样本的路径及其对应标签
:param filepath: 本地文件路径
:return: 无返回值
'''
with open(filepath, 'r') as f:
trainset = []
lines = f.readlines()
for line in lines:
line = line.strip('\n')
image_path = line.split()[0]
label = int(line.split()[1])
trainset.append([image_path, label])
return trainset
def encode_to_tfrecord(filepath, record_name, save_dir):
'''
:description: 将数据集整体转换为 tfrecords 格式
:param filepath: 保存着所有样本路径的文件
:param record_name: 待保存的 tfrecord 文件的文件名
:param save_dir: tfrecord 文件保存的目录
:return: None
'''
record_path = os.path.join(save_dir, record_name + '.tfrecords')
if os.path.exists(record_path):
os.remove(record_path)
trainset = read_dataset(filepath)
writer = tf.python_io.TFRecordWriter(record_path)
pbar = tqdm(trainset)
for train_data in pbar:
try:
image = Image.open(train_data[0])
image_raw = image.tobytes()
label = train_data[1]
example = tf.train.Example(
features=tf.train.Features(feature={
'image_raw': tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[image_raw])),
'label': tf.train.Feature(
int64_list=tf.train.Int64List(
value=[label]))
})
)
writer.write(example.SerializeToString())
except IOError:
print('could not read:', train_data[0])
pbar.set_description('transforming: {}'.format(train_data[0].split('/')[-1]))
writer.close()
if __name__ == 'main':
encode_to_tfrecord(filepath='../data/images/train.txt',
record_name='cat_dog',
save_dir='../data/tfrecords')
程序运行结果:经过大约11分31秒的时间,25000 幅图像被转换转成了 TFRecord 格式的训练集,转换得到的 TFRecord 文件大小为 11.34 GB。
地址集转换为 TFRecord 的步骤和整体转换几乎一致,只不过本地数据集的 Features 有点不一样。在整体转换中,单个样例的 Features 分别为图像数据和标签,而本次转换中,单个样例的 Features 则是图像的本地地址和标签,转换过程如下:
import os
from PIL import Image
from tqdm import tqdm
def read_dataset(filepath):
'''
:description: 从本地文件中读取训练集中所有样本的路径及其对应标签
:param filepath: 本地文件路径
:return: 无返回值
'''
with open(filepath, 'r') as f:
trainset = []
lines = f.readlines()
for line in lines:
line = line.strip('\n')
image_path = line.split()[0]
label = int(line.split()[1])
trainset.append([image_path, label])
return trainset
def convert_to_tfrecord(filepath, record_name,
save_dir='./data/tfrecords', encoding='utf-8'):
'''
:description: 使用 gfile 方式将数据集转换为 tfrecords 格式
:param filepath: 保存着数据集路径的文件
:param record_name: 将要保存的 tfrecord 文件的文件名
:param save_dir: tfrecord 文件保存的目录
:param encoding: 字符串转换为字节类型的编码方式
:return: None
'''
record_path = os.path.join(save_dir, record_name + '.tfrecords')
if os.path.exists(record_path):
os.remove(record_path)
writer = tf.python_io.TFRecordWriter(record_path)
trainset = utils.read_dataset(filepath)
pbar = tqdm(trainset)
for train_data in pbar:
pbar.set_description('transforming {}'.format(train_data[0].split('/')[-1]))
image_path = bytes(train_data[0], encoding=encoding)
image_label = train_data[1]
example = tf.train.Example(
features=tf.train.Features(feature={
'image_path': tf.train.Feature(
bytes_list=tf.train.BytesList(
value=[image_path])),
'image_label': tf.train.Feature(
int64_list=tf.train.Int64List(
value=[image_label]))
}))
writer.write(example.SerializeToString())
writer.close()
if __name__ == 'main':
convert_to_tfrecord(filepath='../data/images/train.txt',
record_name='cat_dog',
save_dir='../data/tfrecords')
程序运行结果:经过大约 4 秒的时间, 25000 幅图像的地址集转换成了 TFRecord 格式,转换后得到的 TFRecord 文件大小为 2.9 MB。
对比以上两种方式可以发现:
转换得到的 TFRecord 格式的数据集大小将远远大于原始数据集的大小;
整体数据集转换为 TFRecord 格式耗费时间较长;
对于大规模数据集而言,将其转换为 TFRecord 格式是一个非常浩大的工程,而且往往由于转换后的TFRecord 格式的数据集容量过于庞大,后续的加载和读取将耗费更多的资源,从而引起一系列问题。
因此,工程上,通常选择将大规模数据集的地址集转换为 TFRecord 格式,每次直接读取生成 batch 后的地址,并通过这些地址,找到训练样本。
TFRecord 格式的数据集生成完成后即可用于神经网络的训练。在训练时,通常需要从 TFRecord 格式的数据集中获取训练样本,从 TFRecord 格式的数据集中获取训练样本按照以下五个步骤进行:
获取文件列表,并使用 tf.train.string_input_producer()
方法创建文件名队列 filename_queue
。(说明:当数据集较大时,通常创建多个 TFRecord 文件来保存数据集)
使用 tf.TFRecordReader()
创建一个 TFRecordReader
对象。
使用 reader.read()
方法从文件名队列中读取样本。
使用 tf.io.parse_single_example()
方法解析出单个样本,并使用 tf.reshape() 对样本进行 reshape。
使用 tf.train.batch()
或者 tf.train.shuffle_batch()
方法将队列中解析到的样本组合成一个批次
def decode_from_tfrecords(filename_list, batch_size):
filename_queue = tf.train.string_input_producer(filename_list)
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
features = tf.io.parse_single_example(serialized_example, features={
'image_raw': tf.FixedLenFeature([], dtype=tf.string),
'label': tf.FixedLenFeature([], tf.int64),
})
image = tf.decode_raw(features['image_raw'], tf.uint8)
image = tf.reshape(image, [227, 227, 3])
label = tf.cast(features['label'], tf.int32)
min_after_dequeue = 1000
capacity = min_after_dequeue * 3 + batch_size
image_batch, label_batch = tf.train.shuffle_batch([image, label], batch_size,
min_after_dequeue=min_after_dequeue,
capacity=capacity,
num_threads=3)
return image_batch, label_batch