import torch
import torch.nn as nn
import torch.nn.functional as F
from core.update import BasicMultiUpdateBlock
from core.extractor import BasicEncoder, MultiBasicEncoder, ResidualBlock
from core.corr import CorrBlock1D, PytorchAlternateCorrBlock1D, CorrBlockFast1D, AlternateCorrBlock
from core.utils.utils import coords_grid, upflow8
import os
import argparse
try:
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 RAFTStereo(nn.Module):
def __init__(self, args):
super().__init__()
self.args = args
context_dims = args.hidden_dims
self.cnet = MultiBasicEncoder(output_dim=[args.hidden_dims, context_dims], norm_fn=args.context_norm, downsample=args.n_downsample)
self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.hidden_dims)
self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, 3, padding=3//2) for i in range(self.args.n_gru_layers)])
if args.shared_backbone:
self.conv2 = nn.Sequential(
ResidualBlock(128, 128, 'instance', stride=1),
nn.Conv2d(128, 256, 3, padding=1))
else:
self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', downsample=args.n_downsample)
def freeze_bn(self):
for m in self.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
def initialize_flow(self, img):
""" Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
N, _, H, W = img.shape
coords0 = coords_grid(N, H, W).to(img.device)
coords1 = coords_grid(N, H, W).to(img.device)
return coords0, coords1
def upsample_flow(self, flow, mask):#关于感受野
""" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
N, D, H, W = flow.shape
factor = 2 ** self.args.n_downsample
mask = mask.view(N, 1, 9, factor, factor, H, W)
mask = torch.softmax(mask, dim=2)
up_flow = F.unfold(factor * flow, [3,3], padding=1)
up_flow = up_flow.view(N, D, 9, 1, 1, H, W)
up_flow = torch.sum(mask * up_flow, dim=2)
up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
return up_flow.reshape(N, D, factor*H, factor*W)
def forward(self, image1, image2, iters=12, flow_init=None, test_mode=False):
""" Estimate optical flow between pair of frames """
image1 = (2 * (image1 / 255.0) - 1.0).contiguous()
image2 = (2 * (image2 / 255.0) - 1.0).contiguous()
# run the context network,交叉预测
with autocast(enabled=self.args.mixed_precision):
if self.args.shared_backbone:
*cnet_list, x = self.cnet(torch.cat((image1, image2), dim=0), dual_inp=True, num_layers=self.args.n_gru_layers)
fmap1, fmap2 = self.conv2(x).split(dim=0, split_size=x.shape[0]//2)
else:
cnet_list = self.cnet(image1, num_layers=self.args.n_gru_layers)
fmap1, fmap2 = self.fnet([image1, image2])
net_list = [torch.tanh(x[0]) for x in cnet_list]
inp_list = [torch.relu(x[1]) for x in cnet_list]
# Rather than running the GRU's conv layers on the context features multiple times, we do it once at the beginning
inp_list = [list(conv(i).split(split_size=conv.out_channels//3, dim=1)) for i,conv in zip(inp_list, self.context_zqr_convs)]
if self.args.corr_implementation == "reg": # Default
corr_block = CorrBlock1D
fmap1, fmap2 = fmap1.float(), fmap2.float()
elif self.args.corr_implementation == "alt": # More memory efficient than reg
corr_block = PytorchAlternateCorrBlock1D
fmap1, fmap2 = fmap1.float(), fmap2.float()
elif self.args.corr_implementation == "reg_cuda": # Faster version of reg
corr_block = CorrBlockFast1D
elif self.args.corr_implementation == "alt_cuda": # Faster version of alt
corr_block = AlternateCorrBlock
corr_fn = corr_block(fmap1, fmap2, radius=self.args.corr_radius, num_levels=self.args.corr_levels)
coords0, coords1 = self.initialize_flow(net_list[0])
if flow_init is not None:
coords1 = coords1 + flow_init
flow_predictions = []
for itr in range(iters):
coords1 = coords1.detach()
corr = corr_fn(coords1) # index correlation volume
flow = coords1 - coords0
with autocast(enabled=self.args.mixed_precision):
if self.args.n_gru_layers == 3 and self.args.slow_fast_gru: # Update low-res GRU
net_list = self.update_block(net_list, inp_list, iter32=True, iter16=False, iter08=False, update=False)
if self.args.n_gru_layers >= 2 and self.args.slow_fast_gru:# Update low-res GRU and mid-res GRU
net_list = self.update_block(net_list, inp_list, iter32=self.args.n_gru_layers==3, iter16=True, iter08=False, update=False)
net_list, up_mask, delta_flow = self.update_block(net_list, inp_list, corr, flow, iter32=self.args.n_gru_layers==3, iter16=self.args.n_gru_layers>=2)
# in stereo mode, project flow onto epipolar
# delta_flow[:,1] = 0.0
# F(t+1) = F(t) + \Delta(t)
coords1 = coords1 + delta_flow
# We do not need to upsample or output intermediate results in test_mode
if test_mode and itr < iters-1:
continue
# upsample predictions
if up_mask is None:
flow_up = upflow8(coords1 - coords0)
else:
flow_up = self.upsample_flow(coords1 - coords0, up_mask)
flow_up = flow_up[:,:1]
flow_predictions.append(flow_up)
if test_mode:
return coords1 - coords0, flow_up
return flow_predictions
parser = argparse.ArgumentParser()
parser.add_argument('--restore_ckpt', help="restore checkpoint")
parser.add_argument('--save_numpy', action='store_true', help='save output as numpy arrays')
parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames")
parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames")
# parser.add_argument('-l', '--left_imgs', help="path to all first (left) frames", default="datasets/Middlebury/MiddEval3/testH/*/im0.png")
# parser.add_argument('-r', '--right_imgs', help="path to all second (right) frames", default="datasets/Middlebury/MiddEval3/testH/*/im1.png")
parser.add_argument('--output_directory', help="directory to save output", default="test_single")
parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
parser.add_argument('--valid_iters', type=int, default=7, help='number of flow-field updates during forward pass')
# Architecture choices
parser.add_argument('--hidden_dims', nargs='+', type=int, default=[128]*3, help="hidden state and context dimensions")
parser.add_argument('--corr_implementation', choices=["reg", "alt", "reg_cuda", "alt_cuda"], default="reg", help="correlation volume implementation")
parser.add_argument('--shared_backbone', action='store_true', help="use a single backbone for the context and feature encoders")
parser.add_argument('--corr_levels', type=int, default=4, help="number of levels in the correlation pyramid")
parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
parser.add_argument('--n_downsample', type=int, default=3, help="resolution of the disparity field (1/2^K)")
parser.add_argument('--context_norm', type=str, default="batch", choices=['group', 'batch', 'instance', 'none'], help="normalization of context encoder")
parser.add_argument('--slow_fast_gru', action='store_true', help="iterate the low-res GRUs more frequently")
parser.add_argument('--n_gru_layers', type=int, default=2, help="number of hidden GRU levels")
parser.set_defaults(save_numpy=True)
parser.set_defaults(mixed_precision=True)
parser.set_defaults(shared_backbone=True)
parser.set_defaults(slow_fast_gru=True)
args = parser.parse_args()
if __name__ == '__main__':
net = RAFTStereo(args)
# 加载预训练权重
# 注意:这里预训练权重是用 DataParallel 训练的
pretrained_weights = torch.load("/home/xxx/RAFT-Stereo-main/models/raftstereo-realtime.pth")
# 解包 DataParallel 包装的权重
# 这一步是为了获取原始模型的权重
if 'module' in list(pretrained_weights.keys())[0]:
new_state_dict = {}
for k, v in pretrained_weights.items():
name = k[7:] # 删除 "module.",假设 "module." 有 7 个字符
new_state_dict[name] = v
else:
# 如果键中没有 "module." 前缀,则不需要做任何事情
new_state_dict = pretrained_weights
net.load_state_dict(new_state_dict)
# net.eval().cpu() # 模式是cpu还是gpu gpu:net.eval().cudu()
net.eval().cuda()
print('to gpu')
input1 = torch.randn(1, 3, 640, 480).cuda() # 同上,gup则 .cudu() 此处的参数必须和模型的一致否则会出错,切记切记
input2 = torch.randn(1, 3, 640, 480).cuda() # 同上,gup则 .cudu() 此处的参数必须和模型的一致否则会出错,切记切记
input_names = ["input1", "input2"]
output_names = ["output"]
torch.onnx.export(net,
(input1 , input2),
"/home/xxx/RKNN_project/rknn_model_zoo-main/RAFT/model.onnx", # 输出的onnx路径
verbose=True,
input_names=input_names,
output_names=output_names,
opset_version=16)
注意:以上代码中,raft模型架构为微调之后,非原文章模型,仅作参考,并且请注意文件结构。