文章目录
0. 前言
- 尝试使用TensorRT,主要内容包括
- 安装TensorRT
- 运行第一个样例,即sampleMNIST
- 分析样例源码
1. 安装
-
个人喜欢通过tar包安装,整体流程参考官方文档
-
第一步:下载tar包,解压,并在
~/.bashrc
中添加export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<TensorRT-${version}/lib>
-
第二步:构建对应的conda环境,并安装各种whl包。
# 注意,这个分python版本 pip install python/tensorrt-7.2.2.3-cp37-none-linux_x86_64.whl # 下面的不分版本 pip install uff/uff-0.6.9-py2.py3-none-any.whl pip install graphsurgeon/graphsurgeon-0.4.5-py2.py3-none-any.whl pip install onnx_graphsurgeon/onnx_graphsurgeon-0.2.6-py2.py3-none-any.whl
-
第三步:运行一个sample验证安装是否成功。
2. 验证(第一个demo)
-
代码可以在
samples/sampleMNIST
中找到,也可以看 github 中对应路径。- 以下过程可以参考 README.md 相关信息。
- 下文中
./
指的是 tensorrt 所在路径,如~/TensorRT-x.x.x.x
。
1.1. 数据准备
- 在
./data/minst
目录下直接运行python generate_pgms.py
就可以了。- 结果就是生成了一些
*.pgm
- 过程中会下载mnist,需要注意网络情况
- 结果就是生成了一些
1.2. 代码编译与运行
- 第一步:在
./samples/sampleMNIST
目录下执行make
命令。可执行文件生成在./bin/
目录下。 - 第二步:在
./
目录下运行./bin/sample_mnist
即可看到预期结果。
&&&& RUNNING TensorRT.sample_mnist # ./bin/sample_mnist
[06/06/2020-16:26:42] [I] Building and running a GPU inference engine for MNIST
[06/06/2020-16:26:48] [I] [TRT] Detected 1 inputs and 1 output network tensors.
[06/06/2020-16:26:48] [I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@:.%@@@@@
@@@@@@@@@@@@@@@%%%%* .@@@@@@
@@@@@@@@@@@%-::. =- :@@@@@@
@@@@@@@@@@@+ -. *@@@@@@
@@@@@@@@@@= -=@@# @@@@@@@
@@@@@@@@@@= .%@@%- +@@@@@@@
@@@@@@@@@@#. -@%: +@@@@@@@@
@@@@@@@@@@@- .*. +@@@@@@@@@
@@@@@@@@@@@= -%@@@@@@@@@
@@@@@@@@@@@+ :%@@@@@@@@@@
@@@@@@@@@@@- :@@@@@@@@@@@@
@@@@@@@@@@- :@@@@@@@@@@@@
@@@@@@@@%+ :- :@@@@@@@@@@@@
@@@@@@@@* +@* -@@@@@@@@@@@@
@@@@@@@# *@%. =@@@@@@@@@@@@
@@@@@@@= :@@+ +@@@@@@@@@@@@
@@@@@@@= :@* -@@@@@@@@@@@@@
@@@@@@@- -: *@@@@@@@@@@@@@@
@@@@@@@+ =@@@@@@@@@@@@@@@
@@@@@@@@+ :+@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
[06/06/2020-16:26:48] [I] Output:
0:
1:
2:
3:
4:
5:
6:
7:
8: **********
9:
&&&& PASSED TensorRT.sample_mnist # ./bin/sample_mnist
3. 样例源码解析
3.1. 基本概念
- 在几乎所有TensorRT MNIST笔记的文章中都有一张图:
- 这张图来自 TensorRT官方文档,用于介绍TensorRT的基本流程,也就是下面源码的基本流程。
- 第一步:将训练好的神经网络模型转换为TensorRT的形式,并用TensorRT Optimizer进行优化。
- 第二步:将在TensorRT Engine中运行优化好的TensorRT网络结构。
- 这张图来自 TensorRT官方文档,用于介绍TensorRT的基本流程,也就是下面源码的基本流程。
3.2. 主函数
- 主要作用:
- 解析输入参数。
- 构造
SampleMNIST
对象,调用相关方法实现Caffe模型转换、TensorRT Engine推理。
int main(int argc, char** argv)
{
// 使用GNU C提供的 getopt_long 获取参数
// 如果所有参数都合法,则返回 true,否则返回 false
samplesCommon::Args args;
bool argsOK = samplesCommon::parseArgs(args, argc, argv);
if (!argsOK)
{
gLogError << "Invalid arguments" << std::endl;
printHelpInfo();
return EXIT_FAILURE;
}
// 如果在输入参数中指定了 help,那就返回帮助文档
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
// 准备工作,如参数初始化、logger
auto sampleTest = gLogger.defineTest(gSampleName, argc, argv);
gLogger.reportTestStart(sampleTest);
samplesCommon::CaffeSampleParams params = initializeSampleParams(args);
// 样例主要就是通过 SampleMNIST 类来实现,主要包括build、infer、teardown三个阶段
SampleMNIST sample(params);
gLogInfo << "Building and running a GPU inference engine for MNIST" << std::endl;
if (!sample.build())
{
return gLogger.reportFail(sampleTest);
}
if (!sample.infer())
{
return gLogger.reportFail(sampleTest);
}
// 结束工作流,释放资源,返回结果
if (!sample.teardown())
{
return gLogger.reportFail(sampleTest);
}
return gLogger.reportPass(sampleTest);
}
3.3. 将caffe模型转换为TensorRT可识别的形式
- 实现方式:通过主函数中的
sample.build()
实现。 build
函数基本流程- 主要工作就是,定义一个
network
对象,用于保存 caffe 模型转换后的结果。- 通过
constructNetwork
实现。
- 通过
- 猜测所谓的 Tensor Optimizer 就是在
builder->buildEngineWithConfig
中实现的。
- 主要工作就是,定义一个
bool SampleMNIST::build()
{
auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(gLogger.getTRTLogger()));
if (!builder)
return false;
// 创建空的network,后面 constructNetwork 中会定义
auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetwork());
if (!network) return false;
auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
if (!config) return false;
// 创建caffe模型的parser,在constructNetwork函数中解析模型,转换为 network
auto parser = SampleUniquePtr<nvcaffeparser1::ICaffeParser>(nvcaffeparser1::createCaffeParser());
if (!parser) return false;
// 解析 caffe 模型,并转换为 network 形式
constructNetwork(parser, network);
// 设置一些参数
builder->setMaxBatchSize(mParams.batchSize);
config->setMaxWorkspaceSize(16_MiB);
config->setFlag(BuilderFlag::kGPU_FALLBACK);
config->setFlag(BuilderFlag::kSTRICT_TYPES);
if (mParams.fp16) config->setFlag(BuilderFlag::kFP16);
if (mParams.int8) config->setFlag(BuilderFlag::kINT8);
samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);
// 构建 tensorrt 引擎
// 注意,这个是成员变量
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config), samplesCommon::InferDeleter());
if (!mEngine) return false;
assert(network->getNbInputs() == 1);
mInputDims = network->getInput(0)->getDimensions();
assert(mInputDims.nbDims == 3);
return true;
}
- caffe模型转换具体实现过程
- 其实模型转换本身,
parser->parse
一个函数就解决了。 - 下面代码的大量篇幅是在:在模型开头添加
输入图片减去平均数
操作上。
- 其实模型转换本身,
void SampleMNIST::constructNetwork(SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser, SampleUniquePtr<nvinfer1::INetworkDefinition>& network)
{
// 解析 caffe 的模型文件与权重文件,将结果写入 network 中
const nvcaffeparser1::IBlobNameToTensor* blobNameToTensor = parser->parse(
mParams.prototxtFileName.c_str(),
mParams.weightsFileName.c_str(),
*network,
nvinfer1::DataType::kFLOAT);
// 标记模型输出
for (auto& s : mParams.outputTensorNames)
{
network->markOutput(*blobNameToTensor->find(s.c_str()));
}
// 在模型开头添加 `输入图片减去平均数` 操作
// add mean subtraction to the beginning of the network
nvinfer1::Dims inputDims = network->getInput(0)->getDimensions();
mMeanBlob = SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob>(parser->parseBinaryProto(mParams.meanFileName.c_str()));
nvinfer1::Weights meanWeights{nvinfer1::DataType::kFLOAT, mMeanBlob->getData(), inputDims.d[1] * inputDims.d[2]};
// For this sample, a large range based on the mean data is chosen and applied to the head of the network.
// After the mean subtraction occurs, the range is expected to be between -127 and 127, so the rest of the network
// is given a generic range.
// The preferred method is use scales computed based on a representative data set
// and apply each one individually based on the tensor. The range here is large enough for the
// network, but is chosen for example purposes only.
float maxMean = samplesCommon::getMaxValue(static_cast<const float*>(meanWeights.values), samplesCommon::volume(inputDims));
// 模型中添加常量(图片channel均值)
auto mean = network->addConstant(nvinfer1::Dims3(1, inputDims.d[1], inputDims.d[2]), meanWeights);
mean->getOutput(0)->setDynamicRange(-maxMean, maxMean);
network->getInput(0)->setDynamicRange(-maxMean, maxMean);
// 添加 减均值 操作
auto meanSub = network->addElementWise(*network->getInput(0), *mean->getOutput(0), ElementWiseOperation::kSUB);
meanSub->getOutput(0)->setDynamicRange(-maxMean, maxMean);
network->getLayer(0)->setInput(0, *meanSub->getOutput(0));
// 执行缩放,输出结果为 [-1, 1]
samplesCommon::setAllTensorScales(network.get(), 127.0f, 127.0f);
}
3.4. 模型推理
- 主要工作:就是将转换好的模型在tensorrt engine上跑一边。
- 里面用到的各种东西比较多,目前也看不懂。
- 主要通过
infer()
函数完成,本函数的主要操作就是:- 读取输入数据(
processInput
)。 - 通过 cuda stream / buffer 等进行推理。
- 判断输出结果是否正确(
verifyOutput
)。
- 读取输入数据(
bool SampleMNIST::infer()
{
// 实现具体推理过程
// 缓存对象管理
// Create RAII buffer manager object
samplesCommon::BufferManager buffers(mEngine, mParams.batchSize);
// 创建上下文
auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
return false;
}
// 随机选择一个数字
// Pick a random digit to try to infer
srand(time(NULL));
const int digit = rand() % 10;
// 读取输入数据到缓存对象中
// 即将 digit 写入 buffers 中,名字为 mParams.inputTensorNames[0]
// Read the input data into the managed buffers
// There should be just 1 input tensor
assert(mParams.inputTensorNames.size() == 1);
if (!processInput(buffers, mParams.inputTensorNames[0], digit))
{
return false;
}
// 创建 cuda 流,准备执行推理
// Create CUDA stream for the execution of this inference.
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// 异步将数据从主机输入缓冲区(buffer)复制到设备输入缓冲区(stream)
// Asynchronously copy data from host input buffers to device input buffers
buffers.copyInputToDeviceAsync(stream);
// 异步将推理任务加入队列中
// Asynchronously enqueue the inference work
if (!context->enqueue(mParams.batchSize, buffers.getDeviceBindings().data(), stream, nullptr))
{
return false;
}
// 异步将模型结果从设备(stream)保存到主机缓冲区(buffers)
// Asynchronously copy data from device output buffers to host output buffers
buffers.copyOutputToHostAsync(stream);
// 等待工作结束,关闭stream
// Wait for the work in the stream to complete
cudaStreamSynchronize(stream);
// Release stream
cudaStreamDestroy(stream);
// 得到结果,判断结果是否准确
// 即从 buffer 中获取名为 mParams.outputTensorNames[0] 的结果,判断与digit是否相同
// Check and print the output of the inference
// There should be just one output tensor
assert(mParams.outputTensorNames.size() == 1);
bool outputCorrect = verifyOutput(buffers, mParams.outputTensorNames[0], digit);
return outputCorrect;
}
- 读取输入数据
- 这部分没啥好说的。
bool SampleMNIST::processInput(const samplesCommon::BufferManager& buffers, const std::string& inputTensorName, int inputFileIdx) const
{
const int inputH = mInputDims.d[1];
const int inputW = mInputDims.d[2];
// Read a random digit file
srand(unsigned(time(nullptr)));
std::vector<uint8_t> fileData(inputH * inputW);
readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);
// Print ASCII representation of digit
gLogInfo << "Input:\n";
for (int i = 0; i < inputH * inputW; i++)
{
gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
}
gLogInfo << std::endl;
float* hostInputBuffer = static_cast<float*>(buffers.getHostBuffer(inputTensorName));
for (int i = 0; i < inputH * inputW; i++)
{
hostInputBuffer[i] = float(fileData[i]);
}
return true;
}
- 验证输出结果
- 还会输出可视化结果,有种梦回当年的感觉。
bool SampleMNIST::verifyOutput(const samplesCommon::BufferManager& buffers, const std::string& outputTensorName, int groundTruthDigit) const
{
// 获取 host buffer 中的输出tensor数值
// 应该是10个数字的概率
const float* prob = static_cast<const float*>(buffers.getHostBuffer(outputTensorName));
// Print histogram of the output distribution
gLogInfo << "Output:\n";
float val{0.0f};
int idx{0};
const int kDIGITS = 10;
for (int i = 0; i < kDIGITS; i++)
{
if (val < prob[i])
{
val = prob[i];
idx = i;
}
gLogInfo << i << ": " << std::string(int(std::floor(prob[i] * 10 + 0.5f)), '*') << "\n";
}
gLogInfo << std::endl;
return (idx == groundTruthDigit && val > 0.9f);
}