【光流估计】【深度学习】Windows11下FastFlowNet代码Pytorch官方实现与源码讲解

【光流估计】【深度学习】Windows11下FastFlowNet代码Pytorch官方实现与源码讲解

提示:最近开始在【光流估计】方面进行研究,记录相关知识点,分享学习中遇到的问题已经解决的方法。

---

---

前言

FastFlowNet是由上海交通大学的Kong, Lingtong等人在《FastFlowNet: A Lightweight Network for Fast Optical Flow Estimation【ICRA-2021】》【论文地址】一文中提出的针对移动设备优化的光流估计网络，通过头部增强池金字塔(HEPP)特征提取、中心密集膨胀相关(CDDC)层和高效的洗牌块解码器(SBD)实现高效计算。该网络在保持大搜索半径的同时降低了参数数量和计算成本，适合在资源有限的平台上运行。
在详细解析FastFlowNet网络之前，首要任务是搭建FastFlowNet【Pytorch-demo地址】所需的运行环境，并完成模型训练和测试工作，展开后续工作才有意义。

---

FastFlowNet模型运行环境搭建

注意需要安装Visual Studio，博主安装的是 Visual Studio 2019版本。
在win11环境下装anaconda环境，方便搭建专用于FastFlowNet模型的虚拟环境。

• 查看主机支持的cuda版本(最高)

  # 打开cmd,执行下面的指令查看CUDA版本号
  nvidia-smi
  

  

• 安装GPU版本的torch【官网】
  博主的cuda版本是12.2，博主选的11.8也没问题。
  
  其他cuda版本的torch在【以前版本】找对应的安装命令。

• 博主安装环境参考

  # 创建虚拟环境
  conda create -n FastFlowNet python=3.10
  # 查看新环境是否安装成功
  conda env list
  # 激活环境
  activate FastFlowNet 
  # 分别安装pytorch和torchvision
  pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  # 安装 opencv
  pip install -i https://pypi.tuna.tsinghua.edu.cn/simple opencv-python
  # 安装Correlation模块,进入工程的correlation_package路径下
  cd ./models/correlation_package/
  python setup.py build
  python setup.py install
  # 查看所有安装的包
  pip list
  conda list
  

  安装Correlation模块出现的问题解决方案：“ “getCurrentCUDAStream”: 不是 “at::Context” 的成员”
  
  解决方案： 找到models/correlation_package路径下的correlation_cuda.cc文件，修改里面内容。

  // 注释 #include <ATen/ATen.h> 
  // 更改为
  #include <ATen/cuda/CUDAContext.h> 
  // 注释 at::globalContext().getCurrentCUDAStream()
  // 更改为
  at::cuda::getCurrentCUDAStream()
  

  
  

---

FastFlowNet模型运行

模型权重

原作者预训练权重放置在当前工程目录下的checkpoints目录下

FastFlowNet测试

源代码未开源训练部分代码

对运行速度进行基准测试并计算模型参数

python benchmark.py

可能出现的问题：“RuntimeError: Legacy autograd function with non-static forward method is deprecated. Please use new-style autograd function with static forward method”。

解决方式： 正在使用旧版本的 PyTorch autograd 系统，需要修改models/correlation_package/correlation.py内容的部分代码内容，使用新的方式定义自定义的 autograd 函数。
修改部分1：修改前的原始代码

class CorrelationFunction(Function):

    def __init__(self, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        super(CorrelationFunction, self).__init__()
        self.pad_size = pad_size
        self.kernel_size = kernel_size
        self.max_displacement = max_displacement
        self.stride1 = stride1
        self.stride2 = stride2
        self.corr_multiply = corr_multiply
        # self.out_channel = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1)

    def forward(self, input1, input2):
        self.save_for_backward(input1, input2)

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()
            output = input1.new()

            correlation_cuda.forward(input1, input2, rbot1, rbot2, output, 
                self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)

        return output

    def backward(self, grad_output):
        input1, input2 = self.saved_tensors

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()

            grad_input1 = input1.new()
            grad_input2 = input2.new()

            correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,
                self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)

        return grad_input1, grad_input2

修改部分1：修改后的新代码

class CorrelationFunction(Function):
    @staticmethod
    def forward(ctx, input1, input2, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, stride2=2, corr_multiply=1):
        ctx.pad_size = pad_size
        ctx.kernel_size = kernel_size
        ctx.max_displacement = max_displacement
        ctx.stride1 = stride1
        ctx.stride2 = stride2
        ctx.corr_multiply = corr_multiply
        ctx.save_for_backward(input1, input2)

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()
            output = input1.new()

            correlation_cuda.forward(input1, input2, rbot1, rbot2, output,
                ctx.pad_size, ctx.kernel_size, ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply)

        return output

    @staticmethod
    def backward(ctx, grad_output):
        input1, input2 = ctx.saved_tensors

        with torch.cuda.device_of(input1):
            rbot1 = input1.new()
            rbot2 = input2.new()

            grad_input1 = input1.new()
            grad_input2 = input2.new()

            correlation_cuda.backward(input1, input2, rbot1, rbot2, grad_output, grad_input1, grad_input2,
                ctx.pad_size, ctx.kernel_size, ctx.max_displacement, ctx.stride1, ctx.stride2, ctx.corr_multiply)

        return grad_input1, grad_input2, None, None, None, None, None

修改部分2：修改前的原始代码

    def forward(self, input1, input2):

        result = CorrelationFunction(self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)(input1, input2)

        return result

修改部分2：修改后的新代码

    def forward(self, input1, input2):

        result = CorrelationFunction.apply(input1, input2, self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)

        return result

预测光流的演示

python demo.py

在代码里可以修改自己想要测试的图片，这里博主使用源代码测试数据，效果如下：

---

总结

尽可能简单、详细的介绍了FastFlowNet的安装流程以及FastFlowNet的使用方法。后续会根据自己学到的知识结合个人理解讲解FastFlowNet的原理和代码。