摘要
在一個(gè)人工智能大爆發(fā)的時(shí)代扮匠,一個(gè)企業(yè)不來(lái)點(diǎn)人工智能都不好意思說(shuō)自己是科技企業(yè)。隨著各公司在各自領(lǐng)域數(shù)據(jù)量的積累凡涩,以及深度學(xué)習(xí)的強(qiáng)擬合特點(diǎn)餐禁,各個(gè)公司都會(huì)訓(xùn)練出屬于自己的模型,那么問(wèn)題就來(lái)了突照,你有模型帮非,我也有模型,那還比什么?對(duì)末盔,就是速度筑舅,誰(shuí)的速度快,誰(shuí)就厲害陨舱。
引言
tensorflow作為谷歌開(kāi)源的深度學(xué)習(xí)神器翠拣,被學(xué)界以及業(yè)界廣泛的推廣,本文就不做過(guò)多的介紹游盲,如有需要可以參考如下鏈接
人們往往在模型迭代指定次數(shù)后误墓,保存一次checkpoint。在模型達(dá)到指定要求的精度后就可以上線(xiàn)應(yīng)用了益缎。NVIDIA tensorrt就是一種專(zhuān)門(mén)用于推理的引擎(不可訓(xùn)練)谜慌,它主要是用于在生產(chǎn)環(huán)境中部署深度學(xué)習(xí)應(yīng)用程序。以此來(lái)提供最大的吞吐量和效率莺奔。他主要是獲取訓(xùn)練好的模型欣范,這些模型通常是32或者16位的,tensorrt降低在這些精度(int8)來(lái)優(yōu)化網(wǎng)絡(luò)令哟。具體的tensorrt介紹就不做過(guò)多的介紹恼琼。
加速步驟(重點(diǎn))
訓(xùn)練好的inceptionV3模型保存的格式是一大堆的ckpt文件:
首先需要將這些文件轉(zhuǎn)化為
.pb
的文件,也就是做模型的持久化屏富,我們只需要模型做前向的推理晴竞,并得到預(yù)測(cè)值即可.這個(gè)步驟主要分為2步
- 保存圖
- 把圖和參數(shù)固定在一起
這一步具體如何操作可以參考這篇博客
接下來(lái)就需要構(gòu)造tenssorrt的推理引擎了(以fp32的為例,int8的話(huà)還需要拿典型圖片構(gòu)造校準(zhǔn)表)
1.導(dǎo)入必要的庫(kù)文件
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)作流傳過(guò)去。
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é)
通過(guò)上述步驟可以將之前的tensorflow模型轉(zhuǎn)化為rt的引擎瘤袖,極大的提高推理速度,但是在flask中使用tensorrt的時(shí)候由于flask的上下文管理方式容易在成引擎找不到設(shè)備上下文昂验,這一點(diǎn)可以用多進(jìn)程的方式得到解決捂敌,在StackOverflow上有人提問(wèn)這個(gè)問(wèn)題,但是解決方案就是每推理一次都初始化一次cuda這樣會(huì)影響整個(gè)應(yīng)用的速度既琴。只有單獨(dú)的開(kāi)一個(gè)進(jìn)程供cuda進(jìn)行初始化占婉,這樣才不會(huì)影響速度懂拾,也不會(huì)造成上下文沖突辣卒,具體的解決方法下一篇博客進(jìn)行介紹