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

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

準(zhǔn)備工作

生成.pb的模型文件

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

春天不是讀書天：[深度學(xué)習(xí)] TensorFlow中模型的freeze_graph

2. 導(dǎo)入必要的庫

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轉(zhuǎn)化為TensorRT引擎支持的uff文件，該文件可以序列化，也可以直接當(dāng)作流傳過去。
• pycyda：用于顯卡cuda編程的，如果要使用TensorRT的python API，這是一個必須的庫
• uffparser ：用于解析uff模型

3. 參數(shù)設(shè)置

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：圖片的通道、高和寬，根據(jù)模型的輸入大小確定
• ENGINE_PATH：等會保存TensorRT引擎的地址
• INPUT_NODE：模型的輸入節(jié)點(diǎn)
• OUTPUT_NODE：模型的輸出節(jié)點(diǎn)，是一個列表，如果有許多個輸出節(jié)點(diǎn)，就將節(jié)點(diǎn)名都列入這個列表中
• INPUT_SIZE：輸入圖片的大小，注意通道在前還是后，這里輸入的是 CHANNEL, HEIGHT ,WIDTH
• MAX_BATCH_SIZE：在推理的時候，每次輸入幾張圖片
• MAX_WORKSPACE：顯存的大小1<<30也就是1GB的大小。有的時候，程序運(yùn)行是會報內(nèi)存溢出的錯，這個時候就可以調(diào)小MAX_WORKSPACE，比如2 << 10

將tensorflow模型轉(zhuǎn)換成TensorRT

pb轉(zhuǎn)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的文件格式轉(zhuǎn)成了uff文件格式。你需要知道的一個概念是，UFF(Universal Framework Format)是一種描述DNN執(zhí)行圖的數(shù)據(jù)格式。綁定執(zhí)行圖的是輸入與輸出，所以parser.register_input和parser.register_output做的事情是將tensorflow模型的輸入輸出在UFF文件中記錄。

注意，對于多個輸出，因為OUTPUT_NODE是一個列表，所以將多個輸出節(jié)點(diǎn)依次放入列表就可以了。

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

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

假設(shè)你的模型在輸入層同時輸入了三張圖片，那么你需要定義3個輸入節(jié)點(diǎn)，并且指定order分別為0、1、2。這里的order指的是模型的輸入在uff結(jié)構(gòu)中的順序，這種order在接下來的binding會得到體現(xiàn)。

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)

以上代碼創(chuàng)建了TensorRT中的engine，即引擎，這個engine將負(fù)責(zé)模型的前向運(yùn)算。TensorRT是一個用于推理的加速工具，所以前向計算就夠了。

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

使用以下語句，我們就保存了一個.plan文件。PLAN文件是運(yùn)行引擎用于執(zhí)行網(wǎng)絡(luò)的序列化數(shù)據(jù)。包含權(quán)重，網(wǎng)絡(luò)中執(zhí)行步驟以及用來決定如何綁定輸入與輸出緩存的網(wǎng)絡(luò)信息。

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

使用TensorRT實現(xiàn)推理

現(xiàn)在，讓我們調(diào)用之前保存的plan文件，啟用引擎，開始使用TensorRT實現(xiàn)推理。

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

引擎叫做engine，而引擎運(yùn)行的上下文叫做context。engine和context在推理過程中都是必須的，這兩者的關(guān)系如下：

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

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

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

現(xiàn)在準(zhǔn)備好一張可以輸入給模型的圖像 img.jpg，并且轉(zhuǎn)換成fp32

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

同時，創(chuàng)建一個array來“接住”輸出數(shù)據(jù)。為什么說“接住”呢，因為之后你就會看到，引擎做前向推理計算的時候，是生成了一個數(shù)據(jù)流，這個數(shù)據(jù)流會寫入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實現(xiàn)內(nèi)存申請。 bindings = [int(d_input), int(d_output)]

現(xiàn)在，我們可以開始TensorRT上的推理計算了！

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

這個時候，如果你將output打印出來，就會發(fā)現(xiàn)output數(shù)組中已經(jīng)有值了，這就是TensorRT計算的結(jié)果。

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

特別注意！

TensorRT和Tensorflow的數(shù)據(jù)格式不一樣，Tensorflow是NHWC格式，即channel_last，而TensorRT中是NCHW格式，即channel_first，比如一張RGB圖像，在Tensorflow中表示為（224， 224， 3），在TensorRT中就是（3，224， 224）。所以使用TensorRT時，請一定確認(rèn)圖像的格式。

參考資料：

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

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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。