tensorflow通過tfrecord高效讀寫資料

利用tensorflow提供的tfrecord資料存儲格式工具，我們可以将我們已經進行過處理的資料儲存起來，以便我們下次更高效地讀取，略過資料處理的過程，提高效率。具體的步驟大概分為以下幾步：

1. 将資料轉化為

   tf.train.Feature

   ，然後存于字典；
2. 接着，将其轉化為

   tf.train.example

   ，然後進行序列化，寫入

   tf.python_io.TFRecordWriter

   ，到這裡就完成了寫入的操作；
3. 讀取的時候，首先通過

   tf.data.TFRecordDataset

   來讀取它；
4. 然後，通過

   tf.parse_single_example

   來解析還原資料原本的結構；
5. 最後就可以結合我們上次提高的

   batch

   來批量擷取。（關于

   batch、shuffle、repeat

   函數的詳細介紹TensorFlow dataset.shuffle、batch、repeat用法）

寫入TFRecordWriter

import tensorflow as tf
import collections
import numpy as np

inputs_1 = np.array([
        [[1,2], [3,4]],
        [[5,6], [7,8]]
        ], dtype=np.int32)  # 這裡的類型int32要與解析時使用的類型保持一緻
inputs_2 = [
        [1.1, 2.2, 3.3], 
        [4.4, 5.5, 6.6]
        ]
lables = [0, 1]

def create_int_feature(values):
  f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))  # 需要注意這裡接受的格式是list，并且隻能是一維的
  return f

def create_float_feature(values):
  f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
  return f

def create_bytes_feature(values):
  f = tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))
  return f

writer = tf.python_io.TFRecordWriter('test.tfrecord')  # test.tfrecord是寫入的檔案路徑

for i1, i2, l in zip(inputs_1, inputs_2, lables):
    features = collections.OrderedDict()  # 這裡是建立一個有序字典
    # 對于多元數組，隻能先将其轉化為byte，才能傳遞給Feature
    features['inputs_1'] = create_bytes_feature(i1.tostring())
    features['inputs_2'] = create_float_feature(i2)
    features['labels'] = create_int_feature([l])
    
    example = tf.train.Example(features=tf.train.Features(feature=features))
    
    writer.write(example.SerializeToString())
writer.close()
           

讀取并解析為dataset

源資料的資料類型最好與解析時使用的類型保持一緻，不然可能會出現錯誤。

如inputs_1（ndarray）時使用的類型時np.int32，那麼在tf.decode_raw進行解析時就相應地使用tf.int32。

name_to_features = {
      "inputs_1": tf.FixedLenFeature([], tf.string),  
      "inputs_2": tf.FixedLenFeature([3], tf.float32),  # 這裡的格式需要與寫入的保持一緻，否則可能出現解析錯誤
      "labels": tf.FixedLenFeature([], tf.int64)
  }

d = tf.data.TFRecordDataset('test.tfrecord')
d = d.repeat() # 這裡repeat不傳遞參數，則會無限重複
d = d.shuffle(buffer_size=2)
# map_and_batch其實就是map和batch結合在一起而已
d = d.apply(tf.contrib.data.map_and_batch(
        lambda record: tf.parse_single_example(record, name_to_features),
        batch_size=1))

iters = d.make_one_shot_iterator()
batch = iters.get_next()

# BytesList解析時會丢失shape資訊，需要自己還原它的shape，是以一般也會将shape資料一同寫入
# 這裡需要将byte解析成原本的資料結構，這裡的tf.int32需要與寫入時的格式保持一緻
inputs_1_batch = tf.decode_raw(batch['inputs_1'], tf.int32)  # tf.int32類型要與源資料的類型保持一緻
inputs_1_batch = tf.reshape(inputs_1_batch, [-1, 2, 2])  
# 因為每次是batch個inputs_1，是以shape是[-1, 2, 2]，原來的shape是[2, 2]
inputs_2_batch = batch['inputs_2']
labels_batch = batch['labels']

sess = tf.Session()
# 這樣我們就可以每次擷取一個batch的資料了
sess.run([inputs_1_batch, inputs_2_batch, labels_batch])
           

歡迎關注同名公衆号：“我就算餓死也不做程式員”。

交個朋友，一起交流，一起學習，一起進步。