首先附上各位大佬的github 地址链接,请多多支持各位开源大佬的工作(动动小手,star一下):
1. rife插帧算法原地址(pytorch实现),作者特别优秀!兼顾了时间和性能的一个插帧算法,尤其是对于分辨率比较低的图像进行插帧,效果一级棒:
https://github.com/megvii-research/ECCV2022-RIFE
2. rife插帧算法的ncnn框架实现(C++代码),基于Vulkan和良好的ncnn开源生态,A卡N卡都可运行,尤其对移动端效果更好:
https://github.com/nihui/rife-ncnn-vulkan
3. 一个基于TensorRT和ncnn等框架加速的超分和插帧模型存储库,包含性能比较好的一些推理代码,基于docker移植方便,作者人也非常nice,非常感谢他提供了rife模型“pytorch->onnx->ncnn”模型的转换流程:
https://github.com/styler00dollar/VSGAN-tensorrt-docker
4. ncnn开源框架的维护者nihui大佬提出了通过pnnx进行pytorch模型转ncnn模型的方式,但是我这边直接对rife 算法进行操作存在问题,大家如果想尝试可以根据nihui 大佬写的文章操作:
https://zhuanlan.zhihu.com/p/427620428
书归正传,本文介绍一下如何通过pytorch转onnx再转ncnn模型的方式,将rife插帧算法应用在ncnn框架上。vs-rife的开发大佬基于原作者的pytorch代码,对V4版本的算法出入参做了一些规范,从而能够适配vapoursynth音视频处理框架:
https://github.com/HolyWu/vs-rife.git
本文就基于上述模型的出入参和基本结构进行V4.6模型后续的转换,环境为CUDA11.3,pytorch版本安装如下:
pip install torch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 --extra-index-url https://download.pytorch.org/whl/cu113
一、模型框架修改
RIFE插帧算法从V4开始引入了timestep的概念,可以直接合成两帧中间任意时刻的视频帧,我测试过0.2-0.25-0.3,生成的每一帧都有细微的差别,确实非常方便。不过在转换模型的时候,在IFNet_HDv3_v4_6.py 中,对于传入的timestep 一定要做如下的修改,添加3行处理:
def forward(self, img0, img1, timestep):
flow_list = []
merged = []
mask_list = []
warped_img0 = img0
warped_img1 = img1
x = torch.cat((img0, img1), 1) # Add
timestep = (x[:, :1].clone() * 0 + 1) * timestep # Add
timestep = timestep.float() # Add
在RIFE模型的最后要进行两个warp 的操作,warp 算子是原作者自己写的,ncnn 框架中肯定也没有,为了避免转换出错,这里就需要把warp 算子替换成RIFE 算法完全没有使用,并且ncnn 中包含的算子,在IFNet_HDv3_v4_6.py 中将warp 算子替换为pow 算子:
mask_list.append(mask)
flow_list.append(flow)
# warped_img0 = warp(img0, flow[:, :2]) # Delete
# warped_img1 = warp(img1, flow[:, 2:4]) #Delete
warped_img0 = img0**flow[:,:3] # Add
warped_img1 = img1**flow[:,1:4] # Add
merged.append((warped_img0, warped_img1))
其它算子ncnn都是支持的,不需要对RIFE框架做太大的改动。
二、模型转onnx
vs-rife 的github 路径下pkl 文件都是空的,需要去原作者的模型仓库中下载V4.6 的模型然后替换到对应路径下。利用torch.onnx.export 进行对应转换即可。我这里自己写了个简单的脚本进行转换,放在IFNet_HDv3_v4_6.py 同级目录下即可,可以参考:
import os
import torch
# modified version by 4.6
model = "4.6"
if model == "4.6":
from IFNet_HDv3_v4_6 import IFNet
if model == "4.4":
from IFNet_HDv3_v4_4 import IFNet
print("Converting")
device = torch.device('cuda', 0)
package_dir = os.path.dirname(os.path.realpath(__file__))
model_name = f'flownet_v{model}.pkl'
# Control the process resolution for optical flow model
# Try scale=0.5 for 4K video
# scale must be 0.25, 0.5, 1.0, 2.0, or 4.0
scale = 1.0
#Smooth predictions in areas where the estimation is uncertain
ensemble = False
checkpoint = torch.load(os.path.join(package_dir, model_name), map_location='cpu')
checkpoint = {k.replace('module.', ''): v for k, v in checkpoint.items() if 'module.' in k}
flownet = IFNet(scale, ensemble)
flownet.load_state_dict(checkpoint, strict=False)
flownet.eval().to(device, memory_format=torch.channels_last)
# export rife pytorch to onnx
test_input = torch.rand(1, 3, 256, 256).cuda()
timestep = torch.Tensor([0.5]).cuda()
dynamic_axes = {
"in0": {2: "width", 3: "height"},
"in1": {2: "width", 3: "height"},
"out0": {2: "width", 3: "height"},
}
torch.onnx.export(
flownet,
(test_input, test_input, timestep),
"rife.onnx",
verbose=False,
opset_version=11,
input_names=["in0", "in1", "in2"],
output_names=["out0"],
# dynamic_axes=dynamic_axes,
)
print("Finished")
执行脚本就可以获取到转换后的rife.onnx 文件。
在转换过程中可能出现下面几个问题,我把解决方式也附一下:
1. 报错:attempted relative import with no known parent package
原因:python脚本在同级目录下,却使用了相对导入。
解决:把所有报这个错的地方相对路径的点去掉即可。
# from .warplayer import warp # Delete
from warplayer import warp # Add
2. 报错:module 'torch' has no attribute 'fx'
原因:我使用的这个版本的pytorch 没有fx 模块。
解决:因为warp 算子我们已经用其它算子代替,所以直接将报错的语句注释掉即可。
# torch.fx.wrap('warp') # Delete
3. 警告:Default upsampling behavior when mode=bilinear is changed to align_corners=False since 0.4.0
原因:从pytorch 0.4.0 版本之后,nn.Upsample 中bilinear 模式align_corners 参数默认设置为False 了,如果之前需要align_corners=True ,则需要手动设置一下参数。
解决:对RIFE模型没有影响,所以我没有做处理。
脚本执行完成后,在目录下可以看到rife.onnx 模型文件:
三、onnx 转ncnn 模型
通过onnx 转换到其它模型的时候,推荐使用一键转换工具,方便快捷:
一键转换 Caffe, ONNX, TensorFlow 到 NCNN, MNN, Tengine
首先将onnx 模型进行simplifier 和optimizer 得到优化后的onnx 模型,然后将优化后的模型转换到ncnn 模型,即可得到.bin 和.param 两个模型文件,点击保存。
然后,我们之前对RIFE 框架做了一个warp 算子替换为pow 算子的操作,在rife-ncnn-vulkan 代码路径下,nihui 大佬手动添加了rife.Warp 算子,所以第二步我们需要将.param 文件中所有的pow 算子都替换成rife.Warp 算子,同时将crop 的维度恢复成二维:
以一个pow 算子的替换为例,实际所有的pow 算子都需要进行替换
原模型参数:
Crop Slice_81 1 1 200_splitncnn_3 210 -23309=1,0 -23310=1,3 -23311=1,0
BinaryOp Pow_82 2 1 in0_splitncnn_3 210 211 0=6
Crop Slice_87 1 1 200_splitncnn_2 216 -23309=1,1 -23310=1,4 -23311=1,0
BinaryOp Pow_88 2 1 in1_splitncnn_3 216 217 0=6
替换后的模型参数:
Crop Slice_81 1 1 200_splitncnn_3 210 -23309=1,0 -23310=1,2 -23311=1,0
rife.Warp warp_82 2 1 in0_splitncnn_3 210 211 0=6
Crop Slice_87 1 1 200_splitncnn_2 216 -23309=1,2 -23310=1,4 -23311=1,0
rife.Warp warp_88 2 1 in1_splitncnn_3 216 217 0=6
将替换完成后的.bin 和.param 模型文件放到rife-ncnn-vulkan 对应目录下,执行进程测试即可。我在V4.6转换的模型执行过程中有一个报错:layer MemoryData not exists or registered
原因:V4.6 之前的版本是不需要编译MemoryData 层的,但是V4.6 需要这个层。
解决:将rife-ncnn-vulkan 代码路径src 目录下的CMakeLists.txt 文件中的
option(WITH_LAYER_memorydata "" OFF)
替换为:
option(WITH_LAYER_memorydata "" ON)
替换完成后重新进行编译即可,如果是使用的ncnn-lib 库,则同样开启memorydata 层,重新编译对应库即可。