RIFE插帧算法(V4.6)-pytorch转ncnn模型

        首先附上各位大佬的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 层，重新编译对应库即可。