[TensorRT] pytorch network를 c++에서 inference

* Enviornment

Jetson Orin NX

Jetpack 5.1 [L4T 35.2.1]

Ubuntu 20.04

cuda 11.4

 

1. python에서 network를 onnx로 변환

# python
neural_network = your_network().to(device)

neural_network.load_state_dict(torch.load("your_weight"))
x = torch.randn("your_input_shape").to(device)
    
torch.onnx.export(neural_network,            
                  x,                         
                  "your_model_path.onnx",   
                  export_params=True,        
                  opset_version=10,          
                  do_constant_folding=True,  
                  input_names = ['input'],   
                  output_names = ['output'], 
                  dynamic_axes={'input' : {0 : 'batch_size'},    
                                'output' : {0 : 'batch_size'}})

 

 

 

2. c++에서 불러와서 실행

 

2.1 onnx를 업로드하는 경우

2.1.1 필요한 함수 및 변수 선언

//tensorRT
using namespace nvinfer1;
void* g_deviceInput;
void* g_deviceOutput;
class Logger : public nvinfer1::ILogger
{
public:
    void log(Severity severity, const char* msg) noexcept override
    {
        if (severity != Severity::kINFO)
        {
            // std::cout << "TensorRT Logger: " << msg << std::endl;
        }
    }
};
static Logger gLogger;

ICudaEngine* createEngine(const std::string& onnxModelPath, int maxBatchSize)
{
    IBuilder* builder = createInferBuilder(gLogger);
    INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);

    if (!parser->parseFromFile(onnxModelPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING)))
    {
        return nullptr;
    }

    builder->setMaxBatchSize(maxBatchSize);
   
    nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
    profile->setDimensions("input", OptProfileSelector::kMIN, Dims4(1, 1, 480, 640));

    profile->setDimensions("input", OptProfileSelector::kOPT, Dims4(1, 1, 480, 640));
    profile->setDimensions("input", OptProfileSelector::kMAX, Dims4(1, 1, 480, 640));



   
    IBuilderConfig* buildConfig = builder->createBuilderConfig();
    buildConfig->setMaxWorkspaceSize(1 << 30);
    buildConfig->addOptimizationProfile(profile);
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *buildConfig);

    if (!engine)
    {
        std::cout<<"Error occuered while building engine"<<std::endl;
        return nullptr;
    }

    parser->destroy();
    network->destroy();
    

    return engine;
}
bool doInference(IExecutionContext* context, cv::Mat& input, cv::Mat& output,size_t inputSize,size_t outputSize)
{
    
    cudaMemcpy(g_deviceInput, input.data, inputSize, cudaMemcpyHostToDevice);

    void* buffers[] = {g_deviceInput, g_deviceOutput};
    context->executeV2(buffers);

    cudaMemcpy(output.data, g_deviceOutput, outputSize, cudaMemcpyDeviceToHost);

    return true;
}

 

 

2.1.2 onnx 불러와서 실행

- input 사이즈는 1x480x640 이라고 가정

nvinfer1::ICudaEngine* engine = createEngine("your model path", 1);

IExecutionContext* context = engine->createExecutionContext();

cv::Mat temp_mat(480, 640, CV_32FC1);

size_t inputSize = temp_mat.total() * temp_mat.elemSize();
size_t outputSize = temp_mat.total() * temp_mat.elemSize();

cudaMalloc(&g_deviceInput, inputSize);
cudaMalloc(&g_deviceOutput, outputSize);

 

2.1.3 inference 수행

if (doInference(context, infer_input, infer_output, inputSize, outputSize))
        {
            std::cout<<"Success"<<std::endl;
        }
        else
        {
            std::cout<<"Fail"<<std::endl;
        }

2.2 trt 파일을 불러오는 경우

 

2.2.1 trt 파일 변환

 

trtexec --onnx=my_model.onnx --fp16 --workspace=1024 --saveEngine=my_model.trt

 

 

2.2.2 c++에서 불러오기

IRuntime* runtime = createInferRuntime(gLogger);

    std::ifstream engineFile(onnxModelPath, std::ios::binary);
    if (!engineFile) {
        std::cerr << "Failed to open engine file." << std::endl;
        return 1;
    }
    engineFile.seekg(0, engineFile.end);
    size_t engineSize = engineFile.tellg();
    engineFile.seekg(0, engineFile.beg);

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

    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), engineSize, nullptr);
    IExecutionContext* context = engine->createExecutionContext();

 

이후로는 2.1.3, 2.1.4를 진행하면 됨

 

 

* Reference

 

 

https://developer.nvidia.com/blog/speed-up-inference-tensorrt/

How to Speed Up Deep Learning Inference Using TensorRT | NVIDIA Technical Blog

 

https://eehoeskrap.tistory.com/414

[TensorRT] NVIDIA TensorRT 개념, 설치방법, 사용하기

 

https://learnopencv.com/how-to-run-inference-using-tensorrt-c-api/

How To Run Inference Using TensorRT C++ API | LearnOpenCV

 

https://discuss.pytorch.kr/t/tensorrt-c/1874/3

TensorRT C++ 에서 엔진파일 로딩하여 추론하기

 

https://da2so.tistory.com/60

[NVIDIA] TensorRT plugin 사용 및 예제 (feat. yolov7)

 

https://blog.si-analytics.ai/33

Pytorch와 TensorRT를 이용한 딥 러닝 추론 최적화