Tensorflow之dataset介紹

雖然搭建模型是算法的主要工作，但在實際應用中我們花在數據處理的時間遠比搭建模型的多。
而且每次由於數據格式不同，需要重複實現數據加載，構造batch以及shuffle等代碼。

記得之前接觸過的Pytorch、Paddle等框架都有dataset的工具，當然tensorflow也不例外，
經過一段時間的瞭解和實踐，準備寫下這篇博文來記錄一下。

TFrecord格式

Tensorflow支持多種輸入格式，當然數據最後都會被處理成tensor。
其中tfrecord格式作爲Tensorflow自家研製的數據格式(使用protocol buffer結構化數據存儲格式)，
有着以下優點：

• 使用二進制格式存儲，數據處理更加高效，佔用空間更少；
• 支持多文件輸入和懶加載，這一點對於大數據量和通過網絡加載數據來說非常重要；
• Tensorflow提供TFRecordDaset API進行數據加載

有優點就有缺點，對於tfrecord來說，缺點就是其他數據類型都要經歷轉化成tfrecord這一過程，
並且不能直接查看數據內容（二進制的格式都有這個缺點）

tfrecord中的每一條樣本叫做Example，其實就是對應protobuf中的message，格式是{“string”: tf.train.Feature}映射。
其中Feature類型有三種取值(.proto file).

1. tf.train.BytesList
2. tf.train.FloatList
3. tf.train.Int64List

支持類型看起來不多，實際上對於我們輸入到模型的數據來說，已經是夠用的了。

其他格式轉化爲TFrecord

我們常用的格式大多是csv、json和libsvm，使用python的話，我們可以通過如下代碼利用tf的api進行轉化。

import tensorflow as tf

out_file = "xxx"
label = 1
index = [1,2,3]
value = [0.5, 0.5, 0.5]
with tf.io.TFRecordWriter(out_file) as writer:
    feature = {
        'label':
            tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
        'index': tf.train.Feature(int64_list=tf.train.Int64List(value=index)),
        'value': tf.train.Feature(float_list=tf.train.FloatList(value=value)),
    }
    proto = tf.train.Example(features=tf.train.Features(feature=feature))
    writer.write(proto.SerializeToString())

但是這種方式可擴展性不強，效率較低。考慮到我們大多時候使用spark來處理數據，
這時可以使用spark-tensorflow-connector這個依賴包把Dataframe保存成tfrecord格式。

以下是保存成tfrecord格式的代碼，其實跟平時保存成parquet沒啥差別。這裏不但保存成tfrecord，
並且以gzip格式壓縮。實際測試中，壓縮率能達到25%，對於二進制格式文件來說，原本就沒有可讀性，
所以個人傾向對tfrecord文件進行壓縮（spark保存成parquet，默認也用了snappy壓縮）。

這次實踐中出現一種情況，pom文件中加入spark-tensorflow-connector依賴打包成jar包，保存結果的tfrecord並沒有被壓縮。
而以--jars spark-tensorflow-connector的方式加入依賴則是正常的。後續有同學找到原因可以一起探討一下。

df.select("label", "index", "value")
        .write.format("tfrecords")
        .option("recordType", "Example")
        .option("codec", "org.apache.hadoop.io.compress.GzipCodec")
        .mode("overwrite").save(outFilePath)

TFrecord讀取

一旦我們生成了tfrecord，首先需要驗證一下數據是否如我們所期望的一致。用以下代碼可以打印第一條數據內容：

import tensorflow as tf

files = ["/path/to/tfrecord"]
tf.compat.v1.enable_eager_execution()
raw_dataset = tf.data.TFRecordDataset(files, compression_type="GZIP")

for raw_record in raw_dataset.take(1):
  example = tf.train.Example()
  example.ParseFromString(raw_record.numpy())
  print(example)

驗證完了我們就可以把tfrecord文件加載成dataset, tfrecord支持多文件輸入，所以傳給dataset的是一個列表,
每個元素是文件路徑名。Tensorflow支持從S3讀取，這樣我們只要用spark把數據保存到S3即可, S3文件名可以通過以下代碼獲取。

def get_tfrecord_files(dir_list, sep=","):
  """
  Get tfrecord file list from directory
  """
  dirs = dir_list.strip().split(sep)
  rlt = []
  for dir_iter in dirs:
    files = tf.io.gfile.listdir(dir_iter)
    rlt.extend([
        os.path.join(dir_iter, f)
        for f in files
        if f.endswith(".tfrecord") or f.endswith(".gz")
    ])
  return rlt

files = get_tfrecord_files("/path/to/s3/dir/")
dataset = tf.data.TFRecordDataset(files, compression_type="GZIP")

此時，我們得到的dataset每一條數據還是序列化之後的Example數據，在實際應用中，需要對其進行解析。
解析時，我們需要構建一個feature_description，原因在於dataset使用graph-execution，所以需要
feature_description來確定Tensor的shape和類型。

for raw_record in dataset.take(10):
  print(repr(raw_record))

# Create a description of the features.
feature_description = {
    'feature0': tf.io.FixedLenFeature([], tf.int64, default_value=0),
    'feature1': tf.io.FixedLenFeature([], tf.int64, default_value=0),
    'feature2': tf.io.FixedLenFeature([], tf.string, default_value=''),
    'feature3': tf.io.FixedLenFeature([], tf.float32, default_value=0.0),
}

def _parse_function(example_proto):
  # Parse the input `tf.Example` proto using the dictionary above.
  return tf.io.parse_single_example(example_proto, feature_description)

parsed_dataset = dataset.map(_parse_function)

Dataset操作

推薦使用dataset API的原因在於，他提供了一套構建數據Pipeline的操作，包括以下三部分

1. 從數據源構造dataset（支持TFRecordDataset，TextLineDataset, CsvDataset等方式)
2. 數據處理操作 (通過map操作進行轉化，tfrecord格式需要對Example進行解析)
3. 迭代數據 （包括batch, shuffle, repeat, cache等操作）

數據讀取部分上面已經提過，只要傳一個file list給TFRecordDataset即可。而數據處理這塊，由於我們
喂入的數據都是處理好的了，只要寫一個example的解析函數即可（對於圖像數據則可能還會涉及到裁剪、變形等操作）。
接下來就是一些dataset常見的操作了.

這裏涉及到dataset的操作順序和一些參數調優的問題：

1. map可以通過num_parallel_calls參數實現並行化，提高數據轉化效率

2. prefetch能夠實現數據加載和模型訓練並行效果，充分利用cpu和gpu資源

3. batch, shuffle, repeat這三者的順序不同，結果也會不一樣。具體可以運行下面代碼並調整三者順序來驗證.

   import tensorflow as tf
   tf.compat.v1.enable_eager_execution()
   dataset = tf.data.Dataset.from_tensor_slices(list(range(0, 10)))
   dataset = dataset.batch(3, drop_remainder=False)
   dataset = dataset.repeat(2)
   dataset = dataset.shuffle(4)
   

   先batch再repeat

   先repeat再batch

   shuffle機制，若先做batch再shuffle，則shuffle的是各個batch，而不是batch裏的樣本。

參考資料

1. Tensorflow Records? What they are and how to use them
2. Google Protocol Buffer 的使用和原理
3. TFRecord and tf.Example
4. tensorflow:input pipeline性能指南