TensorRT 10버전 이후 변환된 모델을 C++에서 Inference 수행 방법(enqueueV3 함수 사용법)

 

※본 글을 작성했던 2025년 8월 당시 TensorRT 코드 작성을 GPT 4o에가 요청시 TensorRT 8버전을 기준으로 C++에서 Inference 수행시 'enqueueV2'를 사용하는 것으로만 코드 작성을 해줘서 온갖 고난 끝에 참고자료 페이지의 자료를 찾아내 enqueueV3() 함수 사용 방법을 작성하였습니다. 그런데 GPT5가 2024년 10월 1일까지의 자료를 학습하게 됨으로서 지금은 TensorRT 10 버전 기준으로 enqueueV3 함수를 잘 작성해줍니다.

 

 TensorRT 10 버전이 등장하면서 변환된 모델을 C++에서 Inference 수행시 enqueueV3() 함수를 사용하여야 하며, 기존 8버전에서 사용했던 enqueueV2() 함수를 더이상 사용되지 않게 되었습니다. TensorRT 버전 변경에 따라 코드 변경 사항이 많아져 방법을 찾다가 아래 링크의 참고자료의 내용대로 코드를 작성하니 TensorRT 10  기준으로 변환된 AI모델이 잘 동작하는 것을 확인할 수 있었습니다.

 

 아래의 소스코드는 퍼플렉시티AI의 도움으로 소스코드 설명을 주석으로 작성하였습니다. 아래의 흐름대로 코드를 작성하시면 TensorRT 10버전 이후 enqueueV3() 함수를 작성하는데 큰 어려움은 없을 것입니다.

 

#include <fstream>
#include <iostream>
#include <memory>
#include <vector>
#include <string>
#include <opencv2/opencv.hpp>
#include <NvInfer.h>
#include <cuda_runtime_api.h>
 
// CUDA 에러 체크 매크로 및 함수
#define CHECK_CUDA_ERROR(val) check((val), #val, __FILE__, __LINE__)
void check(cudaError_t err, const char* const func, const char* const file, const int line) {
    if (err != cudaSuccess) {
        std::cerr << "CUDA Runtime Error at: " << file << ":" << line << std::endl;
        std::cerr << cudaGetErrorString(err) << " " << func << std::endl;
        std::exit(EXIT_FAILURE);
    }
}
 
// TensorRT 로그 출력을 위한 커스텀 Logger 클래스
class CustomLogger : public nvinfer1::ILogger {
    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override {
        // info 이하 레벨만 출력
        if (severity <= nvinfer1::ILogger::Severity::kINFO) {
            std::cout << msg << std::endl;
        }
    }
};
 
// TensorRT 객체 자동 삭제를 위한 Deleter 구조체
struct InferDeleter {
    template <typename T>
    void operator()(T* obj) const {
        delete obj;
    }
};
 
// Pascal VOC 21 클래스 색상 팔레트 (B,G,R 순서, OpenCV는 BGR)
cv::Vec3b getPascalVOCColor(int class_id) {
    static const std::vector<cv::Vec3b> palette = {
        {  0,   0,   0}, // 0: background (black)
        {128,   0,   0}, // 1: aeroplane
        {  0,128,   0}, // 2: bicycle
        {128,128,   0}, // 3: bird
        {  0,   0,128}, // 4: boat
        {128,   0,128}, // 5: bottle
        {  0,128,128}, // 6: bus
        {128,128,128}, // 7: car
        { 64,   0,   0}, // 8: cat
        {192,   0,   0}, // 9: chair
        { 64,128,   0}, // 10: cow
        {192,128,   0}, // 11: diningtable
        { 64,   0,128}, // 12: dog
        {192,   0,128}, // 13: horse
        { 64,128,128}, // 14: motorbike
        {192,128,128}, // 15: person
        {  0, 64,   0}, // 16: potted plant
        {128, 64,   0}, // 17: sheep
        {  0,192,   0}, // 18: sofa
        {128,192,   0}, // 19: train
        {  0, 64,128}  // 20: tv/monitor
    };
    if (class_id < 0 || class_id >= palette.size()) return {0,0,0};
    return palette[class_id];
}
 
int main(int argc, char** argv) {
    // 1. 입력 이미지 경로 및 엔진 파일 경로 지정
    std::string image_path = "input.jpg"; // 입력 이미지 파일명
    std::string engine_file_path = "end2end.engine"; // TensorRT 엔진 파일명
 
    // 2. OpenCV로 이미지 읽기 (BGR)
    cv::Mat img = cv::imread(image_path, cv::IMREAD_COLOR);
    if (img.empty()) {
        std::cerr << "이미지를 읽을 수 없습니다: " << image_path << std::endl;
        return EXIT_FAILURE;
    }
    // RGB로 변환 (딥러닝 모델은 보통 RGB 입력)
    cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
 
    // 3. TensorRT 엔진 역직렬화 및 실행 컨텍스트 생성
    CustomLogger logger{};
    std::ifstream engine_file(engine_file_path, std::ios::binary);
    if (!engine_file) {
        std::cerr << "엔진 파일을 열 수 없습니다: " << engine_file_path << std::endl;
        return EXIT_FAILURE;
    }
 
    // TensorRT 모델을 메모리에 로드
    engine_file.seekg(0, std::ios::end);
    size_t engine_file_size = static_cast<size_t>(engine_file.tellg());
    engine_file.seekg(0, std::ios::beg);
    std::unique_ptr<char[]> engine_data(new char[engine_file_size]);
    engine_file.read(engine_data.get(), engine_file_size);
 
    std::unique_ptr<nvinfer1::IRuntime, InferDeleter> runtime{nvinfer1::createInferRuntime(logger)};
    std::unique_ptr<nvinfer1::ICudaEngine, InferDeleter> engine{
        runtime->deserializeCudaEngine(engine_data.get(), engine_file_size)};
    std::unique_ptr<nvinfer1::IExecutionContext, InferDeleter> context{
        engine->createExecutionContext()};
 
    // 4. 입력/출력 텐서 이름, shape, dtype 확인
    // (엔진에 따라 다를 수 있으니 반드시 확인 필요)
    const char* input_tensor_name = engine->getIOTensorName(0);   // 입력 텐서 이름
    const char* output_tensor_name = engine->getIOTensorName(1);  // 출력 텐서 이름
 
    nvinfer1::Dims input_dims = engine->getTensorShape(input_tensor_name);   // 예: {1, 3, H, W}
    nvinfer1::Dims output_dims = engine->getTensorShape(output_tensor_name); // 예: {1, H, W}
 
    int input_batch = input_dims.d[0];
    int input_channels = input_dims.d[1];
    int input_height = input_dims.d[2];
    int input_width = input_dims.d[3];
 
    int output_batch = output_dims.d[0];
    int output_height = output_dims.d[1];
    int output_width = output_dims.d[2];
 
    // 5. 입력 이미지 전처리 (리사이즈, 정규화, NCHW 변환)
    cv::Mat resized;
    cv::resize(img, resized, cv::Size(input_width, input_height));
    resized.convertTo(resized, CV_32FC3, 1.0 / 255.0); // 0~1 사이의 값으로 Normalization, AI모델이 학습 수행시 설정하였던 값으로 설정
 
    // NCHW로 변환 (OpenCV는 HWC, TensorRT는 NCHW)
    std::vector<float> input_tensor(input_channels * input_height * input_width);
    std::vector<cv::Mat> rgb_channels(input_channels);
    for (int i = 0; i < input_channels; ++i)
        rgb_channels[i] = cv::Mat(input_height, input_width, CV_32FC1, input_tensor.data() + i * input_height * input_width);
    cv::split(resized, rgb_channels);
 
    // 6. 입력/출력 버퍼 할당 (CUDA)
    void* input_device_buffer = nullptr;
    size_t input_bytes = input_tensor.size() * sizeof(float);
    CHECK_CUDA_ERROR(cudaMalloc(&input_device_buffer, input_bytes));
 
    void* output_device_buffer = nullptr;
    size_t output_bytes = output_height * output_width * sizeof(int64_t); // int64 클래스 ID
    CHECK_CUDA_ERROR(cudaMalloc(&output_device_buffer, output_bytes));
 
    // 7. TensorRT 실행 컨텍스트에 입력 버퍼 바인딩
    context->setTensorAddress(input_tensor_name, input_device_buffer);
 
    // 8. 실제 출력 텐서 shape와 타입 확인
    nvinfer1::Dims output_dims = context->getTensorShape(output_tensor_name); // 실제 shape
    size_t output_size = 1;
    for (int i = 0; i < output_dims.nbDims; ++i) {
        output_size *= output_dims.d[i];
    }
 
    nvinfer1::DataType output_dtype = engine->getTensorDataType(output_tensor_name);
 
    // 9. 출력 버퍼 할당 및 추론 실행
    void* output_device_buffer = nullptr;
    size_t output_bytes = 0;
    cudaStream_t stream;
    CHECK_CUDA_ERROR(cudaStreamCreate(&stream));
 
    // 10. 출력 타입에 따라 분기 처리
    if (output_dtype == nvinfer1::DataType::kINT32) {
        output_bytes = output_size * sizeof(int32_t);
        CHECK_CUDA_ERROR(cudaMalloc(&output_device_buffer, output_bytes));
        context->setTensorAddress(output_tensor_name, output_device_buffer);
 
        // 추론 실행
        bool status = context->enqueueV3(stream);
        if (!status) {
            std::cerr << "TensorRT 추론 실행에 실패했습니다." << std::endl;
            return EXIT_FAILURE;
        }
        CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
 
        // 결과 복사 및 마스킹
        std::vector<int32_t> output_tensor(output_size);
        CHECK_CUDA_ERROR(cudaMemcpy(output_tensor.data(), output_device_buffer, output_bytes, cudaMemcpyDeviceToHost));
 
        // 마스크 생성
        int mask_height = output_dims.d[output_dims.nbDims - 2];
        int mask_width  = output_dims.d[output_dims.nbDims - 1];
        cv::Mat mask(mask_height, mask_width, CV_8UC3);
        for (int y = 0; y < mask_height; ++y) {
            for (int x = 0; x < mask_width; ++x) {
                int class_id = static_cast<int>(output_tensor[y * mask_width + x]);
                mask.at<cv::Vec3b>(y, x) = getPascalVOCColor(class_id);
            }
        }
        cv::Mat mask_resized;
        cv::resize(mask, mask_resized, img.size(), 0, 0, cv::INTER_NEAREST);
        cv::Mat blended;
        cv::addWeighted(img, 0.6, mask_resized, 0.4, 0, blended);
        cv::cvtColor(blended, blended, cv::COLOR_RGB2BGR);
        cv::imwrite("output_masked_voc.jpg", blended);
        std::cout << "INT32 타입 결과를 output_masked_voc.jpg로 저장했습니다." << std::endl;
    }
    else if (output_dtype == nvinfer1::DataType::kINT64) {
        output_bytes = output_size * sizeof(int64_t);
        CHECK_CUDA_ERROR(cudaMalloc(&output_device_buffer, output_bytes));
        context->setTensorAddress(output_tensor_name, output_device_buffer);
 
        // 추론 실행
        bool status = context->enqueueV3(stream);
        if (!status) {
            std::cerr << "TensorRT 추론 실행에 실패했습니다." << std::endl;
            return EXIT_FAILURE;
        }
        CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
 
        // 결과 복사 및 마스킹
        std::vector<int64_t> output_tensor(output_size);
        CHECK_CUDA_ERROR(cudaMemcpy(output_tensor.data(), output_device_buffer, output_bytes, cudaMemcpyDeviceToHost));
 
        int mask_height = output_dims.d[output_dims.nbDims - 2];
        int mask_width  = output_dims.d[output_dims.nbDims - 1];
        cv::Mat mask(mask_height, mask_width, CV_8UC3);
        for (int y = 0; y < mask_height; ++y) {
            for (int x = 0; x < mask_width; ++x) {
                int class_id = static_cast<int>(output_tensor[y * mask_width + x]);
                mask.at<cv::Vec3b>(y, x) = getPascalVOCColor(class_id);
            }
        }
        cv::Mat mask_resized;
        cv::resize(mask, mask_resized, img.size(), 0, 0, cv::INTER_NEAREST);
        cv::Mat blended;
        cv::addWeighted(img, 0.6, mask_resized, 0.4, 0, blended);
        cv::cvtColor(blended, blended, cv::COLOR_RGB2BGR);
        cv::imwrite("output_masked_voc.jpg", blended);
        std::cout << "INT64 타입 결과를 output_masked_voc.jpg로 저장했습니다." << std::endl;
    }
    else if (output_dtype == nvinfer1::DataType::kFLOAT) {
        output_bytes = output_size * sizeof(float);
        CHECK_CUDA_ERROR(cudaMalloc(&output_device_buffer, output_bytes));
        context->setTensorAddress(output_tensor_name, output_device_buffer);
 
        // 추론 실행
        bool status = context->enqueueV3(stream);
        if (!status) {
            std::cerr << "TensorRT 추론 실행에 실패했습니다." << std::endl;
            return EXIT_FAILURE;
        }
        CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));
 
        // 결과 복사 및 마스킹
        std::vector<float> output_tensor(output_size);
        CHECK_CUDA_ERROR(cudaMemcpy(output_tensor.data(), output_device_buffer, output_bytes, cudaMemcpyDeviceToHost));
 
        int mask_height = output_dims.d[output_dims.nbDims - 2];
        int mask_width  = output_dims.d[output_dims.nbDims - 1];
        cv::Mat mask(mask_height, mask_width, CV_8UC3);
        for (int y = 0; y < mask_height; ++y) {
            for (int x = 0; x < mask_width; ++x) {
                int class_id = static_cast<int>(output_tensor[y * mask_width + x]);
                mask.at<cv::Vec3b>(y, x) = getPascalVOCColor(class_id);
            }
        }
        cv::Mat mask_resized;
        cv::resize(mask, mask_resized, img.size(), 0, 0, cv::INTER_NEAREST);
        cv::Mat blended;
        cv::addWeighted(img, 0.6, mask_resized, 0.4, 0, blended);
        cv::cvtColor(blended, blended, cv::COLOR_RGB2BGR);
        cv::imwrite("output_masked_voc.jpg", blended);
        std::cout << "FLOAT 타입 결과를 output_masked_voc.jpg로 저장했습니다." << std::endl;
    }
    else {
        std::cerr << "지원하지 않는 출력 데이터 타입입니다." << std::endl;
        return EXIT_FAILURE;
    }
 
    // 11. 자원 해제
    CHECK_CUDA_ERROR(cudaFree(input_device_buffer));
    if (output_device_buffer) CHECK_CUDA_ERROR(cudaFree(output_device_buffer));
    CHECK_CUDA_ERROR(cudaStreamDestroy(stream));
 
    return 0;
}

 

 

 

참고자료: https://leimao.github.io/blog/TensorRT-Custom-Plugin-Example/

TensorRT Custom Plugin Example