tensorflow TFRecords文件的生成和读取的方法

时间:2021-05-22

TensorFlow提供了TFRecords的格式来统一存储数据,理论上,TFRecords可以存储任何形式的数据。

TFRecords文件中的数据都是通过tf.train.Example Protocol Buffer的格式存储的。以下的代码给出了tf.train.Example的定义。

message Example { Features features = 1; }; message Features { map<string, Feature> feature = 1; }; message Feature { oneof kind { BytesList bytes_list = 1; FloatList float_list = 2; Int64List int64_list = 3; } };

下面将介绍如何生成和读取tfrecords文件:

首先介绍tfrecords文件的生成,直接上代码:

from random import shuffle import numpy as np import glob import tensorflow as tf import cv2 import sys import os # 因为我装的是CPU版本的,运行起来会有'warning',解决方法入下,眼不见为净~ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' shuffle_data = True image_path = '/path/to/image/*.jpg' # 取得该路径下所有图片的路径,type(addrs)= list addrs = glob.glob(image_path) # 标签数据的获得具体情况具体分析,type(labels)= list labels = ... # 这里是打乱数据的顺序 if shuffle_data: c = list(zip(addrs, labels)) shuffle(c) addrs, labels = zip(*c) # 按需分割数据集 train_addrs = addrs[0:int(0.7*len(addrs))] train_labels = labels[0:int(0.7*len(labels))] val_addrs = addrs[int(0.7*len(addrs)):int(0.9*len(addrs))] val_labels = labels[int(0.7*len(labels)):int(0.9*len(labels))] test_addrs = addrs[int(0.9*len(addrs)):] test_labels = labels[int(0.9*len(labels)):] # 上面不是获得了image的地址么,下面这个函数就是根据地址获取图片 def load_image(addr): # A function to Load image img = cv2.imread(addr) img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_CUBIC) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 这里/255是为了将像素值归一化到[0,1] img = img / 255. img = img.astype(np.float32) return img # 将数据转化成对应的属性 def _int64_feature(value): return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) def _bytes_feature(value): return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) def _float_feature(value): return tf.train.Feature(float_list=tf.train.FloatList(value=[value])) # 下面这段就开始把数据写入TFRecods文件 train_filename = '/path/to/train.tfrecords' # 输出文件地址 # 创建一个writer来写 TFRecords 文件 writer = tf.python_io.TFRecordWriter(train_filename) for i in range(len(train_addrs)): # 这是写入操作可视化处理 if not i % 1000: print('Train data: {}/{}'.format(i, len(train_addrs))) sys.stdout.flush() # 加载图片 img = load_image(train_addrs[i]) label = train_labels[i] # 创建一个属性(feature) feature = {'train/label': _int64_feature(label), 'train/image': _bytes_feature(tf.compat.as_bytes(img.tostring()))} # 创建一个 example protocol buffer example = tf.train.Example(features=tf.train.Features(feature=feature)) # 将上面的example protocol buffer写入文件 writer.write(example.SerializeToString()) writer.close() sys.stdout.flush()

上面只介绍了train.tfrecords文件的生成,其余的validation,test举一反三吧。。

接下来介绍tfrecords文件的读取:

import tensorflow as tf import numpy as np import matplotlib.pyplot as plt import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' data_path = 'train.tfrecords' # tfrecords 文件的地址 with tf.Session() as sess: # 先定义feature,这里要和之前创建的时候保持一致 feature = { 'train/image': tf.FixedLenFeature([], tf.string), 'train/label': tf.FixedLenFeature([], tf.int64) } # 创建一个队列来维护输入文件列表 filename_queue = tf.train.string_input_producer([data_path], num_epochs=1) # 定义一个 reader ,读取下一个 record reader = tf.TFRecordReader() _, serialized_example = reader.read(filename_queue) # 解析读入的一个record features = tf.parse_single_example(serialized_example, features=feature) # 将字符串解析成图像对应的像素组 image = tf.decode_raw(features['train/image'], tf.float32) # 将标签转化成int32 label = tf.cast(features['train/label'], tf.int32) # 这里将图片还原成原来的维度 image = tf.reshape(image, [224, 224, 3]) # 你还可以进行其他一些预处理.... # 这里是创建顺序随机 batches(函数不懂的自行百度) images, labels = tf.train.shuffle_batch([image, label], batch_size=10, capacity=30, min_after_dequeue=10) # 初始化 init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init_op) # 启动多线程处理输入数据 coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(coord=coord) .... #关闭线程 coord.request_stop() coord.join(threads) sess.close()

好了,就介绍到这里。。,有什么问题可以留言。。大家一起学习。。希望对大家的学习有所帮助,也希望大家多多支持。

声明:本页内容来源网络,仅供用户参考;我单位不保证亦不表示资料全面及准确无误,也不保证亦不表示这些资料为最新信息,如因任何原因,本网内容或者用户因倚赖本网内容造成任何损失或损害,我单位将不会负任何法律责任。如涉及版权问题,请提交至online#300.cn邮箱联系删除。

相关文章