济南网站建设团队,山西城乡和住房建设厅网站,视频生成链接网站,杭州 高端网站建设文章目录安装RealBasicVSR的环境1. 新建一个conda环境2. 安装pytorch(官网上选择合适的版本)版本太低会有问题3. 安装 mim 和 mmcv-full4. 安装 mmedit下载RealBasicVSR源码下载模型文件写一个模型转换的脚步测试生成的模型安装RealBasicVSR的环境
1. 新建一个conda环境
cond…
文章目录安装RealBasicVSR的环境1. 新建一个conda环境2. 安装pytorch(官网上选择合适的版本)版本太低会有问题3. 安装 mim 和 mmcv-full4. 安装 mmedit下载RealBasicVSR源码下载模型文件写一个模型转换的脚步测试生成的模型安装RealBasicVSR的环境
1. 新建一个conda环境
conda create -n RealBasicVSR_to_ONNX python3.8 -y
conda activate RealBasicVSR_to_ONNX2. 安装pytorch(官网上选择合适的版本)版本太低会有问题
conda install pytorch1.12.1 torchvision0.13.1 torchaudio0.12.1 cudatoolkit11.6 -c pytorch -c conda-forge3. 安装 mim 和 mmcv-full
pip install openmim
mim install mmcv-full4. 安装 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 tensor2imgfrom realbasicvsr.models.builder import build_modeldef init_model(config, checkpointNone):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, fbut got {type(config)})config.model.pretrained Noneconfig.test_cfg.metrics Nonemodel build_model(config.model, test_cfgconfig.test_cfg)if checkpoint is not None:checkpoint load_checkpoint(model, checkpoint)model.cfg config # save the config in the model for conveniencemodel.eval()return modeldef 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], dim1)torch.onnx.export(model,input_arg,realbasicvsr.onnx,training True,input_names [input],output_names[output],opset_version11,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 cv2def 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() GPUif is_cuda_available:providers.insert(0, CUDAExecutionProvider)options.insert(0, {device_id: 0})sess.set_providers(providers, options)input_name sess.get_inputs()[0].nameoutput_name sess.get_outputs()[1].nameprint(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), axis0)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()
至此模型转换部分就成功完成了
