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

时间:2022-11-02 10:06:48

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

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

?
1
2
3
4
5
6
7
8
9
10
11
12
13
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文件的生成,直接上代码:

?
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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文件的读取:

?
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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()

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

原文链接:http://blog.csdn.net/u012222949/article/details/72875281