代码见:
https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt
代码整体的一个流程如下:
创建InferBuilder对象——创建Network对象——创建BuilderConfig对象——创建OptimizationProfile对象——创建OptimizationProfile对象——设置输入维度——创建解析器对象——解析ONNX模型文件——使用BuilderConfig对象构建TensorRT引擎——序列化引擎并保存到文件——创建执行上下文
代码部分在featuretensor.cpp文件中
void FeatureTensor::loadOnnx(std::string onnxPath)
{
// 创建InferBuilder对象
auto builder = createInferBuilder(*gLogger);
assert(builder != nullptr);
// 创建Network对象
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto network = builder->createNetworkV2(explicitBatch);
assert(network != nullptr);
// 创建BuilderConfig对象
auto config = builder->createBuilderConfig();
assert(config != nullptr);
// 创建OptimizationProfile对象
auto profile = builder->createOptimizationProfile();
Dims dims = Dims4{1, 3, imgShape.height, imgShape.width};
// 设置输入维度
profile->setDimensions(inputName.c_str(),
OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]});
profile->setDimensions(inputName.c_str(),
OptProfileSelector::kOPT, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]});
profile->setDimensions(inputName.c_str(),
OptProfileSelector::kMAX, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]});
config->addOptimizationProfile(profile);
// 创建解析器对象
nvonnxparser::IParser *parser = nvonnxparser::createParser(*network, *gLogger);
assert(parser != nullptr);
// 解析ONNX模型文件
auto parsed = parser->parseFromFile(onnxPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING));
assert(parsed);
// 设置最大工作空间大小
config->setMaxWorkspaceSize(1 << 20);
// 使用BuilderConfig对象构建TensorRT引擎
engine = builder->buildEngineWithConfig(*network, *config);
assert(engine != nullptr);
// 序列化引擎并保存到文件
std::string enginePath = "./deepsort.engine";
std::ofstream engineFile(enginePath, std::ios::binary);
if (engineFile)
{
nvinfer1::IHostMemory *serializedEngine = engine->serialize();
engineFile.write(static_cast<const char *>(serializedEngine->data()), serializedEngine->size());
serializedEngine->destroy();
engineFile.close();
}
else
{
std::cerr << "Failed to open engine file for writing: " << enginePath << std::endl;
}
// 创建执行上下文
context = engine->createExecutionContext();
assert(context != nullptr);
initResource();
}
上面代码我新增了读取onnx帮保存为engine的代码,这样就不用每次重新生成了。但是换硬件平台的时候,还是需要重新生成。
单独保存engine的代码如下:
// 检查引擎是否已创建
assert(engine != nullptr);
// 序列化引擎并保存到文件
std::string enginePath = "./deepsort.engine";
std::ofstream engineFile(enginePath, std::ios::binary);
if (engineFile) {
nvinfer1::IHostMemory* serializedEngine = engine->serialize();
engineFile.write(static_cast<const char*>(serializedEngine->data()), serializedEngine->size());
serializedEngine->destroy();
engineFile.close();
} else {
std::cerr << "Failed to open engine file for writing: " << enginePath << std::endl;
}