build_engine_tool.cpp

#include <iostream>
#include <fstream>
#include <memory>
#include <chrono>
#include <NvInfer.h>
#include <NvOnnxParser.h>

class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING) {
            std::cout << "[TensorRT] " << msg << std::endl;
        }
    }
};

bool buildEngine(const std::string& onnxPath, const std::string& enginePath, bool enableFP16 = false) {
    Logger logger;
    
    // Create builder
    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
    if (!builder) {
        std::cerr << "Failed to create TensorRT builder" << std::endl;
        return false;
    }
    
    // Create network
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
        builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
    if (!network) {
        std::cerr << "Failed to create network definition" << std::endl;
        return false;
    }
    
    // Create ONNX parser
    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));
    if (!parser) {
        std::cerr << "Failed to create ONNX parser" << std::endl;
        return false;
    }
    
    // Parse ONNX file
    std::cout << "Parsing ONNX model: " << onnxPath << std::endl;
    if (!parser->parseFromFile(onnxPath.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kWARNING))) {
        std::cerr << "Failed to parse ONNX file" << std::endl;
        return false;
    }
    
    // Create builder config
    auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config) {
        std::cerr << "Failed to create builder config" << std::endl;
        return false;
    }
    
    // Set memory pool limit
    config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1U << 30); // 1GB
    
    // Enable FP16 if requested
    if (enableFP16) {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
        std::cout << "FP16 precision enabled" << std::endl;
    } else {
        std::cout << "Using FP32 precision" << std::endl;
    }
    
    // Build engine
    std::cout << "Building TensorRT engine (this may take several minutes)..." << std::endl;
    auto start = std::chrono::high_resolution_clock::now();
    
    auto engine = std::unique_ptr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config));
    
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::seconds>(end - start);
    
    if (!engine) {
        std::cerr << "Failed to build TensorRT engine" << std::endl;
        return false;
    }
    
    std::cout << "Engine built successfully in " << duration.count() << " seconds" << std::endl;
    
    // Print engine info
    std::cout << "Engine info:" << std::endl;
    std::cout << "  Number of I/O tensors: " << engine->getNbIOTensors() << std::endl;
    for (int i = 0; i < engine->getNbIOTensors(); ++i) {
        const char* tensorName = engine->getIOTensorName(i);
        auto dims = engine->getTensorShape(tensorName);
        std::cout << "  Tensor " << i << ": " << tensorName << " [";
        for (int j = 0; j < dims.nbDims; ++j) {
            std::cout << dims.d[j];
            if (j < dims.nbDims - 1) std::cout << ", ";
        }
        std::cout << "]" << std::endl;
    }
    
    // Serialize and save engine
    std::cout << "Saving engine to: " << enginePath << std::endl;
    auto serializedEngine = std::unique_ptr<nvinfer1::IHostMemory>(engine->serialize());
    if (!serializedEngine) {
        std::cerr << "Failed to serialize engine" << std::endl;
        return false;
    }
    
    std::ofstream file(enginePath, std::ios::binary);
    if (!file.good()) {
        std::cerr << "Failed to create engine file: " << enginePath << std::endl;
        return false;
    }
    
    file.write(static_cast<const char*>(serializedEngine->data()), serializedEngine->size());
    file.close();
    
    std::cout << "Engine saved successfully (" << serializedEngine->size() / (1024*1024) << " MB)" << std::endl;
    return true;
}

int main(int argc, char* argv[]) {
    if (argc < 3) {
        std::cout << "Usage: " << argv[0] << " <onnx_path> <engine_path> [fp16]" << std::endl;
        std::cout << "Example: " << argv[0] << " models/yolov5s.onnx models/yolov5s.engine" << std::endl;
        std::cout << "Example: " << argv[0] << " models/yolov5s.onnx models/yolov5s.engine fp16" << std::endl;
        return 1;
    }
    
    std::string onnxPath = argv[1];
    std::string enginePath = argv[2];
    bool enableFP16 = (argc > 3 && std::string(argv[3]) == "fp16");
    
    std::cout << "=== TensorRT Engine Builder ===" << std::endl;
    std::cout << "ONNX model: " << onnxPath << std::endl;
    std::cout << "Engine output: " << enginePath << std::endl;
    std::cout << "Precision: " << (enableFP16 ? "FP16" : "FP32") << std::endl;
    std::cout << "================================" << std::endl;
    
    // Check if ONNX file exists
    std::ifstream onnxFile(onnxPath);
    if (!onnxFile.good()) {
        std::cerr << "ONNX file not found: " << onnxPath << std::endl;
        return 1;
    }
    onnxFile.close();
    
    if (buildEngine(onnxPath, enginePath, enableFP16)) {
        std::cout << "\n✓ Engine build completed successfully!" << std::endl;
        std::cout << "Engine location: " << enginePath << std::endl;
        return 0;
    } else {
        std::cerr << "\n✗ Engine build failed!" << std::endl;
        return 1;
    }
}