Tensorflow之TFRecord的原理和使用心得
本文首發於微信公眾號「對白的演算法屋」
大家好，我是對白。
目前，越來越多的網際網路公司內部都有自己的一套框架去訓練模型，而模型訓練時需要的資料則都儲存在分散式檔案系統（HDFS）上。Hive作為構建在HDFS上的一個數據倉庫，它本質上可以看作是一個翻譯器，可以將HiveSQL語句翻譯成MapReduce程式或Spark程式，因此模型需要的資料例如csv/libsvm檔案都會儲存成Hive表並存放在HDFS上，那麼問題就來了，如何大規模地把HDFS中的資料直接喂到Tensorflow中呢？Tensorflow提供了一種解決方法：spark-tensorflow-connector，支援將spark DataFrame格式資料直接儲存為TFRecords格式資料，接下來就帶大家瞭解一下TFRecord的原理、構成和如何生成TFRecords檔案。
TFRecord介紹
TFRecord是Tensorflow訓練和推斷標準的資料儲存格式之一，將資料儲存為二進位制檔案（二進位制儲存具有佔用空間少，拷貝和讀取（from disk）更加高效的特點），而且不需要單獨的標籤檔案了，其本質是一行行位元組字串構成的樣本資料。
一條TFRecord資料代表一個Example，一個Example就是一個樣本資料，每個Example內部由一個字典構成，每個key對應一個Feature，key為欄位名，Feature為欄位名所對應的資料，Feature有三種資料型別：ByteList、FloatList，Int64List。
TFRecord構成
它實質上是由protobuf定義的一種資料協議，其中tensorflow提供了兩種Example表示形式 Example和SequenceExample。它的定義程式碼位於[tensroflow/core/example/example.proto & feature.proto]。
Example和SequenceExample的定義：
message Example { Features features = 1; }; message SequenceExample { Features context = 1; FeatureLists feature_lists = 2; }; message Features { // Map from feature name to feature. map<string, Feature> feature = 1; }; // Containers for non-sequential data. message Feature { // Each feature can be exactly one kind. oneof kind { BytesList bytes_list = 1; FloatList float_list = 2; Int64List int64_list = 3; } }; // Containers for sequential data. // // A FeatureList contains lists of Features. These may hold zero or more // Feature values. // // FeatureLists are organized into categories by name. The FeatureLists message // contains the mapping from name to FeatureList. // message FeatureList { repeated Feature feature = 1; }; message FeatureLists { // Map from feature name to feature list. map<string, FeatureList> feature_list = 1; };
我們這裡以最常用的Example來進行解釋。從圖中可以看到，在樣本生產環節，每個Example內部由一個dict構成，每個key（string）對應著一個Feature結構，這個Feature結構有三種具體形式，分別是ByteList，FloatList，Int64List三種。這三種形式便可以承載string，bytes，float，double，int，long等多種樣本結構，並且基於list的表示，使得我們既可以表達scalar，也可以表達vector型別的資料（注意如果想要將一個matrix儲存到到一個Feature內，其值需要時按照Row-Major拍平的1-D array, 行列資料需使用額外欄位儲存，方便反序列化）。這裡需要注意的是，我們在序列化的時候，並未將格式資訊序列化進去，實質上，序列化後的，每條tfrecord中的資料，只具有以下資料：
TFRecord中每條資料的格式：
uint64 length
uint32 masked_crc32_of_length
byte data[length]
uint32 masked_crc32_of_data因此我們可以看出來，TFRecord並不是一個self-describing的格式，也就是說，tfrecord的write和read都需要額外指明schema。從上圖我們也能看出來，在實際訓練的時候，樣本都需要經過一個知曉了Schema的Parser來進行解析，然後才能傳遞給Tensorflow進行實際的訓練。
注：這裡只展示了CTR場景常使用的Example，當然也有影象等場景需要使用SequenceExample進行一些樣本的結構化表達，這裡不做展開。根據官方文件來看，SequenceExample主要是使用在時序特徵和視訊特徵。其中context欄位描述的是和當期時間和特徵不相關的共性資料，而feature_list則持有和時間或者視訊幀相關的資料。感興趣可以參考youtube-8M這個資料集中關於樣本資料的表示。
TFRecord的生成（小規模）
TFRecord的生成=Example序列化+寫入TFRecord檔案
構建Example時需要指定格式資訊（字典）key是特徵，value是BytesList/FloatList/Int64List值，但Example序列化時並未將格式資訊序列化進去，因此讀取TFRecord檔案需要額外指明schema。
每個Example會序列化成位元組字串並寫入TFRecord檔案中，程式碼如下：
import tensorflow as tf
# 回憶上一小節介紹的，每個Example內部實際有若干種Feature表達，下面
# 的四個工具方法方便我們進行Feature的構造
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]))
def _int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def _int64list_feature(value_list):
return tf.train.Feature(int64_list=tf.train.Int64List(value=value_list))
# Example序列化成位元組字串
def serialize_example(user_id, city_id, app_type, viewd_pois, avg_paid, comment):
# 注意我們需要按照格式來進行資料的組裝，這裡的dict便按照指定Schema構造了一條Example
feature = {
'user_id': _int64_feature(user_id),
'city_id': _int64_feature(city_id),
'app_type': _int64_feature(app_type),
'viewd_pois': _int64list_feature(viewd_pois),
'avg_paid': _float_feature(avg_paid),
'comment': _bytes_feature(comment),
}
# 呼叫相關api將Example序列化為位元組字串
example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
return example_proto.SerializeToString()
# 樣本的生產，這裡展示了將2條樣本資料寫入到了TFRecord檔案中
def write_demo(filepath):
with tf.python_io.TFRecordWriter(filepath) as writer:
writer.write(serialize_example(1, 10, 1, [658, 325], 36.3, "yummy food."))
writer.write(serialize_example(2, 20, 2, [897, 568,