Search code examples
c++deep-learningpytorchonnxonnxruntime

how can load ONNX model in C++


I have convert PyTorch SuperResouloution (.pth) Model to ONNX model (.onnx) with code below:

import torch
import onnx
from onnx_tf.backend import prepare
from basicsr.archs.rrdbnet_arch import RRDBNet

pth_path = "RESRGAN_x4plus_Main.pth"

Model = RRDBNet(num_in_ch=3, num_out_ch=3, scale=4, num_feat=64, num_block=23, num_grow_ch=32)

Model.load_state_dict(torch.load(pth_path)\["params_ema"\])
Model.eval()

X = torch.randn((1, 3, 64, 64))

torch.onnx.export(Model, X,
"Model.onnx",
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=\["input"\],
output_names=\["output"\])

and now I want to load this model into C++ and run it... I have seen many examples of onnxruntime for loading models but I am very confused. sample codes are very complex and not for Super Resolution models.

I emphasize I want to load the model directly in the C++ program (with onnx-runtime package)

I don't want to convert onnx model into another model format...

I have to load the .onnx model into c++ and pass the image into it and I expect to receive an Image from the model output. (model is a GAN)


Solution

  • You can upload ONNX model to C++ using OpenCV library.

    cv::dnn::Net model = cv::dnn::readNetFromONNX("Model.onnx");
    
    // prepare model input
    cv::Mat blob = cv::dnn::blobFromImage(originalImage,
        1.f/255.f,
        {64,64});
    
    // set model input
    model.setInput(blob);
    
    // get model results
    cv::Mat output = model.forward();
    
    // get pointer to raw data
    float* ptr = output.ptr<float>();
    

    You can refer to this code to find out how to load GAN model to C++ and create an image from raw data pointer.