tensorRT加速caffe模型的部署
一、简介及其tensorRT加速原理
在计算资源并不丰富的嵌入式设备上,TensorRT之所以能加速神经网络的的推断主要得益于两点。首先是TensorRT支持INT8和FP16的计算,通过在减少计算量和保持精度之间达到一个理想的trade-off,达到加速推断的目的。
更为重要的是TensorRT对于网络结构进行了重构和优化,主要体现在一下几个方面。
- tensorRT通过解析网络模型将网络中无用的输出层消除以减小计算。
- 对于网络结构的垂直整合,即将目前主流神经网络的conv、BN、Relu三个层融合为了一个层,例如将图1所示的常见的Inception结构重构为图2所示的网络结构。
- 对于网络的水平组合,水平组合是指将输入为相同张量和执行相同操作的层融合一起,如图2向图3的转化。
- 对于concat层,将contact层的输入直接送入下面的操作中,不用单独进行concat后在输入计算,相当于减少了一次传输吞吐。
- 图 1
图 2
图 3
二、改写deploy.prototxt
将deploy.prototxt文件名改写为deploy_plugin.prototxt
- convolution层的param{}全部去掉,convolution_param中的weight_filter{}去掉,bias_filter{}去掉
- 将自定义层的名字改写为IPlugin,自定义层的参数不用写在prototxt中,全部写在IPlugin的实现中。例如SSD的priorbox层,permute层,都是需要改写为IPlugin
- SSD的最后输出是detection_out层,需要将detection_out_param去掉,然后新增加一个输出detection_out2,因为在tensorRT中默认的输出是两个,如果只有一个top,那么程序会报错,“Plugin layer output count is not equal to caffe output count“
三、完成自定义层的代码
- 先实现PluginFactory的class,继承IPluginFactory,class中定义需要使用的自定义层。
- 完善判断函数
//将所有的plugin层的名字全部写到strcmp()中
//通过此函数来判断是否是自定义的IPlugin层
bool PluginFactory::isPlugin(const char* name)
{
return (!strcmp(name,"conv4_3_norm") ||!strcmp(name,"conv4_3_norm_mbox_conf_perm"));
}
- 完善创建plugin函数
//
nvinfer1::IPlugin* PluginFactory::createPlugin(const char* layerName,
const nvinfer1::Weights* weights,
int nbWeights)
{
assert(isPlugin(layerName));
//通过if语句判断是哪一个plugin,然后实现相对应的功能
//下面这个是normalize层的实现,使用tensorRT中自带的createSSDNormalizePlugin函数来实现
if(!strcmp(layerName,"conv4_3_norm"))
{
assert(mNormalizerLayer.get() == nullptr);
mNormalizerLayer = std::unique_ptr<INvPlugin,decltpe(nvPluginDeleter)>(createSSDNormalizePlugin(weights,false,false,0.0001),nvPluginDeleter);
return mNormalizeLayer.get();
}
//下面的这个SSD的permute层
//0,2,3,1的数据就是deploy.prototxt中的conv4_3_norm_mbox_conf_perm这一层的数据
else if(!strcmp(layerName,"conv4_3_norm_mbox_conf_perm"))
{
assert(mConv4_3_norm_mbox_conf_perm_layer.get() == nullptr);
mConv4_3_norm_mbox_conf_perm_layer = std::unique_ptr<INvPlugin,decltype(nvPluginDelter)>(createSSDPermutePlugin({{0,2,3,1}}),nvPluginDelter);
return mConv4_3_norm_mbox_conf_perm_layer.get();
}
//下面这个是priorbox层
//priorbox层的每一个参数都对应于deploy.prototxt中的参数
//需要仔细对比deploy.prototxt文件
else if(!strcmp(layerName,"pool6_mbox_priorbox"))
{
assert(mPool6_mbox_priorbox_layer.get() == nullptr);
PriorBoxParameters params;
float minsize[1] = {200};
float maxsize[1] = {240};
float aspect_ratio[2] = {1.0,2.0};
params.minSize = minsize;
params.maxSize = maxsize;
params.aspectRatio = aspect_raio;
params.numMinSize = 1;
params.numAspectRatios = 2;
params.numMaxSize = 1;
params.flip = true;
params.clip = false;
params.variance[0] = 0.1;
params.variance[1] = 0.1;
params.variance[2] = 0.2;
params.variance[3] = 0.2;
params.imgH = 0;
params.imgW = 0;
params.stepH = 300;
params.stepW = 300;
params.offset = 0.5;
mPool6_mbox_priorbox_layer = std::unique_ptr<INvPlugin,decltype(nvPluginDeleter)>(createSSDPriorBoxPlugin(params),nvPluginDeleter);
return mPool6_mbox_priorbox_layer.get();
}
}
- 如果tensorRT中没有对应的实现,那就需要自己手撸代码了。例如SSD中的softmax(axive = 2)这种情况,tensorRT中原生不支持,需要自己用代码实现
//softmax layer
else if(!strcmp(layerName,"mbox_conf_softmax"))
{
assert(mPluginSoftmax == nullptr);
assert(nbWeights == 0 && weights == nullptr);
mPluginSoftmax = std::unique_ptr<SoftmaxPlugin>(new SoftmaxPlugin());
return mPluginSoftmax.get();
}
//手动实现softmax
class SoftmaxPlugin : public IPlugin
{
public:
int initialize() override {return 0;}
inline void terminate() override {}
SoftmaxPlugin(){}
SoftmaxPlugin(const void* buffer,size_t size)
{
assert(size == sizeof(mCopySize));
mCopySize = *reinterpret_cast<const size_t*>(buffer);
}
inline int getNbOutputs() const override
{
return 1;
}
Dims getOutputDimensions(int index,const Dims* inputs,int nbInputDims) override
{
assert(nbInputDims == 1);
assert(index == 0);
assert(inputs[index].nbDims == 3);
//这里的5是最后检测的类别数
assert((inputs[0].d[0]) * (inputs[0].d[1] % 5 == 0);
return DimsCHW(inputs[0].d[0] , inputs[0].d[1],inputs[0].d[2]);
}
size_t getWorkspaceSize(int) const override
{
return nCopySize * 1;
}
//最重要的是这里的函数
int enqueue(int batchSize,const void*const *input,void** outputs,void* workspace,cudaStream_t stream) override
{
cudaSoftmax(1917*5,5,(float*)*inputs,static_cast<float *>(*outputs));
return 0;
}
}
cudaSoftmax函数用cuda代码实现,查看.cu文件
tensorRT的使用
- caffeToTRTModel
void TensorNet::caffeToTRTModel(const std::string& deployFile, const std::string& modelFile, const std::vector<std::string>& outputs,
unsigned int maxBatchSize)
{
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();
ICaffeParser* parser = createCaffeParser();
parser->setPluginFactory(&pluginFactory);
bool useFp16 = builder->platformHasFastFp16();
//useFp16 = false;
DataType modelDataType = useFp16 ? DataType::kHALF : DataType::kFLOAT;
std::cout << deployFile.c_str() <<std::endl;
std::cout << modelFile.c_str() <<std::endl;
//std::cout << (*network) <<std::endl;
std::cout << "Here : 1"<<std::endl;
const IBlobNameToTensor* blobNameToTensor = parser->parse(deployFile.c_str(),modelFile.c_str(),*network, DataType::kFLOAT);
std::cout << "Here : 2" <<std::endl;
assert(blobNameToTensor != nullptr);
std::cout << "Here : 3" <<std::endl;
for (auto& s : outputs) network->markOutput(*blobNameToTensor->find(s.c_str()));
builder->setMaxBatchSize(maxBatchSize);
builder->setMaxWorkspaceSize(10 << 20);
std::cout << "Here : 4"<< std::endl;
ICudaEngine* engine = builder->buildCudaEngine( *network );
std::cout << "Here : 5"<<std::endl;
assert(engine);
network->destroy();
parser->destroy();
gieModelStream = engine->serialize();
engine->destroy();
builder->destroy();
pluginFactory.destroyPlugin();
shutdownProtobufLibrary();
}
这里的函数中
bool useFp16 = builder->platformHasFastFp16();
显卡是不支持fp16的计算的,但是tx2等嵌入式平台支持fp16,调用此函数可以得到是否支持fp16.
四、需要注意的地方
- 修改prototxt文件一定要仔细,确认没有prototxt的语法错误,然后根据IPlugin的数量来改写代码,IPlugin的数量和
isPlugin
函数的数量保持一致 - priorbox层中的参数,按照原始的deploy改写的时候需要注意,在deploy.prototxt中的priorbox的参数,例如
aspect_ratio: 2.0
,则在代码中需要写成float aspect_ratio[2] = {1.0,2.0};
如果是aspect_ratio: 3.0
,则需要写成`float aspect_ratio[3] = {1.0,2.0,3.0}; - detection_out层中的参数,tensorRT的版本不同,命名规则是不一样的,需要根据tensorRT版本来改变。
- 在softmax的参数,需要有priorbox的数量,例如vgg6-ssd的数量的8732,resnet20-ssd的数量是1917,红绿灯ssd的数量是4992。这个参数可以在reshape层中打印出来看到。在reshape类定义中,getOutputDimensions函数中打印出来。