BasicVSR++模型转JIT并用c++libtorch推理

安装BasicVSR++ 环境

1.下载源码

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

2. 新建一个conda环境

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

3. 安装pytorch

pytorch官网 安装合适的版本
我这里是CUDA11.6版本

conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia

4. 安装 mim 和 mmcv-full

pip install openmim
mim install mmcv-full

5. 安装 mmedit

pip install mmedit

6. 下载模型文件

下载模型文件放在这里chkpts/basicvsr_plusplus_reds4.pth

7. 测试一下能否正常运行

python demo/restoration_video_demo.py configs/basicvsr_plusplus_reds4.py chkpts/basicvsr_plusplus_reds4.pth data/demo_000 results/demo_000

OK ! 环境正常！下面开始转换工作

转换为JIT模型

在demo下新建一个转换脚本

import os
import cv2
import mmcv
import numpy as np
import torch
from mmedit.core import tensor2img
from mmedit.apis import init_model


def main():
    # 加载模型并设置为评估模式
    model = init_model("configs/basicvsr_plusplus_reds4.py",
                       "chkpts/basicvsr_plusplus_reds4.pth", device=torch.device('cuda', 0))
    model.eval()

    # 准备一个示例输入
    src1 = cv2.imread("./data/img/00000000.png")
    src = cv2.cvtColor(src1, cv2.COLOR_BGR2RGB)
    src = torch.from_numpy(src / 255.).permute(2, 0, 1).float()
    src = src.unsqueeze(0)

    input_arg = torch.stack([src], dim=1)
    input_arg = input_arg.to(torch.device('cuda', 0))  # 确保输入在GPU上

    # # 执行模型推理
    # with torch.no_grad():  # 在推理时不需要计算梯度
    #     result = model(input_arg, test_mode=True)['output'].cpu()
    # output_i = tensor2img(result)
    # mmcv.imwrite(output_i, "./test.png")

    # 模型转换
    traced_model = torch.jit.trace(model.generator, input_arg)
    torch.jit.save(traced_model, "basicvsrPP.pt")

    # 测试
    res = traced_model(input_arg)
    out = tensor2img(res)
    mmcv.imwrite(out, "./testoo.png")


if __name__ == '__main__':
    main()

用c++ libtorch推理

/*
 * @Author: Liangbaikai
 * @LastEditTime: 2024-03-29 11:28:42
 * @Description: 视频超分
 * Copyright (c) 2024 by Liangbaikai, All Rights Reserved.
 */

#pragma once
#include <iostream>
#include <torch/script.h>
#include <torch/torch.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <c10/cuda/CUDACachingAllocator.h>

namespace LIANGBAIKAI_BASE_MODEL_NAME
{
    class lbk_video_super_resolution_basicPP
    {
    public:
        lbk_video_super_resolution_basicPP() = default;

        virtual ~lbk_video_super_resolution_basicPP()
        {
            c10::cuda::CUDACachingAllocator::emptyCache();
            // cudaDeviceReset();
        }

        /**
         * @description: 初始化
         * @param {string} &modelpath 模型文件
         * @param {int} gpuid GPU的id
         * @return {*}成功返回0，失败返回-1
         */
        int init(const std::string &modelpath, int gpuid = 0)
        {
            try
            {
                _mymodule = std::make_unique<torch::jit::script::Module>(torch::jit::load(modelpath));
            }
            catch (const c10::Error &e)
            {
                std::cerr << "Error loading the model " << modelpath << std::endl;
                std::cerr << "Error " << e.what() << std::endl;
                return -1;
            }
            _gpuid = gpuid;

            if ((_gpuid < 0) || (!torch::cuda::is_available()))
            {
                _device = std::make_unique<torch::Device>(torch::kCPU);
                _mymodule->to(at::kCPU);
            }
            else
            {
                _device = std::make_unique<torch::Device>(torch::kCUDA, _gpuid);
                _mymodule->to(at::kCUDA, _gpuid);
            }

            _mymodule->eval();
            _modelsuccess = true;
            return 0;
        }

        /**
         * @description: 推理
         * @param {Mat} &inputpic 输入图片
         * @param {Mat} &outputpic 输出结果
         * @param {bool} showlog  是否打印日志
         * @return {*} 成功返回0，失败返回-1
         */
        int inference(cv::Mat &inputpic, cv::Mat &outputpic, bool showlog = false)
        {
            if (inputpic.empty() || (inputpic.channels() != 3))
            {
                std::cout << "input data ERROR" << std::endl;
                return -1;
            }

            if (!_modelsuccess)
            {
                std::cout << "model has not been inited!" << std::endl;
                return -1;
            }

            // torch::DeviceGuard 是一个类，它的作用是确保在使用完设备（如CPU或GPU）后，能够正确地将设备恢复到使用前的状态。
            torch::DeviceGuard device_guard(*_device); // 作用域内所有操作都在指定设备上运行，离开此作用域恢复

            cv::transpose(inputpic, inputpic); // 顺时针旋转

            // 将图片转换为tensor
            cv::Mat img_float;
            inputpic.convertTo(img_float, CV_32FC3, 1.0 / 255);
            torch::Tensor img_tensor = torch::from_blob(img_float.data, {img_float.rows, img_float.cols, 3}, torch::kFloat32).permute({2, 1, 0});
            img_tensor = (img_tensor - 0.5) / 0.5;
            img_tensor = (img_tensor + 1) / 2;
            img_tensor = torch::clamp(img_tensor, 0, 1);

            torch::Tensor src_unsqueezed = img_tensor.unsqueeze(0).to(*_device); // 将tensor转移到GPU上

            std::vector<torch::Tensor> tensors_to_stack = {src_unsqueezed}; // 创建一个包含 src 的 vector

            torch::Tensor input_arg = torch::stack(tensors_to_stack, 1); // 沿着维度1堆叠tensors

            if (showlog)
            {
                std::cout << input_arg.sizes() << std::endl;
            }

            torch::NoGradGuard no_grad; // 暂时禁用梯度计算
            auto output_dict = _mymodule->forward({input_arg});

            torch::Tensor output_data;
            if (output_dict.isTensor())
            {
                output_data = output_dict.toTensor().to(at::kCPU); // 如果是Tensor，则通过toTensor()方法获取它
                if (showlog)
                {
                    std::cout << "out shape: " << output_data.sizes() << std::endl;
                }
            }
            else
            {
                if (showlog)
                {
                    std::cerr << "The IValue does not contain a Tensor." << std::endl;
                }
            }

            float *f = output_data.data_ptr<float>();

            int output_width = output_data.size(3);
            int output_height = output_data.size(4);
            int size_pic = output_width * output_height;

            std::vector<cv::Mat> rgbChannels(3);
            rgbChannels[0] = cv::Mat(output_width, output_height, CV_32FC1, f);
            rgbChannels[1] = cv::Mat(output_width, output_height, CV_32FC1, f + size_pic);
            rgbChannels[2] = cv::Mat(output_width, output_height, CV_32FC1, f + size_pic + size_pic);

            rgbChannels[0].convertTo(rgbChannels[0], CV_8UC1, 255);
            rgbChannels[1].convertTo(rgbChannels[1], CV_8UC1, 255);
            rgbChannels[2].convertTo(rgbChannels[2], CV_8UC1, 255);

            merge(rgbChannels, outputpic);

            return 0;
        }

    private:
        bool _modelsuccess = false;
        int _gpuid = 0;
        std::unique_ptr<torch::Device> _device;
        std::unique_ptr<torch::jit::script::Module> _mymodule;
    };

}

#include <unistd.h>
#include"lbk_video_super_resolution.hpp"
using namespace LIANGBAIKAI_BASE_MODEL_NAME;
int main(int argc,char *argv[])
{
    if(argc < 5){
        std::cout << "./test 模型  GPUid(cpu传-1) 输入图片 输出图片" << std::endl;
        return -1;
    }
    std::string modelfile = argv[1];
    int gpuid = atoi(argv[2]);
    std::string imgfile = argv[3];
    std::string outfile = argv[4];

    cv::Mat src = cv::imread(imgfile);


    lbk_video_super_resolution_basicPP test;
    if(0 > test.init(modelfile,gpuid)){
        std::cout << "init failed" << std::endl;
        return -1;
    }
    cv::Mat out;
    int rec = test.inference(src,out,true);
    if(rec >= 0){
        cv::imwrite(outfile, out);
    }

    return 0;
}

效果