개인적으로 읽고 쓰는 공부용 리뷰입니다.
틀린 점이 있을 수도 있으니 감안하고 읽어주세요. 피드백은 댓글로 부탁드립니다.
[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할수있어서 오버로딩해놨다.
끝
