
使用TensorRT能夠有效加速tensorflow模型的推理,同時C++相比于其他語言要更加的高效,是以在追求模型推斷速度時,用TensorRT的C++APi來部署tensorflow模型是一種不錯的方式。本文旨在提供一個用TensorRT的C++APi來部署tensorflow模型的執行個體,講解部署過程中可能遇到的問題和處理方法,以提取inception_resnet_v2最後全局平均層輸出特征為例。TensorRT版本5.1.5
準備pb模型
images
首先構模組化型,然後load權重,之後調用tf.graph_util.convert_variables_to_constants函數将模型持久化,其中outputs是需要作為輸出的tensor的清單,最後用pb_graph.SerializeToString()将圖序列化并寫入到pb檔案當中,這樣就生成了pb模型。
生成uff檔案
有了pb模型,需要将其轉換為tensorRT可用的uff模型,隻需調用uff自帶的convert腳本即可
python /usr/lib/python2.7/site-packages/uff/bin/convert_to_uff.py pbmodel_name.pb
如轉換成功會輸出如下資訊
用C++API調用模型
調用uff模型的代碼主要參考TensorRT官方的samples中的sampleUffMNIST代碼,下面我會主要講解其中幾個重要的部分
TensorRT要進行推斷首先需要定義網絡結構,然後生成Cuda engine,從uff生成engine的代碼如下,輸入uff檔案名和輸入輸出op的名字就能生成engine
ICudaEngine
首先要定義一個IBuilder* builder,以及一個用來解析uff檔案的parser,之後通過builder建立network, parser會将uff中定義的網絡結構和權值解析出來放到network當中,需要注意的是uffparser解析時需要先給定網絡輸入輸出輸出的節點。之後就能夠用builder根據network中的結構生成 ICudaEngine* engine。生成engine前我們需要定義最大的batchsize數,在之後使用engine時輸入的batchsize不能超過這個數值否則就會出錯,這個值最好設定的和實際使用時的batch size數一緻,這時候tensorRT的效率最高,輸入batchsize小于最大batchsize時不能達到最高的效率。另外還需指定一個工作空間的最大大小。
生成engine之後就可以開始進行推斷,在推斷前首先需要生成一個執行上下文IExecutionContext* context,直接用engine->createExecutionContext()獲得。執行推斷的代碼如下
std
總體來說有三個步驟,首先在裝置(顯示卡)上開辟記憶體用來裝網絡的輸入和輸出,位址存在buffer當中。輸出部分直接按照輸出大小開辟空間,輸入部分通過createCudaBuffer開辟記憶體并且将輸入的資料拷貝到裝置上去。之後可以調用context->execute(batchSize,&buffers[0])來執行推斷的運算,最後通過verifyOutput函數将裝置上的輸出結果拷貝出來到。createCudabuffer代碼如下
void* createCudaBuffer(int64_t eltCount, DataType dtype,const std::vector<cv::Mat>& srcs)
{
/* in that specific case, eltCount == INPUT_H * INPUT_W */
int batchSize = srcs.size();
assert(eltCount ==batchSize * INPUT_H * INPUT_W * INPUT_C);
assert(elementSize(dtype) == sizeof(float));
std::cerr<<"start to generate input"<<std::endl;
size_t memSize = eltCount * elementSize(dtype);
float* inputs = new float[eltCount];
int threadNum=batchSize;
std::thread threads[threadNum];
//uint8_t fileData[batchSize * INPUT_H * INPUT_W * INPUT_C];
int part = batchSize/threadNum;
if(batchSize%threadNum!=0) part=part+1;
for(int i =0;i<threadNum;i++){
threads[i]=std::thread(preprocess,i*part,min(i*part+part,batchSize),std::ref(srcs),inputs);
}
for(int i =0;i<threadNum;i++)
threads[i].join();
void* deviceMem = safeCudaMalloc(memSize);
CHECK(cudaMemcpy(deviceMem, inputs, memSize, cudaMemcpyHostToDevice));
//std::cerr<<"finish to generate input"<<std::endl;
delete[] inputs;
return deviceMem;
}
将圖檔進行預處理後存到數組inputs中然後将inputs中内容拷貝到裝置上(預處理部分采用多線程加速)。verifyOutput代碼和sampleUffMNIST示例中類似,将裝置記憶體上網絡的輸出結果拷貝出來,這樣就是執行了一次完整推斷。加速效果在不同GPU上有所不同,在Tesla M40上相比于tensorflow中推斷可以穩定獲得3倍的速度提升
簡單的插件層示例
如果tensorflow模型中有一些TensorRT不支援的操作,那麼要成功的在TensorRT中調用模型就需要自行編寫這些操作的實作作為TensorRT的插件,我以Sign取符号函數為例講解一些TensorRT插件的編寫
class
首先要定義Sign插件從IPluginV2繼承,代碼中列出來繼承後需要實作的類方法,核心函數enqueue中調用了實作sign功能的cuda核函數進行計算。
為了能讓TensorRT的uffparser能認出我們編寫的層,我們還需要定義一個Creator用來建立插件層示例代碼如下
class
完成這些之後我們就可以在TensorRt中使用Sign操作了,在網絡定義中直接調用Sign可以通過如下代碼實作,先定義一個Sign操作,然後通過addPluginV2添加到網絡中
auto sl = Sign();
auto sign_layer=network->addPluginV2(input,1,sl);
如果想直接讓uffparser直接解析含有Sign操作的網絡則需要在轉換uff檔案是進行一些步驟。
out, _ = inception_resnet_v2(images, is_training=False)
out = tf.sign(out)
存在warning說不能識别Sign操作,自動轉換為名為Sign的插件(tensorRT預設起的名字實際沒有對應op),我們需要告訴轉換器在轉換時将tensorflow的Sign操作轉換為我們編寫的Sign插件(Sign_TRT),是以需要編寫一個config.py檔案給uff轉換器這部分知識可參考TensorRT samples sampleUffSSD
本例子中config.py如下
import graphsurgeon as gs
import tensorflow as tf
sign_node = gs.create_plugin_node("Sign", op="Sign_TRT", dtype=tf.float32)
namespace_plugin_map = {
"Sign": sign_node
}
def preprocess(dynamic_graph):
dynamic_graph.collapse_namespaces(namespace_plugin_map)
首先注冊一個sign node然後給出一個映射将tensorflow圖中的sign op對應到sign node上,告訴轉換器按照映射将op轉換為對應的tensorRT 層。按照如下方式調用
python /usr/lib/python2.7/site-packages/uff/bin/convert-to-uff pbmodel_name.
則輸出變為如下結果
将Sign操作轉化為了我們之前編寫的插件Sign,之後按照之前步驟就能夠正常用C++API調用uff檔案進行推斷。
心得體會
1.雖然用tensorRT能夠縮短很多的GPU推斷時間,但是部署的服務最終響應時間并不僅僅由GPU推斷時間決定,整個流程中的資料預處理、資料的編解碼都會存在影響,在其他部分也應該注意效率。一個例子就是前文中提到的多線程預處理batch圖檔,如果單線程操作就會在預處理的部分直接耗時很長,成為整個服務響應時間的瓶頸。
2. 生成pb模型、轉換uff模型以及調用uffparser時register Input,output,這三個過程中輸入輸出節點的名字一定要注意保持一緻,否則最終在parser進行解析時會出現錯誤。
3.運作程式出現cuda failure一般情況下是由于将記憶體資料拷貝到磁盤時出現了非法記憶體通路,注意檢查buffer開辟的空間大小和拷貝過去資料的大小是否一緻
4.tensorRT中插件其實有比較多種IPlugin,IPluginExt,IPluginV2,注意在添加自己寫的層到網絡時要用對應的函數addPlugin,addPluginExt,addPluginV2,否則有些預設調用的類方法時不會調用的,比如用addPlugin添加的層是不會調用configureWithFormat方法的。
參考資料
Nvidia TensorRT Samples
tensorrt-developer-guide
TensorRT API Docs