def __init__(self, args, batchNorm=False, div_flow=20.):
super(FlowNet2CS, self).__init__()
self.batchNorm = batchNorm
self.div_flow = div_flow
self.rgb_max = args.rgb_max
self.args = args
self.channelnorm = ChannelNorm()
# First Block (FlowNetC)
self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)
self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')
if args.fp16:
self.resample1 = nn.Sequential(tofp32(), Resample2d(), tofp16())
else:
self.resample1 = Resample2d()
# Block (FlowNetS1)
self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)
self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear')
for m in self.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None:
init.uniform(m.bias)
init.xavier_uniform(m.weight)
if isinstance(m, nn.ConvTranspose2d):
if m.bias is not None:
init.uniform(m.bias)
init.xavier_uniform(m.weight)
def __init__(self,
args,
batchNorm=False,
div_flow=20.,
requires_grad=False):
super(FlowNet2, self).__init__()
self.batchNorm = batchNorm
self.div_flow = div_flow
self.rgb_max = args.rgb_max
self.args = args
self.channelnorm = ChannelNorm()
# First Block (FlowNetC)
self.flownetc = FlowNetC.FlowNetC(args, batchNorm=self.batchNorm)
self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')
if args.fp16:
self.resample1 = nn.Sequential(tofp32(), Resample2d(), tofp16())
else:
self.resample1 = Resample2d()
# Block (FlowNetS1)
self.flownets_1 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)
self.upsample2 = nn.Upsample(scale_factor=4, mode='bilinear')
if args.fp16:
self.resample2 = nn.Sequential(tofp32(), Resample2d(), tofp16())
else:
self.resample2 = Resample2d()
# Block (FlowNetS2)
self.flownets_2 = FlowNetS.FlowNetS(args, batchNorm=self.batchNorm)
# Block (FlowNetSD)
self.flownets_d = FlowNetSD.FlowNetSD(args, batchNorm=self.batchNorm)
self.upsample3 = nn.Upsample(scale_factor=4, mode='nearest')
self.upsample4 = nn.Upsample(scale_factor=4, mode='nearest')
if args.fp16:
self.resample3 = nn.Sequential(tofp32(), Resample2d(), tofp16())
else:
self.resample3 = Resample2d()
if args.fp16:
self.resample4 = nn.Sequential(tofp32(), Resample2d(), tofp16())
else:
self.resample4 = Resample2d()
# Block (FLowNetFusion)
self.flownetfusion = FlowNetFusion.FlowNetFusion(
args, batchNorm=self.batchNorm)
for m in self.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None:
init.uniform_(m.bias)
init.xavier_uniform_(m.weight)
if isinstance(m, nn.ConvTranspose2d):
if m.bias is not None:
init.uniform_(m.bias)
init.xavier_uniform_(m.weight)
# init_deconv_bilinear(m.weight)
if not requires_grad:
for param in self.parameters():
param.requires_grad = False