基於TensorFlow模型檔案的靜態和動態分析
簡介
在往期文章中 ofollow,noindex" target="_blank">我們給你推薦一種TensorFlow模型格式 ，我們介紹過TensorFlow眾多的模型格式，不同格式在服務端或者移動端場景下有特殊的優化，但無論用什麼工具轉都一定包含了Graph（計算圖）和Variables（模型權重）的資訊。有了Graph和Variables，理論上我們就可以重建模型訓練時的場景，甚至通過Graph editor等工具去拓展模型功能做更多有趣的事情。
目前TensorFlow提供了SavedModelBuilder和tf.saved_model.loader.load介面，載入後的Graph可以更新後新建的Session中，進而實現檢查模型簽名、統計Op資訊等靜態功能。而由於TensorFlow模型簽名輸入輸出是強型別和固定Shape的，因此我們還能在本地模擬模型上線的Inference介面，從而實現基於本地檔案的效能Benchmark等動態功能。
本文將以tfmodel專案為例，介紹基於TensorFlow模型檔案的靜態和動態分析，專案地址 tobegit3hub/tfmodel 。
模型靜態分析
首先TensorFlow的SavedModel的模型格式有一定的規範，同一個目錄下可以有多個模型版本（Model version），但需要約定使用相同的模型簽名，通過檔案的目錄結構也可以檢查模型是否有效，如下圖。
當然檔案存在不能保證模型可用，例如saved_model.pb和variables.data都是特殊序列化介面的protobuf檔案，要保證模型可用需要用TensorFlow API去載入測試。我們驗證過TensorFlow的Python API、C++ API和Java API都提供了可用的tf model loader介面，如果你想寫程式碼只需要下面幾行就可以實現了。
try: session = tf.Session(graph=tf.Graph()) tf.saved_model.loader.load(session, [tf.saved_model.tag_constants.SERVING], self.model_version_path) return True except Exception as e: logging.error("Fail to validate and get error: {}".format(e)) return False
靜態檢查除了檢查模型有效性，還可以檢視模型的模型簽名和Op細節。以檢查模型簽名為例，TensorFlow Serving官方也推出過簡單的Python工具來列印Signature資訊，saved_model_cli可以指定模型版本的tag和method列印對應資訊（嗯，一個TensorFlow模型目錄可以有多個模型版本，一個模型版本可以有多個簽名，每個簽名還有對應的tag）。
當然，我們也可以自己寫程式碼實現定製化的列印格式，所有的模型簽名信息都在tf.saved_model.loader.load返回的meta_graph物件中，是一個protobuf序列化的Python物件，從中獲取對應的signature_name、method_name、input和output tensors即可，例如用tfmodel工具返回的工具如下。
如果只是列印模型簽名，所有的工具看起來都差不多，實際上載入模型後的Session還可以做更多有用的事情，後面動態分析中就會介紹到。
模型動態分析
所謂模型的動態分析，就是讓模型可以執行起來，進而分析出模型的預估效能以及視訊記憶體佔用等指標資訊。目前Google的TensorFlow Serving tensorflow/serving 和我們的Simple TensorFlow Serving tobegit3hub/simple_tensorflow_serving 都可載入TensorFlow模型提供API服務，但基於Server的Benchmark在一定程度上受限於框架本身實現的效能。實際上看過兩者的實現程式碼，分別是呼叫了C++和Python的Session.run()介面而已，底層當然還是TensorFlow的MKL或者CUDA庫等實現。
既然如此，我們就可以不依賴TensorFlow Serving或者Simple TensorFlow Serving，在模型分析工具中載入模型，並通過Session.run()去做Inference和預估模型的效能，如果有配套工具還可以檢查模型預估過程中每個Op的耗時和視訊記憶體的使用變化情況。
除此之外，如果訓練過程中忘記匯出TensorBoard event files了，也可以在分析模型時載入完整的Graph，然後匯出到TensorBoard中，由於支援Session.run()功能，還可以通過構造不同分佈的測試資料來展示更多訓練時未被考慮的指標。
我們組同事還在做基於Graph editor的功能，可以將匯出的單機版SavedModel模型修改成分散式版本上線，也可以動態修改Op執行的device從而適配離線訓練與線上預估不同的GPU資源情況，當然這個與模型分析無關暫時不細講了。
使用tfmodel
前面介紹了靜態和動態分析模型的基本原理，後面介紹的tfmodel就可以實現幾乎所有可能做到的功能（主要是TensorFlow Serving命令列工具也沒有做的事情）。tfmodel的安裝也很簡單，可以使用pip install tensorflow-model或者通過docker run tobegit3hub/tfmodel來執行。
首先是檢查TensorFlow模型有效性的功能，通過前面介紹的tf.saved_model.loader.load即可做到，通過命令列的檢查可以快速判斷模型檔案是否可以上線，因為載入模型程式碼與預估服務是一致的，可以保證檢查通過的模型檔案都符合上線規範。
其次是檢查TensorFlow模型的簽名和Op資訊，這個其實是受到一個Torch專案的啟發 Swall0w/torchstat ，雖然不是官方專案但可以展示模型每一層的input shape、output shape、params、memory、flops等資訊。目前tfmodel inspect子命令提供了模型簽名的視覺化和“偽造”請求資料的功能，未來也可以通過TensorFlow API做更多事情。
然後是對TensorFlow模型的本地Benchmark功能，是通過本地起程序執行Session.run()實現的。前面提到的資料構造功能，在MXNet、PyTorch、Scikit-learn等框架下確實比較難實現的，因為無論還是開源框架的模型格式還是ONNX的通用格式，都很少有把模型簽名作為模型一部分匯出，目前MXNet和ONNX可以通過額外的一個JSON檔案提供。
tfmodel benchmark提供了不同batch size下的效能測試報告，這是在本地無網路通訊條件下測試的，理論上效能最接近模型本身的運算效能，而且還可以支援GPU。目前測試用佔用本地所有資源，因此我們還提供了docker run tobegit3hub/tfmodel容器，可以在容器中限定運算資源來跑Benchmark。
最後是生成TensorBoard event files功能，TensorBoard絕對是深度學習模型結構視覺化以及模型效果視覺化的神器，如果訓練過程中沒有保留event files，可以使用tfmodel工具重新載入模型Graph然後生成TensorBoard依賴的event files。
最近，Google還把以前基於Chrome和Timeline檔案的Op cost time視覺化功能整合到TensorBoard中，也就是說在TensorBoard視覺化的Graph中，還可以看到每一個step的每個op的耗時情況，以及op的device還有哪些op未被呼叫也畫得清清楚楚。之前我們就用於定位一個上游依賴多個op但其中一個被feed_dict攔截傳值的問題，這部分程式碼會影響訓練效能但在SavedModel的離線分析中加入分析就非常合適了，示例如下。
tf.summary.scalar('test', tf.constant(1)) merged = tf.summary.merge_all() with tf.Session() as sess: writer = tf.summary.FileWriter("./tensorboard", sess.graph) run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) run_metadata = tf.RunMetadata() for i in range(10): summary, result = sess.run( [merged, h], feed_dict={a: 10, d: 100}, options=run_options, run_metadata=run_metadata) writer.add_run_metadata(run_metadata, 'step%d' % i) writer.add_summary(summary, i)
總結
最後總結一下，TensorFlow因為實現了save和load後Session和Graph物件一致的介面，讓開發者只基於模型檔案也可以做很多靜態和動態的事情，讓模型檔案本身的校驗和效能評估有了更好的做法。當然TensorFlow模型也有一定的缺點，基於Op的粒度就不能像Keras、Torch這種逐個layer展示模型資訊，而且手動的device placement在多GPU和分散式的預估中增加了不少開發成本。
大家如果有更多模型評估的點子也可以在評論或者Github中提出，當然除了TensorFlow其他框架要支援類似功能也是相當容易的，首先我們還是考慮把MXNet的signature json寫到模型檔案中吧。