python如何读取tfrecord_tensorflow TFRecords文件的生成和读取的方法

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

TFRecords文件中的数据都是通过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;

}

};

下面将介绍如何生成和读取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()

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