[TensorRT] 2. Serializing & Deserializing tensorrt engine (tensorRT 7.2.3)

개인적으로 읽고 쓰는 공부용 리뷰입니다.

틀린 점이 있을 수도 있으니 감안하고 읽어주세요. 피드백은 댓글로 부탁드립니다.

[TensorRT] 1. Build tensorrt engine (tensorRT 7.2.3)

---

Serializie는 나중에 재사용을 위해 저장하기위한 포맷으로 바꾸는 것을 의미한다. inference에 사용하기 위해서는 그냥 deserializie한 뒤 쓰면 된다. 보통 빌드과정이 시간을 소요하기 때문에 매번 빌드하는 것을 피하기 위해 이 과정을 한다.

// code for serializing ICudaEngine
IHostMemory *serializedModel = engine->serialize();
// store model to disk
// <…>
 serializedModel->destroy();

-

Deserialize code, The final argument is a plugin layer factory for applications using custom layers. For more information, see [Extending TensorRT With Custom Layers]
별거아니고, 마지막에 nullptr는 iplugin for custom layer인데 없다면 그냥 nullptr넣으면 된다.

// code for deserializing ICudaEngine
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(modelData, modelSize, nullptr)

---

주의 할점은 trt version, gpu, platform을 항시 잘 체크해야한다.

• Serialized engines are not portable across platforms or TensorRT versions.
• Engines are specific to the exact GPU model they were built on.

 

위에는 tensorrt reference의 코드인데 처음에보고 어떻게 더 추가해야할지 몰라서 막막했다.
아래는 실제로 내가 사용하는 serializing & save 코드다.

여기서 engine_ 은 빌드가 성공적으로 된 ICudaEngine이다.

bool saveEngine( std::string &fileName ) const 
{
    std::ofstream engineFile( fileName, std::ios::binary );
    if ( !engineFile ) {
        gLogFatal << "Cannot open engine file : " << fileName << std::endl;
        return false;
    }

    if ( engine_ == nullptr ) {
        gLogError << "Engine is not defined" << std::endl;
        return false;
    }
    nvinfer1::IHostMemory *serializedEngine{engine_->serialize()};
    if ( serializedEngine == nullptr ) {
        gLogError << "Engine serialization failed" << std::endl;
        return false;
    }

    engineFile.write( static_cast<char *>( serializedEngine->data() ),
                      serializedEngine->size() );
    if ( engineFile.fail() ) {
        gLogError << "Failed to save Engine." << std::endl;
        return false;
    }
    std::cout << "Successfully save to : " << fileName << std::endl;
    return true;
}

 

다음은 저장된 ICudaEngine을 load후 다시 deserializing하는 코드다. 

bool Load( const std::string &fileName ) {
    std::ifstream engineFile( fileName, std::ios::binary );
    if ( !engineFile ) {
        std::cout << "can not open file : " << fileName << std::endl;
        return false;
    }
    engineFile.seekg( 0, engineFile.end );
    auto fsize = engineFile.tellg();
    engineFile.seekg( 0, engineFile.beg );

    std::vector<char> engineData( fsize );
    engineFile.read( engineData.data(), fsize );

    Load( engineData.data(), ( long int )fsize );
    return true;
}

bool Load( const void *engineData, const long int fsize ) {
    nvinfer1::IRuntime *runtime = nvinfer1::createInferRuntime( gLogger.getTRTLogger() );
    engine_ = runtime->deserializeCudaEngine( engineData, fsize, nullptr );
    // if u want DLA core setting, then u shoud write code here
    runtime->destroy();
    return true;
}

두가지 방법으로 load할수있어서 오버로딩해놨다.

끝