適合大量資料I/O的儲存格式: TFRecord簡介與操作教學

適合大量資料I/O的儲存格式: TFRecord簡介與操作教學

在訓練模型的時候，資料輸入的流程直接影響了模型訓練或推理的效率。其中，資料格式 (data format) 相當地影響了I/O的速度，除了更換更高階的固態硬碟外，選擇binary形式的資料格式也是提升I/O效率的方法之一。

在許多binary資料格式中，TFRecord 的最大優勢在於支援部分資料的載入 (lazy loading) 與tf.data的讀取支援。TFRecord的每一筆資料解碼後都是一個類似Protocol buffers 的binary message，包裝的流程相較pickle或HDF都來的繁雜許多，讓許多使用者望而卻步。其實，TFRecord的操作就是將資料一層一層地打包與打開protocal message，並不是非常的困難。

這篇文章會從頭到尾，詳細地介紹各種資料型態儲存到TFRecord的操作，包含基本的純量 (整數、浮點數等)，到影像與高維矩陣這些進階卻有高實務需求的資料型態。

TFRecord簡介

**在訓練模型的時候，資料輸入的流程直接影響了模型訓練或推理的效率。**一般來說，資料會以某種資料格式 (data format)，比如說圖片或是文字檔，儲存在檔案系統 (file system) 中，在透過程式讀取建立輸入流程 (input pipeline)，送到模型 (model) 之中。

**資料格式相當地影響了I/O的速度，除了更換更高階的固態硬碟外，選擇binary形式的資料格式也是提升I/O效率的方法之一。**Python常見的binary檔案格式包含HDF5、pickle與npy，

**而TFRecord也是一種binary檔案格式，由Google的Tensorflow團隊開發，遵循Google Protocol Buffer，實現跨平台、跨語言間保持結構的序列資料。**因此TFRecord可以說是為了Tensorflow而生，但又不專屬於Tensorflow的格式。

也就是用 Tensorflow不代表一定要用tf.data與TFRecord。反之，寫 Pytorch或其他框架也可以搭配tf.data與TFRecord。

在許多binary資料格式中，TFRecord **的最大優勢在於支援部分資料的載入 (lazy loading) 與tf.data的讀取支援，**假設現在有一個40GB的pickle檔案，而device的內存 (memory) 只有32G，那就無法讀取這個pickle檔案，必須自己先將這個pickle拆成兩個小於內存容量的檔案才能讀取。但TFRecord的資料是允許一筆一筆讀入的，也就是無關原始檔案大小，需要多少就讀多少。

其實將資料建成 TFRecord時，通常還是會拆分成多個小檔案，讓 I/O可以平行操作。經驗法則上維持同時打開 10個以上的TFRecord讀取是比較有效率的。每個TFRecord檔案大小至少10+MB，理想100+MB。

TFRecord專門儲存序列式binary資料，而其中的每一筆資料 (或稱為record)，解碼binary後都是一個Protocol buffers **的protocal message。**詳細解讀Protocal buffers有點麻煩，簡單藉由一個例子來看。

TFRecord中許多繁雜的class其實都是protocol message在Python中的形式，比如說tf.train.Example與tf.train.Feature。

在操作TFRecord時，需要藉由tf.train.Example將資料封裝成protocol message。tf.train.Example本質就是一個protocal messgae，由{“string”: value}組成，而value的部分是另一個protocal messagetf.train.Feature。tf.train.Feature內部也是由protocal message組成 (BytesList,FloatList或Int64List)。

這一串聽起來很複雜，簡單說就是所有資訊都要一層一層地包成protocal message。

上圖包含了儲存資料大部分的資料型態，包含影像、矩陣與各種純量等等。以下將介紹上圖中的所有元件，從最簡單的純量資訊讀取開始，到影像與矩陣的存取。

基本資料讀寫(Scalar)

在寫一個TFRecord時，任何資料格式，不論是影像、矩陣、字串等，都是要走一樣的protocal message打包流程，才能放進TFRecord裡。

1. 轉換成serialized的浮點數、整數、字串或是byte等格式 (實際支援string,byte,float (float32),double (float64),bool,enum,int32,uint32,int64,uint64)。
2. 將這些資訊轉換成三種Tensorflow定義的List資料 (BytesList,FloatList,Int64List)，這三個資料型態都是protocol message。
3. 將每個List資料裝進tf.train.Feature。
4. 將tf.train.Feature裝進tf.train.Example裡。

整個流程看起來很複雜，因為 TFRecord的實作上希望可以是跨平台、跨語言保持結構的序列資料，設計上非常貼合 Goggle Protocol Buffer，才會形成這樣的多層封裝結構。

假設現在有一個簡單的case，我們要儲存：學生的名子string、GPAfloat與排名int。第一步要做的事將這三個資料轉成protocol message，也就是對應的BytesList、loatList、Int64List。

通常這三類格式轉換非常的固定，寫法大多都長這樣。

• BytesList: 可接受包含string與byte。
• FloatList: 可接受包含float (float32)與double (float64)。
• Int64List: 可接受包含bool、enum、int32、uint32、int64與uint64。

接下來就是將轉換好的protocol message打包進tf.train.Feature，然後再進一步封裝成tf.train.Example，一般來說以上的三個步驟都習慣一氣呵成地完成。

**讀取TFRecord時通常會使用tf.data.TFRecordDataset的API來讀取，透過tf.data.Dataset使用。**簡單點來看，tf.data.Dataset可以視為一個Python generator，依序地將資料讀出。

以讀取TFRecord來說，每次讀取出來的資料都是一個tf.train.Example，搭配tf.io.parse_single_example與給定dtype的tf.io.FixedLenFeature就可以將原本的資訊讀出。

進階資料讀寫(Image & Array)

影像資訊是資料格式要處理的常見型態，在training input pipeline一種常用的安排是指儲存影像的路徑，要使用時再去讀取存在file system裡原始的影像。與之相對的另外一種就是把影像也包進資料格式之中，這邊又分為先解碼影像 (轉成array) 與不先解碼影像 (維持binary) 兩種。

如果影像保持binary形式，TFRecord讀寫是非常簡單的。也就是在讀取時直接用binary的方法將影像當成string讀進來，然後包成BytesList，收工。

讀取時直接讀出來的資訊會是string型態的圖片，還需要在經過decode才會變成常用的矩陣型態 (直接使用tf.image.decode_jpeg是較有效率的方法)。

如果今天有有二維或更高維的矩陣需要讀寫，或是說影像想要先解碼成矩陣在儲存，那就需要考慮高維矩陣的讀寫，因為TFRecord只接受序列的scalar形式。

不過其實高維矩陣的讀寫跟binary影像讀寫一樣換湯不換藥，透過tf.io.serialize_tensor將矩陣encode成string的形式，讀取時再透過tf.io.parse_tensor來decode回原本的矩陣。

[2021.08.03補充]

其實高維矩陣讀寫也可以直接把tf.io.serialize_tensor的numpy結果 (string)寫到TFRecord裡。已經是 string型態就不必再包Example這一層。

這個資訊感謝 @Chien Chi Liao提醒修正。

這邊要註記一下，tf.data與TFRecord非必須。發現 pipeline是 bottleneck時，可以嘗試使用tf.data。發現 data I/O是 bottleneck時，可以嘗試使用TFRecord。

Take-away

1. TFRecord專門儲存序列式binary資料，而其中的每一筆資料 (或稱為record)，解碼binary後都是一個Protocol buffers的protocal message。
2. TFRecord的最大優勢在於支援部分資料的載入與原生tf.data的讀取支援。
3. TFRecord其實就是將資訊一層一層地包成protocal message，從底層的三大list message (BytesList,FloatList,Int64List)，到Feature，最後到Example。
4. tf.data與TFRecord非必須。發現pipeline是bottleneck時，可以嘗試使用tf.data。發現data I/O是bottleneck時，可以嘗試使用TFRecord。
5. 用Tensorflow不代表一定要用tf.data與TFRecord。反之，Pytorch或其他框架也可以搭配tf.data與TFRecord。
6. 在將資料建成TFRecord時，可以建成多個檔案，允許在I/O上平行操作。經驗法則上維持同時打開10個以上的TFRecord讀取是比較有效率的。每個TFRecord檔案大小至少10+MB，理想100+MB。

好了～這篇文章就先到這邊。老話一句，Tensorflow更新的速度很快，我的資訊或許也有錯誤或過時，若有發現什麼錯誤或值得討論的地方，歡迎回覆文章或來信一起討論 :)

Reference

1. tf.data: Build TensorFlow input pipelines[TensorFlow Core]
2. TFRecord and tf.train.Example[TensorFlow Core]
3. Protocol Buffers[Google Developers]