標籤：

[深度學習]TensorRT加速tensorflow實例

04-29

使用TensorRT加速tensorflow

模型的推理應該是很有市場的一種應用了，但是使用Python的、易懂的例子並不多，官方的文檔在這方面也是很不友好。

所以，本文旨在提供一個能把原理講明白，代碼能跑的起來的實例，本例中用到模型是inception V3

準備工作

  1. 生成.pb的模型文件

首先我們需要從保存模型的chekpoint文件中，生成.pb的模型文件。這一步叫做模型的持久化，具體的做法可以參考之前寫的這篇文章：

春天不是讀書天：[深度學習] TensorFlow中模型的freeze_graph

2. 導入必要的庫

import tensorflow as tf
import uff
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from tensorrt.parsers import uffparser

  • uff：是將剛才的pb轉化為TensorRT引擎支持的uff文件，該文件可以序列化，也可以直接當作流傳過去。
  • pycyda：用於顯卡cuda編程的，如果要使用TensorRT的python API，這是一個必須的庫
  • uffparser ：用於解析uff模型

3. 參數設置

MODEL_DIR = ./model_seg/model.pb
CHANNEL = 3
HEIGHT = 299
WIDTH = 299
ENGINE_PATH = ./model_seg/model_.pb.plan
INPUT_NODE = input
OUTPUT_NODE = [InceptionV3/Logits/SpatialSqueeze]
INPUT_SIZE = [CHANNEL, HEIGHT ,WIDTH]
MAX_BATCH_SIZE = 1
MAX_WORKSPACE = 1<<30

  • MODEL_DIR：第一步中生成的pb模型地址
  • CHANNEL、HEIGHT、WIDTH：圖片的通道、高和寬，根據模型的輸入大小確定
  • ENGINE_PATH：等會保存TensorRT引擎的地址
  • INPUT_NODE：模型的輸入節點
  • OUTPUT_NODE：模型的輸出節點，是一個列表，如果有許多個輸出節點，就將節點名都列入這個列表中
  • INPUT_SIZE：輸入圖片的大小，注意通道在前還是後，這裡輸入的是 CHANNEL, HEIGHT ,WIDTH
  • MAX_BATCH_SIZE：在推理的時候，每次輸入幾張圖片
  • MAX_WORKSPACE：顯存的大小1<<30也就是1GB的大小。有的時候，程序運行是會報內存溢出的錯，這個時候就可以調小MAX_WORKSPACE，比如2 << 10

將tensorflow模型轉換成TensorRT

  1. pb轉uff 並解析模型

G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
uff_model = uff.from_tensorflow_frozen_model(FROZEN_GDEF_PATH, OUTPUT_NODE)
parser = uffparser.create_uff_parser()
parser.register_input(INPUT_NODE, INPUT_SIZE, 0)
parser.register_output(OUTPUT_NODE)

這裡做的事情是將pb的文件格式轉成了uff文件格式。你需要知道的一個概念是，UFF(Universal Framework Format)是一種描述DNN執行圖的數據格式。綁定執行圖的是輸入與輸出，所以parser.register_input和parser.register_output做的事情是將tensorflow模型的輸入輸出在UFF文件中記錄。

注意，對於多個輸出，因為OUTPUT_NODE是一個列表，所以將多個輸出節點依次放入列表就可以了。

如果是多個輸入的話，則需要將輸入節點名一個個的記錄在uff中。register_input()需要3個參數：

  • name – Input name.
  • shape – Input shape.
  • order – Input order on which the framework input was originally.

假設你的模型在輸入層同時輸入了三張圖片，那麼你需要定義3個輸入節點，並且指定order=0，因為這三張圖片是同時輸入的。

parser.register_input(INPUT_NODE1, INPUT_SIZE, 0)
parser.register_input(INPUT_NODE2, INPUT_SIZE, 0)
parser.register_input(INPUT_NODE3, INPUT_SIZE, 0)

2. 保存模型

engine = trt.utils.uff_to_trt_engine(
G_LOGGER,
uff_model,
parser,
MAX_BATCH_SIZE,
MAX_WORKSPACE,
datatype=trt.infer.DataType.FLOAT)

以上代碼創建了TensorRT中的engine，即引擎，這個engine將負責模型的前向運算。TensorRT是一個用於推理的加速工具，所以前向計算就夠了。

在engine創建成功之後，就可以使用了。不過，一個建議是將結果保存下來。畢竟到目前為止，雖然代碼很少，但是將pb文件成功轉換成uff文件是不容易的（誰用誰知道！）

使用以下語句，我們就保存了一個.plan文件。PLAN文件是運行引擎用於執行網路的序列化數據。包含權重，網路中執行步驟以及用來決定如何綁定輸入與輸出緩存的網路信息。

trt.utils.cwrite_engine_to_file(./model_.pb.plan,engine.serialize())

使用TensorRT實現推理

現在，讓我們調用之前保存的plan文件，啟用引擎，開始使用TensorRT實現推理。

engine = trt.utils.load_engine(G_LOGGER, ./model_.pb.plan)

引擎叫做engine，而引擎運行的上下文叫做context。engine和context在推理過程中都是必須的，這兩者的關係如下：

context = engine.create_execution_context()
engine = context.get_engine()

在運行前向運算前，我們還需要做一次確認。get_nb_bindings()是為了獲取與這個engine相關的輸入輸出tensor的數量。對於本例中單輸入輸出的模型，tensor的數量是2。如果有多個輸入輸出，這個確認值就要相應的變化，比如3個輸入，1個輸出的模型，tensor的數量就是4。我們需要知道這個數量，是為了之後的顯存分配做準備。

print(engine.get_nb_bindings())
assert(engine.get_nb_bindings() == 2)

現在準備好一張可以輸入給模型的圖像 img.jpg，並且轉換成fp32

img = cv2.imread(img.jpg)
img = img.astype(np.float32)

同時，創建一個array來「接住」輸出數據。為什麼說「接住」呢，因為之後你就會看到，引擎做前向推理計算的時候，是生成了一個數據流，這個數據流會寫入output array中

#create output array to receive data
OUTPUT_SIZE = 10
output = np.zeros(OUTPUT_SIZE , dtype = np.float32)

我們需要為輸入輸出分配顯存，並且綁定。

# 使用PyCUDA申請GPU顯存並在引擎中註冊
# 申請的大小是整個batchsize大小的輸入以及期望的輸出指針大小。
d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)

# 引擎需要綁定GPU顯存的指針。PyCUDA通過分配成ints實現內存申請。
bindings = [int(d_input), int(d_output)]

現在，我們可以開始TensorRT上的推理計算了！

# 建立數據流
stream = cuda.Stream()
# 將輸入傳給cuda
cuda.memcpy_htod_async(d_input, img, stream)
# 執行前向推理計算
context.enqueue(1, bindings, stream.handle, None)
# 將預測結果傳回
cuda.memcpy_dtoh_async(output, d_output, stream)
# 同步
stream.synchronize()

這個時候，如果你將output列印出來，就會發現output數組中已經有值了，這就是TensorRT計算的結果。

如過你使用tensorflow的方法，對同一組輸入數據做預測，看看計算的結果是否一致 ，因為精度的差異會有一些差異，但是大體上來說，使用tensorflow和TensorRT，會得到一致的結果。

參考資料：

tensorRt加速tensorflow模型推理（inception V3為例）

blog.csdn.net/abrams90/


推薦閱讀：

TAG:深度學習（DeepLearning） | TensorFlow |