FlowNet2 torch 代码运行经验

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)
  • 4
    点赞
  • 62
    收藏
    觉得还不错? 一键收藏
  • 29
    评论
以下是Flownet中的correlation代码示例: ```python from torch.nn.functional import unfold def correlation(input1, input2, kernel_size=1, max_displacement=1, stride1=1, stride2=1, padding=0, dilation=1): b, c, h, w = input1.size() # Pad input2 padded_input2 = F.pad(input2, (padding, padding, padding, padding)) # Unfold input2 unfolded_input2 = unfold(padded_input2, kernel_size=kernel_size, dilation=dilation, stride=stride2) # Reshape unfolded input2 unfolded_input2 = unfolded_input2.view(b, c, kernel_size**2, h, w) # Create empty output tensor output = torch.zeros(b, (2 * max_displacement + 1) ** 2, h // stride1, w // stride1).to(input1.device) # Loop over displacements for i in range(-max_displacement, max_displacement + 1): for j in range(-max_displacement, max_displacement + 1): # Shift input2 shifted_input2 = padded_input2[:, :, padding + i:h + padding + i, padding + j:w + padding + j] # Unfold shifted input2 unfolded_shifted_input2 = unfold(shifted_input2, kernel_size=kernel_size, dilation=dilation, stride=stride2) # Reshape unfolded shifted input2 unfolded_shifted_input2 = unfolded_shifted_input2.view(b, c, kernel_size**2, h, w) # Compute correlation corr = (input1 * unfolded_shifted_input2).sum(dim=1) # Add to output tensor output[:, (i + max_displacement) * (2 * max_displacement + 1) + j + max_displacement, :, :] = corr return output ``` 这段代码实现了correlation操作,它是Flownet中的关键操作之一。在此操作中,两个输入张量被移位并相乘,然后求和以计算它们之间的相似性。这个操作在光流估计中非常有用,因为它可以帮助找到两个图像中的相同区域。
评论 29
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值