摘要

在一個人工智慧大爆發的時代，一個企業不來點人工智慧都不好意思說自己是科技企業。隨著各公司在各自領域資料量的積累，以及深度學習的強擬合特點，各個公司都會訓練出屬於自己的模型，那麼問題就來了，你有模型，我也有模型，那還比什麼？對，就是速度，誰的速度快，誰就厲害。

引言

tensorflow作為谷歌開源的深度學習神器，被學界以及業界廣泛的推廣，本文就不做過多的介紹，如有需要可以參考如下連結

• link1
• link2

人們往往在模型迭代指定次數後，儲存一次checkpoint。在模型達到指定要求的精度後就可以上線應用了。NVIDIA tensorrt就是一種專門用於推理的引擎(不可訓練)，它主要是用於在生產環境中部署深度學習應用程式。以此來提供最大的吞吐量和效率。他主要是獲取訓練好的模型，這些模型通常是32或者16位的，tensorrt降低在這些精度(int8)來優化網路。具體的tensorrt介紹就不做過多的介紹。

加速步驟(重點)

訓練好的inceptionV3模型儲存的格式是一大堆的ckpt檔案：

首先需要將這些檔案轉化為.pb的檔案，也就是做模型的持久化，我們只需要模型做前向的推理，並得到預測值即可.這個步驟主要分為2步

• 儲存圖
• 把圖和引數固定在一起

這一步具體如何操作可以參考這篇部落格
接下來就需要構造tenssorrt的推理引擎了(以fp32的為例，int8的話還需要拿典型圖片構造校準表)

1.匯入必要的庫檔案

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

uff是將剛才的pb轉化為引擎支援的uff檔案，該檔案可以序列化，也可以直接當作流傳過去。
trt則是用於加速推理的tensorrt
pycyda則是用於顯示卡cuda程式設計的
uffparser 解析uff模型的
2.設定一些常量

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 圖片的通道RGB 3 通道
HEIGHT 圖片的高
WIDTH 圖片的寬
ENGINE_PATH 等會儲存引擎的地址
INPUT_NODE 模型的輸入節點
OUTPUT_NODE 模型的輸出節點
INPUT_SIZE 輸入圖片的大小，注意輸入的是 CHANNEL, HEIGHT ,WIDTH
MAX_BATCH_SIZE 每次輸入幾張圖片吧
MAX_WORKSPACE 視訊記憶體的大小1<<30也就是1GB的大小，不信你試試
3.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)

4.構建引擎

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

5.儲存模型

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

6.呼叫引擎實現推理

def infer32(context, input_img, batch_size):
    engine = context.get_engine()
    assert(engine.get_nb_bindings() == 2)
    start = time.time()
    dims = engine.get_binding_dimensions(1).to_DimsCHW()
    elt_count = dims.C() * dims.H() * dims.W() * batch_size
    input_img = input_img.astype(np.float32)
    output = cuda.pagelocked_empty(elt_count, dtype=np.float32)
    d_input = cuda.mem_alloc(batch_size * input_img.size * input_img.dtype.itemsize)
    d_output = cuda.mem_alloc(batch_size * output.size * output.dtype.itemsize)
    bindings = [int(d_input), int(d_output)]
    stream = cuda.Stream()
    cuda.memcpy_htod_async(d_input, input_img, stream)
    context.enqueue(batch_size, bindings, stream.handle, None)
    cuda.memcpy_dtoh_async(output, d_output, stream)
    end = time.time()
    return output

是不是很好懂？

總結

通過上述步驟可以將之前的tensorflow模型轉化為rt的引擎，極大的提高推理速度，但是在flask中使用tensorrt的時候由於flask的上下文管理方式容易在成引擎找不到裝置上下文，這一點可以用多程序的方式得到解決，在StackOverflow上有人提問這個問題，但是解決方案就是每推理一次都初始化一次cuda這樣會影響整個應用的速度。只有單獨的開一個程序供cuda進行初始化，這樣才不會影響速度，也不會造成上下文衝突，具體的解決方法下一篇部落格進行介紹