TensorRT学习(1) 使用C++API构建引擎
1. tensorrt轮子代码示例
首先我们先将tensorRT的函数简单封装一下,方便后面搭建网络,以下为部分示例:
// size() 获取张量形状
static vector<int> trt_size(ITensor *input) {
auto shape = input->getDimensions().d;
auto dim_num = input->getDimensions().nbDims;
vector<int> d(dim_num, 0);
for (int i = 0; i < dim_num; ++i)
d[i] = shape[i];
return d;
}
// matmul() transpose为true则表示在乘法前先进行转置
static ITensor* trt_matmul(INetworkDefinition *network, ITensor *input1, bool transpose1, ITensor *input2, bool transpose2) {
return network->addMatrixMultiply(*input1, transpose1, *input2, transpose2)->getOutput(0);
}
// 张量切片 后三个参数依次为切片开始时各维度的索引、切片完成后张量的形状以及切片时各维度的步长
static ITensor* trt_slice(INetworkDefinition *network, ITensor *input, Dims start, Dims size, Dims step){
return network->addSlice(*input, start, size, step)->getOutput(0);
}
// view() 不可乱用,张量填充是最后一个维度对应填充,会导致数据混乱,所以一定要确保数据对应后在填充(先做维度变换) 例如:(n,c)->(c,1,n),应该先(n,c)->(c,n),然后(c,n)->(c,1,n)
static ITensor* trt_view(INetworkDefinition *network, ITensor *input, Dims dims){
IShuffleLayer *reshape = network->addShuffle(*input);
reshape->setReshapeDimensions(dims);
return reshape->getOutput(0);
}
// transpose() 维度变换
static ITensor* trt_transpose(INetworkDefinition *network, ITensor *input, Permutation pmt){
IShuffleLayer *premute = network->addShuffle(*input);
premute->setFirstTranspose(pmt);
return premute->getOutput(0);
}
// softmax()
static ITensor* trt_softmax(INetworkDefinition *network, ITensor *input, int axes) {
auto sm = network->addSoftMax(*input);
sm->setAxes(axes);
return sm->getOutput(0);
}
// 对张量做一元运算,如取对数、绝对值、三角函数等
static ITensor* trt_unary(INetworkDefinition *network, ITensor *input, UnaryOperation op) {
return network->addUnary(*input, op)->getOutput(0);
}
// log_softmax()
static ITensor* trt_log_softmax(INetworkDefinition *network, ITensor *input, int axes) {
input = trt_softmax(network, input, axes);
return trt_unary(network, input, UnaryOperation::kLOG);
}
// cat() 多个张量进行拼接,所有张量的dimisions必须相同!即一个2*3*4,其它也得是2*3*4
static ITensor* trt_cat(INetworkDefinition *network, vector<ITensor *> inputs, int axis){
int num_inputs = inputs.size();
ITensor** input = new ITensor*[num_inputs];
for (int i = 0; i < num_inputs; i++)
input[i] = inputs[i];
IConcatenationLayer *cat = network->addConcatenation(input, num_inputs);
cat->setAxis(axis);
return cat->getOutput(0);
}
// T = (scale * T + shift) ^ power T至少三个维度
static ITensor* trt_scale(INetworkDefinition *network, ITensor *input, ScaleMode sm, Weights &shift, Weights &scale, Weights &power) {
return network->addScale(*input, sm, shift, scale, power)->getOutput(0);
}
// 对张量进行某个维度求最值,求和等操作。 reduceAxes采用的是bitmask表示法,下方max()为一个示例
static ITensor* trt_reduce(INetworkDefinition *network, ITensor *input, uint32_t reduceAxes, ReduceOperation op, bool keepDims){
/* 若TRT的输入是 C * H * W
*reduceAxes W H C
* 0X01 0 0 1
* 0X02 0 1 0
* 0X03 0 1 1
* 0X04 1 0 0
* 0X05 1 0 1
* 0X06 1 1 0
* 0X07 1 1 1
*/
return network->addReduce(*input, op, reduceAxes, keepDims)->getOutput(0);
}
// max()
static ITensor* trt_max(INetworkDefinition *network, ITensor *input, uint32_t reduceAxes, bool keepDims) {
/* 若TRT的输入是 C * H * W
*reduceAxes W H C
* 0X01 0 0 1
* 0X02 0 1 0
* 0X03 0 1 1
* 0X04 1 0 0
* 0X05 1 0 1
* 0X06 1 1 0
* 0X07 1 1 1
*/
return network->addReduce(*input, ReduceOperation::kMAX, reduceAxes, keepDims)->getOutput(0);
}
// 常量层
static ITensor* trt_constant(INetworkDefinition *network, Dims dims, float *wgt, int count){
Weights weights{ DataType::kFLOAT,wgt,count };
return network->addConstant(dims, weights)->getOutput(0);
}
// 激活层
static ITensor* trt_act(INetworkDefinition *network, ITensor *input, ActivationType op) {
return network->addActivation(*input, op)->getOutput(0);
}
// relu()
static ITensor* trt_relu(INetworkDefinition *network, ITensor *input) {
return trt_act(network, input, ActivationType::kRELU);
}
还有很多其它层的写法就不展示了,感兴趣的同学可以看这位大佬的github:tensorRT_Wheels.
2. 创建项目
打开项目属性页,添加所需包含目录、库目录以及链接器输入。
包含目录:C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include
库目录:C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\lib\x64
链接器输入:nvinfer.lib nvinfer_plugin.lib cudart.lib
3. 构建engine(以pointNet为例)
-
使用 C++ API 搭建pointNet网络
模仿pointNet分类的python代码,编写出C++版本的pointNetCla类,以下为部分代码对比:
-
加载权重文件(如不了解wts文件结构,可参考上篇文章)
std::map<std::string, Weights> loadWeights(const std::string file) {
std::cout << "加载权重文件: " << file << std::endl;
std::map<std::string, Weights> weightMap;
std::ifstream input(file);
assert(input.is_open() && "无法打开文件,请检查文件路径是否正确!");
int count;
input >> std::dec >> count; // 指定读入格式 dec(十进制)
while (count--){
Weights wt{ DataType::kFLOAT, nullptr, 0 };
uint32_t size;
std::string name;
input >> name >> std::dec >> size;
wt.count = size;
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
input >> std::hex >> val[x]; // 指定读入格式 dec(十六进制)
wt.values = val;
weightMap[name] = wt;
}
return weightMap;
}
- 构建引擎文件
// func 表示要构建的是分类还是分割网络
ICudaEngine* createEngine(string wts_name, string func) {
map<string, Weights> &&weightMap = loadWeights(wts_name);
Logger logger;
IBuilder *builder = createInferBuilder(logger);
IBuilderConfig *config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(B * (1 << 20)); // 这个值设置的要大于网络中最大的中间变量占用的空间,否则可能出现无法找到某个函数的实现的情况
builder->setMaxBatchSize(B);
if (use_fp == "USE_FP16")
config->setFlag(BuilderFlag::kFP16);
INetworkDefinition *network = builder->createNetworkV2(1);
// 点云是三维文件,tensorRT默认第四个维度是batchsize的大小,所以给点云升一维
ITensor* input = network->addInput(inputName, DataType::kFLOAT, Dims4{ B, C, N, 1 });
if (func == "cla") {
shared_ptr<pointNetCla> pn_cla = make_shared<pointNetCla>(network, input, &weightMap, K);
input = pn_cla->forword();
}else if (func == "sem") {
shared_ptr<pointNetSem> pn_sem = make_shared<pointNetSem>(network, input, &weightMap, K);
input = pn_sem->forword();
}
input->setName(outputName);
network->markOutput(*input);
cout << "build engine........" << endl;
auto engine = builder->buildEngineWithConfig(*network, *config);
config->destroy();
network->destroy();
builder->destroy();
return engine;
}
- 引擎文件的保存
void saveEngine(ICudaEngine *engine,string engine_name) {
IHostMemory* modelStream = engine->serialize();
engine->destroy();
ofstream ofs(engine_name, ios::binary);
ofs.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
ofs.close();
cout << "save engine successful!" << endl;
modelStream->destroy();
}