Method _startONNXSession

sloam/src/segmentation/inference.cpp:21–78

}

void Segmentation::_startONNXSession(const std::string sessionName, const std::string modelFilePath, bool useCUDA, size_t numThreads){
    Ort::SessionOptions sessionOptions;
    sessionOptions.SetIntraOpNumThreads(numThreads);
    // Sets graph optimization level
    // Available levels are
    // ORT_DISABLE_ALL -> To disable all optimizations
    // ORT_ENABLE_BASIC -> To enable basic optimizations (Such as redundant node
    // removals) ORT_ENABLE_EXTENDED -> To enable extended optimizations
    // (Includes level 1 + more complex optimizations like node fusions)
    // ORT_ENABLE_ALL -> To Enable All possible optimizations
    sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
    if (useCUDA){
        // Using CUDA backend
        // https://github.com/microsoft/onnxruntime/blob/v1.8.2/include/onnxruntime/core/session/onnxruntime_cxx_api.h#L329
        OrtCUDAProviderOptions cuda_options{0};
        sessionOptions.AppendExecutionProvider_CUDA(cuda_options);
    }

    auto env = boost::make_shared<Ort::Env>(ORT_LOGGING_LEVEL_ERROR, sessionName.c_str());
    _env = boost::move(env);

    auto session = boost::make_shared<Ort::Session>(*_env, modelFilePath.c_str(), sessionOptions);
    _session = boost::move(session);

    auto memInfo = boost::make_shared<Ort::MemoryInfo>(Ort::MemoryInfo::CreateCpu(
        OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault));
    _memoryInfo = boost::move(memInfo);

    // Ort::AllocatorWithDefaultOptions allocator;
    // INPUT 
    const char* inputName = _session->GetInputName(0, _allocator);
    std::cout << "[Segmentation] Input Name: " << inputName << std::endl;

    Ort::TypeInfo inputTypeInfo = _session->GetInputTypeInfo(0);
    auto inputTensorInfo = inputTypeInfo.GetTensorTypeAndShapeInfo();
    _inputDims = inputTensorInfo.GetShape();
    // _inputDims = {1, 64, 2048, 2};
    std::cout << "[Segmentation] Input Dimensions: "; printVector(_inputDims);

    // OUTPUT 
    const char* outputName = _session->GetOutputName(0, _allocator);
    std::cout << "[Segmentation] Output Name: " << outputName << std::endl;

    Ort::TypeInfo outputTypeInfo = _session->GetOutputTypeInfo(0);
    auto outputTensorInfo = outputTypeInfo.GetTensorTypeAndShapeInfo();
    // ONNXTensorElementDataType outputType = outputTensorInfo.GetElementType();

    _outputDims = outputTensorInfo.GetShape();
    // _outputDims = {1, 64, 2048, 2};
    std::cout << "[Segmentation] Output Dimensions: "; printVector(_outputDims);

    // _inputTensorSize = _img_w * _img_h * _img_d;
    // _outputTensorSize = _img_w * _img_h * _img_d;
    _inputTensorSize = vectorProduct(_inputDims);
    _outputTensorSize = vectorProduct(_outputDims);
    _inputNames = {inputName};
    _outputNames = {outputName};
}