摘要

在一個(gè)人工智能大爆發(fā)的時(shí)代，一個(gè)企業(yè)不來點(diǎn)人工智能都不好意思說自己是科技企業(yè)。隨著各公司在各自領(lǐng)域數(shù)據(jù)量的積累，以及深度學(xué)習(xí)的強(qiáng)擬合特點(diǎn)，各個(gè)公司都會(huì)訓(xùn)練出屬于自己的模型，那么問題就來了，你有模型，我也有模型，那還比什么？對(duì)，就是速度，誰的速度快，誰就厲害。

引言

tensorflow作為谷歌開源的深度學(xué)習(xí)神器，被學(xué)界以及業(yè)界廣泛的推廣，本文就不做過多的介紹，如有需要可以參考如下鏈接

• link1
• link2

人們往往在模型迭代指定次數(shù)后，保存一次checkpoint。在模型達(dá)到指定要求的精度后就可以上線應(yīng)用了。NVIDIA tensorrt就是一種專門用于推理的引擎(不可訓(xùn)練)，它主要是用于在生產(chǎn)環(huán)境中部署深度學(xué)習(xí)應(yīng)用程序。以此來提供最大的吞吐量和效率。他主要是獲取訓(xùn)練好的模型，這些模型通常是32或者16位的，tensorrt降低在這些精度(int8)來優(yōu)化網(wǎng)絡(luò)。具體的tensorrt介紹就不做過多的介紹。

加速步驟(重點(diǎn))

訓(xùn)練好的inceptionV3模型保存的格式是一大堆的ckpt文件：

在這里插入圖片描述

• 保存圖
• 把圖和參數(shù)固定在一起

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

1.導(dǎo)入必要的庫文件

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

uff是將剛才的pb轉(zhuǎn)化為引擎支持的uff文件，該文件可以序列化，也可以直接當(dāng)作流傳過去。
trt則是用于加速推理的tensorrt
pycyda則是用于顯卡cuda編程的
uffparser 解析uff模型的
2.設(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 剛轉(zhuǎn)的pb模型地址
CHANNEL 圖片的通道RGB 3 通道
HEIGHT 圖片的高
WIDTH 圖片的寬
ENGINE_PATH 等會(huì)保存引擎的地址
INPUT_NODE 模型的輸入節(jié)點(diǎn)
OUTPUT_NODE 模型的輸出節(jié)點(diǎn)
INPUT_SIZE 輸入圖片的大小，注意輸入的是 CHANNEL, HEIGHT ,WIDTH
MAX_BATCH_SIZE 每次輸入幾張圖片吧
MAX_WORKSPACE 顯存的大小1<<30也就是1GB的大小，不信你試試
3.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)

4.構(gòu)建引擎

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.調(diào)用引擎實(shí)現(xiàn)推理

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

是不是很好懂？

總結(jié)

通過上述步驟可以將之前的tensorflow模型轉(zhuǎn)化為rt的引擎，極大的提高推理速度，但是在flask中使用tensorrt的時(shí)候由于flask的上下文管理方式容易在成引擎找不到設(shè)備上下文，這一點(diǎn)可以用多進(jìn)程的方式得到解決，在StackOverflow上有人提問這個(gè)問題，但是解決方案就是每推理一次都初始化一次cuda這樣會(huì)影響整個(gè)應(yīng)用的速度。只有單獨(dú)的開一個(gè)進(jìn)程供cuda進(jìn)行初始化，這樣才不會(huì)影響速度，也不會(huì)造成上下文沖突，具體的解決方法下一篇博客進(jìn)行介紹