一.方法：预测当前帧各像素’运动’（‘位移量’ ΔF）
简化架构：
补充：
方法：
3.1特征提取
key：相邻帧；映射；低分辨率（单分辨率）；残差块
消融实验1-组件性能
单分辨率提取特征
key：简化网络；大偏移匹配（‘感受野一样’）
上下文网络
key：更新融合空间信息；
3.2计算视觉相关性（key）
key：序列模型输入之一；通过‘局部位移信息’，推理出该点位移信息
key:特征向量对的点积
[c,hw]*[hw,c]=[hw,hw]
key：后两维均池化（保留高分辨率像素信息；获得’不同对象‘大小位移信息）
key：映射1（点->点**【光流估计序列(f1,f2)】）；映射2（点->区域【领域】**）；四种不同的局部大小位移信息
消融实验2-组件性能
全像素对点积：满足对指定范围的需要；便于矩阵运算
3.3 迭代更新
实际架构：
目标：预测当前帧各像素‘运动’（‘位移量’ΔF）
key-key-key：光流估计序列（获得‘局部位移信息’）；更新预测‘位移’量；更新光流估计序列
实现：
输入：上下文特征；相关性特征(（‘局部位移信息’）【领域】；光流特征(点->预测光流点)【初始化/迭代中的预测‘位移’量Δf】（最后，通过输出ΔF与label建立监督）
输出：ΔF=（Δf1，Δf2，…,Δfn） n:像素点数
key：各像素点的更新位移量Δf
key：上采样；凸组合；‘反卷积’
上采样：获得原输入分辨度
3.4 监督问题
key：l1 distance；loss；BP；验证训练有效性
二.代码解毒
整体框图
raft.py
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from core.extractor import BasicEncoder#, SmallEncoder#feature network, context network
from core.corr import CorrBlock, AlternateCorrBlock#相关性查找表
from core.update import BasicUpdateBlock#, SmallUpdateBlock#ConVGRU
from core.utils.utils import bilinear_sampler, coords_grid, upflow8#计算flow、上采样
#
import matplotlib.pyplot as plt
from core.utils import flow_viz
import cv2
#https://www.cnblogs.com/jimchen1218/p/14315008.html
try:
#利用with语句，在autocast实例的上下文范围内，进行模型的前向推理和loss计算
autocast = torch.cuda.amp.autocast
except:
# dummy autocast for PyTorch < 1.6
class autocast:
def __init__(self, enabled):
pass
def __enter__(self):
pass
def __exit__(self, *args):
pass
class RAFT(nn.Module):
def __init__(self, args):
super(RAFT, self).__init__()
self.args = args
#①初始化参数
"""
if args.small:
self.hidden_dim = hdim = 96
self.context_dim = cdim = 64
args.corr_levels = 4
args.corr_radius = 3
"""
#这里选用4.8M的RAFT框架
self.hidden_dim = hdim = 128
self.context_dim = cdim = 128
#相关体数
args.corr_levels = 4
args.corr_radius = 4
#dropout值
if 'dropout' not in self.args:
self.args.dropout = 0
if 'alternate_corr' not in self.args:
self.args.alternate_corr = False
#②子网络架构
# feature network, context network, and update block
"""
if args.small:
self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout)
self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout)
self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
"""
self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout)
self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
# def freeze_bn(self):
# for m in self.modules():
# if isinstance(m, nn.BatchNorm2d):#如果m的类型与nn.BatchNorm2d的类型相同则返回 True，否则返回 False
# m.eval()
#③初始化光流f0
#feature encoder network outputs:features at 1/8 resolution(R^H*W*3->G^H/8*W/8*256)
# Flow is represented as difference between two coordinate grids flow = coords1 - coords0
def initialize_flow(self, img):
N, C, H, W = img.shape#这里cam[0]:torch.Size([1, 3, 160, 320])
coords0 = coords_grid(N, H//8, W//8).to(img.device)# (1，1,20,40)
coords1 = coords_grid(N, H//8, W//8).to(img.device)# (1，1,20,40)
return coords0, coords1
#④上采样：mask * up_flow
#Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination
#提取每个像素点以及周围的 8 邻域像素点特征（总共 9 个像素点）重新排列到 channel 维度上
def upsample_flow(self, flow, mask):
#1)mask
N, _, H, W = flow.shape#这里torch.Size([1, 2, 20, 40])
mask = mask.view(N, 1, 9, 8, 8, H, W) # (N,9*8*8,20,40) -> (N,1,9,8,8,20,40)
mask = torch.softmax(mask, dim=2) # 权重归一化
#2)up_flow
#8*flow:上采样后图像的尺度变大了，为了匹配尺度增大的像素坐标，光流(flow=coords1 - coords0)也要按同样的倍率（8 倍）上采样
up_flow = F.unfold(8 * flow, [3,3], padding=1) #每一列的元素为滑动窗口(只卷不积)依次所覆盖的内容 (b,2,h,w) -> (b,2*3*3,h*w)
up_flow = up_flow.view(N, 2, 9, 1, 1, H, W) # (b,2*3*3,h*w) -> (b,2,9,1,1,h,w)
up_flow = torch.sum(mask * up_flow, dim=2) # (b,1,9,8,8,h,w) * (b,2,9,1,1,h,w) -> (b,2,9,8,8,h,w) ->(sum) (b,2,8,8,h,w)
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3) # (b,2,8,8,h,w) -> (b,2,h,8,w,8)
return up_flow.reshape(N, 2, 8*H, 8*W) # (b,2,h,8,w,8) -> (b,2,8h,8w)
##Estimate optical flow between pair of frames
def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
# step 1：网络输入-图像预处理
image1 = 2 * (image1 / 255.0) - 1.0#图像归一化
image2 = 2 * (image2 / 255.0) - 1.0#图像归一化
image1 = image1.contiguous()#类似深拷贝；调用contiguous()时，会强制拷贝一份tensor，让它的布局和从头创建的一模一样，但是两个tensor完全没有联系
image2 = image2.contiguous()
hdim = self.hidden_dim#这里128
cdim = self.context_dim#这里128
#step 2：Feature Encoder 提取两图特征（权值共享）
with autocast():
fmap1, fmap2 = self.fnet([image1, image2])
fmap1 = fmap1.float()
fmap2 = fmap2.float()
# step 3：初始化相关性查找表时，调用 __init__() 函数；
if self.args.alternate_corr:
corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)#这里self.args.corr_radius=4
else:
corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
# step 4：Context Encoder提取第一帧图特征
#GRU输入特征之一：net为GRU的隐状态，inp后续与其他特征结合作为 GRU 的一般输入
with autocast():
cnet = self.cnet(image1)
net, inp = torch.split(cnet, [hdim, cdim], dim=1)
net = torch.tanh(net)
inp = torch.relu(inp)
# step 5：更新光流
# 初始化光流的坐标信息
# coords0 为初始时刻的坐标，coords1 为当前迭代的坐标
coords0, coords1 = self.initialize_flow(image1)#此处两坐标数值相等
if flow_init is not None:
coords1 = coords1 + flow_init
flow_predictions = [] #key：对每次迭代的光流都进行了上采样
for itr in range(iters):
#1）初始的coords1
coords1 = coords1.detach()#key:使之切断反向传播，不具有更新属性
#2）更新的coords1（coords0做基准一直不变）
corr = corr_fn(coords1) #从相关性查找表中获取当前坐标的对应特征（查找对应特征时，调用 __call__() 函数）
flow = coords1 - coords0
#a）self.update_block的输入：context网络的输出（net, inp）-帧1信息；从相关性查找表中获取各坐标的对应特征-帧1、2信息；flow-包含抽象更新趋势delta_flow
with autocast():
net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)# 计算当前迭代的光流
#b）F(t+1) = F(t) + \Delta(t)
coords1 = coords1 + delta_flow#更新光流
#3）upsample predictions（目的：训练网络）
# step 6：上采样光流（此处为了训练网络，对每次迭代的光流都进行了上采样，实际 inference 时，只需要保留最后一次迭代后的上采样）
if up_mask is None:
flow_up = upflow8(coords1 - coords0)
else:
flow_up = self.upsample_flow(coords1 - coords0, up_mask)
flow_predictions.append(flow_up)
if test_mode:
return coords1 - coords0, flow_up
return flow_predictions
extractor.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class ResidualBlock(nn.Module):
def __init__(self, in_planes, planes, norm_fn='group', stride=1):
super(ResidualBlock, self).__init__()
#1)定义残差内块：卷积运算、标准化函数和激活函数
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
self.relu = nn.ReLU(inplace=True)#inplace = True 时,会修改输入对象的值,所以打印出对象存储地址相同,类似于C语言的址传递
# [N,C,H*W]
# nn.BatchNorm2d():统计C_i上所有N张图片(H*W)的均值和方差
# nn.LayerNorm2d():统计N_i上所有C通道图片(H*W)的均值和方差
# nn.InstanceNorm2d():统计(C_i,N_j)上图片(H*W)的均值和方差
# nn.GroupNorm():统计(C_i1,C_i2,······;N_j)上图片(H*W)的均值和方差,介于<IN,LN>
num_groups = planes // 8
if norm_fn == 'group':
self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)#GroupNorm 不会改变输入张量的shape，它只是按照group做normalization
self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
if not stride == 1:
self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
elif norm_fn == 'batch':
self.norm1 = nn.BatchNorm2d(planes)
self.norm2 = nn.BatchNorm2d(planes)
if not stride == 1:
self.norm3 = nn.BatchNorm2d(planes)
elif norm_fn == 'instance':
self.norm1 = nn.InstanceNorm2d(planes)
self.norm2 = nn.InstanceNorm2d(planes)
if not stride == 1:
self.norm3 = nn.InstanceNorm2d(planes)
elif norm_fn == 'none':
self.norm1 = nn.Sequential()
self.norm2 = nn.Sequential()
if not stride == 1:
self.norm3 = nn.Sequential()
#2)定义残差短路块：nn.Sequential() create a small sequential model
if stride == 1:
self.downsample = None
else:
self.downsample = nn.Sequential(
nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride)
, self.norm3)
def forward(self, x):
#残差内块
y = x
y = self.relu(self.norm1(self.conv1(y)))
y = self.relu(self.norm2(self.conv2(y)))
#残差短路块
if self.downsample is not None:
x = self.downsample(x)
#输出
return self.relu(x+y)
class BasicEncoder(nn.Module):
def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
super(BasicEncoder, self).__init__()
#定义标准化函数
self.norm_fn = norm_fn
if self.norm_fn == 'group':
self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
elif self.norm_fn == 'batch':
self.norm1 = nn.BatchNorm2d(64)
elif self.norm_fn == 'instance':
self.norm1 = nn.InstanceNorm2d(64)
elif self.norm_fn == 'none':
self.norm1 = nn.Sequential()
#定义dropout
self.dropout = None
if dropout > 0:
self.dropout = nn.Dropout2d(p=dropout)
#1)key-key-key:网络架构
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
self.relu1 = nn.ReLU(inplace=True)
self.in_planes = 64
self.layer1 = self._make_layer(64, stride=1)
self.layer2 = self._make_layer(96, stride=2)
self.layer3 = self._make_layer(128, stride=2)
#output convolution
self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
#2)key-key-key：参数初始化的不同选择（对于卷积结构、标准化结构）
for m in self.modules():
if isinstance(m, nn.Conv2d):
#using a normal distribution. The resulting tensor will have values sampled from N(0,std^2)
#mode='fan_out/in':要么选择保留了向前传递中权重方差的大小。选择保留了向后传递的大小。
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
if m.weight is not None:
nn.init.constant_(m.weight, 1)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
def _make_layer(self, dim, stride=1):
#1)调用ResidualBlock，并完成（输出-下次输入）的维度转换
layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
self.in_planes = dim
#2)序列化Block
layers = (layer1, layer2)
return nn.Sequential(*layers)
def forward(self, x):
# if input is list, combine batch dimension
is_list = isinstance(x, tuple) or isinstance(x, list)
if is_list:
batch_dim = x[0].shape[0]
x = torch.cat(x, dim=0)
#1)
x = self.conv1(x)
x = self.norm1(x)
x = self.relu1(x)
#2)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
#3)
x = self.conv2(x)
if self.training and self.dropout is not None:
x = self.dropout(x)
# if input is list, devide batch dimension
if is_list:
x = torch.split(x, [batch_dim, batch_dim], dim=0)
return x
update.py
import torch
import torch.nn as nn
import torch.nn.functional as F
#Flow与Corr经过卷积而结合
class BasicMotionEncoder(nn.Module):
def __init__(self, args):
super(BasicMotionEncoder, self).__init__()
#
cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
#
self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
#
self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1)
def forward(self, flow, corr):
#1)
cor = F.relu(self.convc1(corr))
cor = F.relu(self.convc2(cor))
#2)
flo = F.relu(self.convf1(flow))
flo = F.relu(self.convf2(flo))
#3)
cor_flo = torch.cat([cor, flo], dim=1)
out = F.relu(self.conv(cor_flo))
#4)
return torch.cat([out, flow], dim=1)
class SepConvGRU(nn.Module):
def __init__(self, hidden_dim=128, input_dim=192+128):
super(SepConvGRU, self).__init__()
self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
#h=net 注：context网络的输出（net, inp）-帧1信息；
#x=[inp, motion_features] 注：motion_features从BasicMotionEncoder模块得到
#两层GRU：ht-1->ht
def forward(self, h, x):
# 将 3x3 卷积替换成 1x5 和 5x1 的两次卷积，在不提高参数量的情况下增大感受野，下面使用的数学计算见 GRU 公式
# horizontal
hx = torch.cat([h, x], dim=1)
z = torch.sigmoid(self.convz1(hx))
r = torch.sigmoid(self.convr1(hx))
q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))
h = (1-z) * h + z * q
# vertical
hx = torch.cat([h, x], dim=1)
z = torch.sigmoid(self.convz2(hx))
r = torch.sigmoid(self.convr2(hx))
q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))
h = (1-z) * h + z * q
return h
class FlowHead(nn.Module):
def __init__(self, input_dim=128, hidden_dim=256):
super(FlowHead, self).__init__()
self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
return self.conv2(self.relu(self.conv1(x)))
#key-key-key
class BasicUpdateBlock(nn.Module):
def __init__(self, args, hidden_dim=128, input_dim=128):
super(BasicUpdateBlock, self).__init__()
self.args = args
#
self.encoder = BasicMotionEncoder(args)
#
self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
#
self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
self.mask = nn.Sequential(
nn.Conv2d(128, 256, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, 64*9, 1, padding=0))
#self.update_block的输入：
# context网络的输出（net, inp）-帧1信息；
# 从相关性查找表中获取各坐标的对应特征 - 帧1、2信息；
# flow - 包含抽象更新趋势delta_flow
def forward(self, net, inp, corr, flow, upsample=True):
#1)BasicMotionEncoder模块
motion_features = self.encoder(flow, corr)# 结合光流和相关性图提取特征
#2)SepConvGRU模块
inp = torch.cat([inp, motion_features], dim=1)# 连接 Context Encoder 提取的特征和上面提取的特征
net = self.gru(net, inp) # GRU 迭代(iters次)，更新隐状态 net
#3)FlowHead模块
delta_flow = self.flow_head(net) # 由隐状态得到光流残差
#4)上采样
# scale mask to balence gradients
mask = .25 * self.mask(net) # 由隐状态得到上采样 mask
return net, mask, delta_flow
4960
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
