FlowNet2 torch 代码运行经验

最新推荐文章于 2024-01-25 12:51:01 发布

原创

最新推荐文章于 2024-01-25 12:51:01 发布 · 6.3k 阅读

66 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习

本文档记录了在NVIDIA提供的FlowNet2 PyTorch实现中运行和优化代码的经验。主要内容包括环境配置、权重文件下载、Stintel数据集获取、多GPU运行注意事项、数据集特征提取以及.flow文件可视化方法。适合对深度学习和光流估计感兴趣的开发者参考。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

FlowNet2是目前最流行的网络，原文中使用的是CAFFE进行训练的网络。在GITHUB上最火的是NIVDIA官方给出的torch代码。运行的时候需要一些操作技巧，对于linux小白的我，可能不太友好，在这记录一下。

文章目录

NIVDIA官方代码

https://github.com/NVIDIA/flownet2-pytorch

运行环境要求

由于是NIVDIA官方给出的网络，程序代码都封装的比较好了，主要通过main.py的参数，既可以实现网络的各类调整。
首先操作系统的运行环境最好是Linux，因为其中不少代码写成的.sh。windows需要git

python3
numpy
PyTorch ( == 0.4.1, for <= 0.4.0 see branch python36-PyTorch0.4)
scipy
scikit-image
tensorboardX
colorama, tqdm, setproctitle

代码下载

# get flownet2-pytorch source
git clone https://github.com/NVIDIA/flownet2-pytorch.git
cd flownet2-pytorch

# install custom layers
bash install.sh

下载权重文件

网址上给出了好几个pre-trained网络的weight，大家无法可科学上网下载权重的，从如下百度网盘中下载

链接: https://pan.baidu.com/s/1aiKc80Q-DI4p3MVODcPsyQ 提取码: apz9

主要用到的就是FlowNet2这个文件
在这里插入图片描述

Stintel数据集下载

http://sintel.is.tue.mpg.de/

运行main.py

测试数据(inference)

python main.py --inference --model FlowNet2 --save_flow --inference_dataset ImagesFromFolder --resume ./checkpoints/FlowNet2_checkpoint.pth.tar --save ./output --inference_dataset_root ./MPI-Sintel-complete/test/clean/wall

除了以上的初始化参数外，还有几个参数这里也介绍一下：

number_gpus

之前用单GPU的电脑时没有发现问题，最近用4*GPU时发现输出的.flo会除以4，去issue上看也有人提出问题，但是没有人解答。所以在使用多GPU时记得将其改为1

inference_dataset_iext

这里需要根据输入图像的后缀格式定义，比如png，jpg等

inference_visualize

同步可以输出png格式的可视化图

修改输出.flo图像大小

由于.flo格式的数据相对很大，比如一个[1280, 720]尺寸的.flo图像居然有7+M。而输入到网络中的图像尺寸为[224, 398]，因此对utils.frame_util中利用resize模块进行相应修改如下：

import numpy as np
from os.path import *
#from scipy.misc import imread
from imageio import imread
from . import flow_utils 
import cv2

def read_gen(file_name):
    ext