RealBasicVSR模型转成ONNX以及JIT并用c++推理

一. RealBasicVSR模型转成ONNX

安装RealBasicVSR的环境

1. 新建一个conda环境

conda create -n RealBasicVSR_to_ONNX  python=3.8 -y
conda activate RealBasicVSR_to_ONNX

2. 安装pytorch(官网上选择合适的版本)版本太低会有问题

conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge

3. 安装 mim 和 mmcv-full

pip install openmim
mim install mmcv-full

4. 安装 mmedit

pip install mmedit

下载RealBasicVSR源码

git clone https://github.com/ckkelvinchan/RealBasicVSR.git

下载模型文件

模型文件下载 (Dropbox / Google Drive / OneDrive) ，随便选一个渠道下载就行

cd RealBasicVSR
#然后新建文件夹model
将模型文件放在model文件夹下

写一个模型转换的脚步

import cv2
import mmcv
import numpy as np
import torch
from mmcv.runner import load_checkpoint
from mmedit.core import tensor2img

from realbasicvsr.models.builder import build_model

def init_model(config, checkpoint=None):
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {
     type(config)}')
    config.model.pretrained = None
    config.test_cfg.metrics = None
    model = build_model(config.model, test_cfg=config.test_cfg)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)

    model.cfg = config  # save the config in the model for convenience
    model.eval()
    return model


def main():
    model = init_model("./configs/realbasicvsr_x4.py","./model/RealBasicVSR_x4.pth")
    src = cv2.imread("./data/img/test1.png")
    src = torch.from_numpy(src / 255.).permute(2, 0, 1).float()
    src = src.unsqueeze(0)
    input_arg = torch.stack([src], dim=1)

    torch.onnx.export(model,
        input_arg,
        'realbasicvsr.onnx',
        training= True,
        input_names= ['input'],
        output_names=['output'],
        opset_version=11,
        dynamic_axes={
   'input' : {
   0 : 'batch_size', 3 : 'w', 4 : 'h'}, 'output' : {
   0 : 'batch_size', 3 : 'dstw', 4 : 'dsth'}})


if __name__ == '__main__':
    main()

这里报错：

ValueError: SRGAN model does not support `forward_train` function.

直接将这个test_mode默认值改为Ture,让程序能走下去就行了。

测试生成的模型

这里已经得到了 realbasicvsr.onnx 模型文件了.

import onnxruntime as ort
import numpy as np
import onnx
import cv2


def main():
    onnx_model = onnx.load_model("./realbasicvsr.onnx")
    onnxstrongmodel = onnx_model.SerializeToString()
    sess = ort.InferenceSession(onnxstrongmodel)
    
    providers = ['CPUExecutionProvider']
    options = [{
   }]
    is_cuda_available = ort.get_device() == 'GPU'
    if is_cuda_available:
        providers.insert(0, 'CUDAExecutionProvider')
        options.insert(0, {
   'device_id': 0})
    sess.set_providers(providers, options)

    input_name = sess.get_inputs()[0].name
    output_name = sess.get_outputs()[1].name
    print(sess.get_inputs()[0])
    print(sess.get_outputs()[0])
    print(sess.get_outputs()[0].shape)
    print(sess.get_inputs()[0].shape)

    img = cv2.imread("./data/img/test1.png")
    img = np.expand_dims((img/255.0).astype(np.float32).transpose(2,0,1), axis=0)
    imgs = np.array([img])
    print(imgs.shape)
    print(imgs)
    output = sess.run([output_name], {
   input_name : imgs})

    print(output)

    print(output[0].shape)
    output = np.clip(output, 0, 1)

    res = output[0][0][0].transpose(1, 2, 0)
    cv2.imwrite("./testout.png", (res * 255).astype(np.uint8))

if __name__ == '__main__':
    main()

至此模型转换部分就成功完成了

用ONNX Runtime c++推理

根据cuda版本选择合适的onnxruntime版本

下载onnx runtime的运行环境 onnxruntime
我这里下载这个：

#include <iostream>
#include <assert.h>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>

class ONNX_RealBasicVSR{
   
  public:
    ONNX_RealBasicVSR():session(nullptr){
   };
    virtual ~ONNX_RealBasicVSR() = default;
    /*初始化
    * @param model_path 模型
    * @param gpu_id 选择用那块GPU
    */
    void Init(const char * model_path,int gpu_id = 0);

    /**执行模型推理
     * @param src : 输入图
     * @param inputid : 输入id
     * @param outputid : 输出的id
     * @return 输出结果图
    */
    cv::Mat Run(cv::Mat src,unsigned inputid = 0,unsigned outputid = 0,bool show_log = false);

  private:
    /*获取模型的inputname 或者 outputname
    * @param input_or_output  选择要获取的是input还是output
    * @param id 选择要返回的是第几个name
    * @param show_log 是否打印信息
    * @return 返回name
    */
    std::string GetInputOrOutputName(std::string input_or_output = "input",unsigned id = 0,bool show_