編譯 TensorFlow 模型

超神經HyperAI發表於2023-05-15
編譯模型

本篇文章譯自英文文件 Compile Tensorflow Models。

更多 TVM 中文文件可訪問 →TVM 中文站。

本文介紹瞭如何用 TVM 部署 TensorFlow 模型。

首先安裝 TensorFlow Python 模組（可參考 https://www.tensorflow.org/install）。

# 匯入 tvm 和 relay
import tvm
from tvm import te
from tvm import relay

# 匯入 os 和 numpy
import numpy as np
import os.path

# 匯入 TensorFlow
import tensorflow as tf

# 讓 TensorFlow 將 GPU 記憶體限制為實際需要的記憶體，而非佔用所有可用的記憶體。
# https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth
# 本教程這樣做，對 sphinx-gallery 更友好。
gpus = tf.config.list_physical_devices("GPU")
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("tensorflow will use experimental.set_memory_growth(True)")
    except RuntimeError as e:
        print("experimental.set_memory_growth option is not available: {}".format(e))

try:
    tf_compat_v1 = tf.compat.v1
except ImportError:
    tf_compat_v1 = tf

# TensorFlow 實用函式
import tvm.relay.testing.tf as tf_testing

# 模型相關檔案的基本位置
repo_base = "https://github.com/dmlc/web-data/raw/main/tensorflow/models/InceptionV1/"

# 測試影像
img_name = "elephant-299.jpg"
image_url = os.path.join(repo_base, img_name)

教程​

參考 docs/frontend/tensorflow.md，獲取 TensorFlow 中各種模型的更多資訊。

model_name = "classify_image_graph_def-with_shapes.pb"
model_url = os.path.join(repo_base, model_name)

# 影像標籤圖
map_proto = "imagenet_2012_challenge_label_map_proto.pbtxt"
map_proto_url = os.path.join(repo_base, map_proto)

# 可讀的標籤文字
label_map = "imagenet_synset_to_human_label_map.txt"
label_map_url = os.path.join(repo_base, label_map)

# target 設定
# 用下面這些註釋為 cuda 構建
# target = tvm.target.Target("cuda", host="llvm")
# layout = "NCHW"
# dev = tvm.cuda(0)
target = tvm.target.Target("llvm", host="llvm")
layout = None
dev = tvm.cpu(0)

下載所需檔案

下載上述列出的檔案：

from tvm.contrib.download import download_testdata

img_path = download_testdata(image_url, img_name, module="data")
model_path = download_testdata(model_url, model_name, module=["tf", "InceptionV1"])
map_proto_path = download_testdata(map_proto_url, map_proto, module="data")
label_path = download_testdata(label_map_url, label_map, module="data")

匯入模型

從 protobuf 檔案建立 TensorFlow 計算圖定義。

with tf_compat_v1.gfile.GFile(model_path, "rb") as f:
    graph_def = tf_compat_v1.GraphDef()
    graph_def.ParseFromString(f.read())
    graph = tf.import_graph_def(graph_def, name="")
    # 呼叫函式將計算圖定義匯入預設計算圖。
    graph_def = tf_testing.ProcessGraphDefParam(graph_def)
    # 給計算圖新增 shape
    with tf_compat_v1.Session() as sess:
        graph_def = tf_testing.AddShapesToGraphDef(sess, "softmax")

解碼影像

備註
TensorFlow 前端匯入不支援 JpegDecode 等預處理操作。 JpegDecode 被繞過（只返回源節點），因此我們只向 TVM 提供解碼後的幀。

from PIL import Image

image = Image.open(img_path).resize((299, 299))

x = np.array(image)

將計算圖匯入 Relay

將 TensorFlow 計算圖定義匯入到 Relay 前端。

結果：

  • sym：給定 TensorFlow protobuf 的 Relay 表示式。
  • params：從 TensorFlow 引數 (tensor protobuf) 轉換而來的引數。
shape_dict = {"DecodeJpeg/contents": x.shape}
dtype_dict = {"DecodeJpeg/contents": "uint8"}
mod, params = relay.frontend.from_tensorflow(graph_def, layout=layout, shape=shape_dict)

print("Tensorflow protobuf imported to relay frontend.")

輸出結果：

/workspace/python/tvm/relay/frontend/tensorflow.py:535: UserWarning: Ignore the passed shape. Shape in graphdef will be used for operator DecodeJpeg/contents.
  "will be used for operator %s." % node.name
/workspace/python/tvm/relay/frontend/tensorflow_ops.py:1009: UserWarning: DecodeJpeg: It's a pass through, please handle preprocessing before input
  warnings.warn("DecodeJpeg: It's a pass through, please handle preprocessing before input")
Tensorflow protobuf imported to relay frontend.

Relay 構建

用給定的輸入規範，將計算圖編譯為 LLVM target。

結果：

  • graph：編譯後的最終計算圖。
  • params：編譯後的最終引數。
  • lib：target 庫（可用 TVM runtime 部署到 target 上） 。
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target, params=params)

輸出結果：

/workspace/python/tvm/driver/build_module.py:268: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.
  "target_host parameter is going to be deprecated. "

在 TVM 上執行可移植計算圖

接下來在 target 上部署編譯好的模型：

from tvm.contrib import graph_executor

dtype = "uint8"
m = graph_executor.GraphModule(lib["default"](dev))
# 設定輸入
m.set_input("DecodeJpeg/contents", tvm.nd.array(x.astype(dtype)))
# 執行
m.run()
# 得到輸出
tvm_output = m.get_output(0, tvm.nd.empty(((1, 1008)), "float32"))

處理輸出

將 InceptionV1 模型的輸出處理為人類可讀文字。

predictions = tvm_output.numpy()
predictions = np.squeeze(predictions)

# 建立節點 ID --> 英文字串查詢
node_lookup = tf_testing.NodeLookup(label_lookup_path=map_proto_path, uid_lookup_path=label_path)

# 列印 TVM 輸出的前 5 個預測。
top_k = predictions.argsort()[-5:][::-1]
for node_id in top_k:
    human_string = node_lookup.id_to_string(node_id)
    score = predictions[node_id]
    print("%s (score = %.5f)" % (human_string, score))

輸出結果：

African elephant, Loxodonta africana (score = 0.61481)
tusker (score = 0.30387)
Indian elephant, Elephas maximus (score = 0.03343)
banana (score = 0.00023)
rapeseed (score = 0.00021)

在 TensorFlow 上推理

在 TensorFlow 上執行對應的模型：

def create_graph():
    """從已儲存的 GraphDef 檔案建立一個計算圖，並返回 saver。"""
    # 從已儲存的 graph_def.pb 建立圖形
    with tf_compat_v1.gfile.GFile(model_path, "rb") as f:
        graph_def = tf_compat_v1.GraphDef()
        graph_def.ParseFromString(f.read())
        graph = tf.import_graph_def(graph_def, name="")
        # 呼叫函式將計算圖定義匯入預設計算圖。
        graph_def = tf_testing.ProcessGraphDefParam(graph_def)

def run_inference_on_image(image):
    """在影像上進行推理。

    引數
    ----------
    image: String 型別
        影像檔名。

    返回值
    -------
        無
    """
    if not tf_compat_v1.gfile.Exists(image):
        tf.logging.fatal("File does not exist %s", image)
    image_data = tf_compat_v1.gfile.GFile(image, "rb").read()

    # 從已儲存的 GraphDef 建立計算圖。
    create_graph()

    with tf_compat_v1.Session() as sess:
        softmax_tensor = sess.graph.get_tensor_by_name("softmax:0")
        predictions = sess.run(softmax_tensor, {"DecodeJpeg/contents:0": image_data})

        predictions = np.squeeze(predictions)

        # 建立節點 ID --> 英文字元查詢
        node_lookup = tf_testing.NodeLookup(
            label_lookup_path=map_proto_path, uid_lookup_path=label_path
        )

        # 列印 TensorFlow 的前 5 個預測。
        top_k = predictions.argsort()[-5:][::-1]
        print("===== TENSORFLOW RESULTS =======")
        for node_id in top_k:
            human_string = node_lookup.id_to_string(node_id)
            score = predictions[node_id]
            print("%s (score = %.5f)" % (human_string, score))

run_inference_on_image(img_path)

輸出結果：

===== TENSORFLOW RESULTS =======
African elephant, Loxodonta africana (score = 0.58394)
tusker (score = 0.33909)
Indian elephant, Elephas maximus (score = 0.03186)
banana (score = 0.00022)
desk (score = 0.00019)

指令碼總執行時長： （1 分 6.352 秒）

下載 Python 原始碼：from_tensorflow.py

下載 Jupyter Notebook：from_tensorflow.ipynb

相關文章