TensorRT C++ 序列化serialize / 反序列化deserialize

序列化Engine

// serialize the engine, then close everything down
const std::string engine_file{ "tensorrt_mnist.trt" };
nvinfer1::IHostMemory* gieModelStream = mEngine->serialize(); // GIE model
fprintf(stdout, "allocate memory size: %d bytes\n", gieModelStream->size());
std::ofstream outfile(engine_file.c_str(), std::ios::out | std::ios::binary);
if (!outfile.is_open()) {
    fprintf(stderr, "fail to open file to write: %s\n", engine_file.c_str());
    return -1;
}
unsigned char* p = (unsigned char*)gieModelStream->data();
outfile.write((char*)p, gieModelStream->size());
outfile.close();

反序列化Engine

// deserialize the engine 
auto runtime = SampleUniquePtr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(gLogger.getTRTLogger()));
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
    runtime->deserializeCudaEngine((const void*)engine_data.get(), size, nullptr),
    samplesCommon::InferDeleter());

注意这里的engine是以智能指针的形式进行定义，官网教程中的反序列化得到的是普通指针，智能指针同sample中的engine全局变量的类型吻合，故使用智能指针。