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

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

準備工作

生成.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

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、1、2。這里的order指的是模型的輸入在uff結構中的順序，這種order在接下來的binding會得到體現。

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

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的數據格式不一樣，Tensorflow是NHWC格式，即channel_last，而TensorRT中是NCHW格式，即channel_first，比如一張RGB圖像，在Tensorflow中表示為（224， 224， 3），在TensorRT中就是（3，224， 224）。所以使用TensorRT時，請一定確認圖像的格式。

參考資料：

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

https://blog.csdn.net/abrams90/article/details/80410308

總結

以上是生活随笔為你收集整理的tensorflow打印模型图_[深度学习]TensorRT加速tensorflow实例的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。