def __init__(self,args, batchNorm=True, div_flow = 20):
super(FlowNetC,self).__init__()
self.batchNorm = batchNorm
self.div_flow = div_flow
self.conv1 = conv(self.batchNorm, 3, 64, kernel_size=7, stride=2)
self.conv2 = conv(self.batchNorm, 64, 128, kernel_size=5, stride=2)
self.conv3 = conv(self.batchNorm, 128, 256, kernel_size=5, stride=2)
self.conv_redir = conv(self.batchNorm, 256, 32, kernel_size=1, stride=1)
if args.fp16:
self.corr = nn.Sequential(
tofp32(),
Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1),
tofp16())
else:
self.corr = Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1)
self.corr_activation = nn.LeakyReLU(0.1,inplace=True)
self.conv3_1 = conv(self.batchNorm, 473, 256)
self.conv4 = conv(self.batchNorm, 256, 512, stride=2)
self.conv4_1 = conv(self.batchNorm, 512, 512)
self.conv5 = conv(self.batchNorm, 512, 512, stride=2)
self.conv5_1 = conv(self.batchNorm, 512, 512)
self.conv6 = conv(self.batchNorm, 512, 1024, stride=2)
self.conv6_1 = conv(self.batchNorm,1024, 1024)
self.deconv5 = deconv(1024,512)
self.deconv4 = deconv(1026,256)
self.deconv3 = deconv(770,128)
self.deconv2 = deconv(386,64)
self.predict_flow6 = predict_flow(1024)
self.predict_flow5 = predict_flow(1026)
self.predict_flow4 = predict_flow(770)
self.predict_flow3 = predict_flow(386)
self.predict_flow2 = predict_flow(194)
self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
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)
self.upsample1 = nn.Upsample(scale_factor=4, mode='bilinear')
def __init__(self, args):
super(Net, self).__init__()
self.args = args
self.feature_pyramid_extractor = FeaturePyramidExtractor(args).to(args.device)
self.warping_layer = WarpingLayer(args)
if args.corr == 'CostVolumeLayer':
self.corr = CostVolumeLayer(args)
else:
self.corr = Correlation(pad_size = args.search_range * 2 + 1, kernel_size = 1, max_displacement = args.search_range * 2 + 1, stride1 = 1, stride2 = 2, corr_multiply = 1).to(args.device)
self.flow_estimators = []
for l, ch in enumerate(args.lv_chs[::-1] + [3]):
layer = OpticalFlowEstimator(args, ch + (args.search_range*2+1)**2 + 2).to(args.device)
self.add_module(f'FlowEstimator(Lv{l})', layer)
self.flow_estimators.append(layer)
self.context_network = ContextNetwork(args, 3 + 2).to(args.device)
# self.context_networks = []
# for l, ch in enumerate(args.lv_chs[::-1]):
# layer = ContextNetwork(args, ch + 2).to(args.device)
# self.add_module(f'ContextNetwork(Lv{l})', layer)
# self.context_networks.append(layer)
# init
for m in self.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None: nn.init.uniform_(m.bias)
nn.init.xavier_uniform_(m.weight)
if isinstance(m, nn.ConvTranspose2d):
if m.bias is not None: nn.init.uniform_(m.bias)
nn.init.xavier_uniform_(m.weight)
def __init__(self, pretrained=None):
super(Model, self).__init__()
idx = [list(range(n, -1, -9)) for n in range(80,71,-1)]
idx = list(np.array(idx).flatten())
self.idx_fwd = Variable(torch.LongTensor(np.array(idx)).cuda(), requires_grad=False)
self.idx_bwd = Variable(torch.LongTensor(np.array(list(reversed(idx)))).cuda(), requires_grad=False)
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
self.softmax2d = nn.Softmax2d()
self.conv1a = conv_feat_block(3,16)
self.conv1b = conv_feat_block(3,16)
self.conv1c = conv_feat_block(3,16)
self.conv2a = conv_feat_block(16,32)
self.conv2b = conv_feat_block(16,32)
self.conv2c = conv_feat_block(16,32)
self.conv3a = conv_feat_block(32,64)
self.conv3b = conv_feat_block(32,64)
self.conv3c = conv_feat_block(32,64)
self.conv4a = conv_feat_block(64,96)
self.conv4b = conv_feat_block(64,96)
self.conv4c = conv_feat_block(64,96)
self.conv5a = conv_feat_block(96,128)
self.conv5b = conv_feat_block(96,128)
self.conv5c = conv_feat_block(96,128)
self.conv6a = conv_feat_block(128,192)
self.conv6b = conv_feat_block(128,192)
self.conv6c = conv_feat_block(128,192)
self.corr = Correlation(pad_size=4, kernel_size=1, max_displacement=4, stride1=1, stride2=1, corr_multiply=1)
self.decoder_fwd6 = conv_dec_block(162)
self.decoder_bwd6 = conv_dec_block(162)
self.decoder_fwd5 = conv_dec_block(292)
self.decoder_bwd5 = conv_dec_block(292)
self.decoder_fwd4 = conv_dec_block(260)
self.decoder_bwd4 = conv_dec_block(260)
self.decoder_fwd3 = conv_dec_block(228)
self.decoder_bwd3 = conv_dec_block(228)
self.decoder_fwd2 = conv_dec_block(196)
self.decoder_bwd2 = conv_dec_block(196)
self.decoder_occ6 = conv_dec_block(354)
self.decoder_occ5 = conv_dec_block(292)
self.decoder_occ4 = conv_dec_block(260)
self.decoder_occ3 = conv_dec_block(228)
self.decoder_occ2 = conv_dec_block(196)
if pretrained is not None:
self.load_state_dict(torch.load(pretrained))
print('Model loaded from ', pretrained)
def __init__(self, input_channels = 6, output_channels=2, batch_norm=False,
div_flow=1):
super(FlowNetC,self).__init__()
self.batch_norm = batch_norm
self.input_channels = input_channels
self.div_flow = div_flow
self.conv1 = conv(input_channels/2, 64, kernel_size=7, stride=2, batch_norm=self.batch_norm)
self.conv2 = conv(64, 128, kernel_size=5, stride=2, batch_norm=self.batch_norm)
self.conv3 = conv(128, 256, kernel_size=5, stride=2, batch_norm=self.batch_norm)
self.conv_redir = conv(256, 32, kernel_size=1, stride=1, batch_norm=self.batch_norm)
self.corr = Correlation(pad_size=20, kernel_size=1, max_displacement=20, stride1=1, stride2=2, corr_multiply=1)
self.corr_activation = nn.LeakyReLU(0.1,inplace=False)
self.conv3_1 = conv(473, 256, batch_norm=self.batch_norm)
self.conv4 = conv(256, 512, stride=2, batch_norm=self.batch_norm)
self.conv4_1 = conv(512, 512, batch_norm=self.batch_norm)
self.conv5 = conv(512, 512, stride=2, batch_norm=self.batch_norm)
self.conv5_1 = conv(512, 512, batch_norm=self.batch_norm)
self.conv6 = conv(512, 1024, stride=2, batch_norm=self.batch_norm)
self.conv6_1 = conv(1024, 1024, batch_norm=self.batch_norm)
self.deconv5 = deconv(1024,512)
self.deconv4 = deconv(1024 + output_channels, 256)
self.deconv3 = deconv(768 + output_channels, 128)
self.deconv2 = deconv(384 + output_channels, 64)
self.predict_flow6 = predict_flow(1024, output_channels)
self.predict_flow5 = predict_flow(1024 + output_channels, output_channels)
self.predict_flow4 = predict_flow(768 + output_channels, output_channels)
self.predict_flow3 = predict_flow(384 + output_channels, output_channels)
self.predict_flow2 = predict_flow(192 + output_channels, output_channels)
self.upsampled_flow6_to_5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow5_to_4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow4_to_3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsampled_flow3_to_2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
init.uniform_(m.weight.data, 1.0, 0.02)
init.constant_(m.bias.data, 0.0)
def __init__(self, input_channels=6, div_flow=1, md=4, deconv_relu=False):
"""
input: md --- maximum displacement (for correlation. default: 4), after warpping
"""
super(PWCDCNet_archarticle, self).__init__()
self.input_channels = input_channels
self.div_flow = div_flow
self.deconv_relu = deconv_relu
self.conv1a = conv(int(input_channels / 2),
16,
kernel_size=3,
stride=2)
self.conv1aa = conv(16, 16, kernel_size=3, stride=1)
self.conv1b = conv(16, 16, kernel_size=3, stride=1)
self.conv2a = conv(16, 32, kernel_size=3, stride=2)
self.conv2aa = conv(32, 32, kernel_size=3, stride=1)
self.conv2b = conv(32, 32, kernel_size=3, stride=1)
self.conv3a = conv(32, 64, kernel_size=3, stride=2)
self.conv3aa = conv(64, 64, kernel_size=3, stride=1)
self.conv3b = conv(64, 64, kernel_size=3, stride=1)
self.conv4a = conv(64, 96, kernel_size=3, stride=2)
self.conv4aa = conv(96, 96, kernel_size=3, stride=1)
self.conv4b = conv(96, 96, kernel_size=3, stride=1)
self.conv5a = conv(96, 128, kernel_size=3, stride=2)
self.conv5aa = conv(128, 128, kernel_size=3, stride=1)
self.conv5b = conv(128, 128, kernel_size=3, stride=1)
self.conv6aa = conv(128, 196, kernel_size=3, stride=2)
self.conv6a = conv(196, 196, kernel_size=3, stride=1)
self.conv6b = conv(196, 196, kernel_size=3, stride=1)
self.corr = Correlation(pad_size=md,
kernel_size=1,
max_displacement=md,
stride1=1,
stride2=1,
corr_multiply=1)
self.leakyRELU = nn.LeakyReLU(0.1)
nd = (2 * md + 1)**2
dd = np.cumsum([128, 128, 96, 64, 32])
od = nd
self.conv6_0 = conv(od, 128, kernel_size=3, stride=1)
self.conv6_1 = conv(od + dd[0], 128, kernel_size=3, stride=1)
self.conv6_2 = conv(od + dd[1], 96, kernel_size=3, stride=1)
self.conv6_3 = conv(od + dd[2], 64, kernel_size=3, stride=1)
self.conv6_4 = conv(od + dd[3], 32, kernel_size=3, stride=1)
self.predict_flow6 = predict_flow(od + dd[4])
self.deconv6 = deconv_pwc_article(2,
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
self.upfeat6 = deconv_pwc_article(od + int(dd[4]),
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
od = nd + 128 + 4
self.conv5_0 = conv(od, 128, kernel_size=3, stride=1)
self.conv5_1 = conv(od + dd[0], 128, kernel_size=3, stride=1)
self.conv5_2 = conv(od + dd[1], 96, kernel_size=3, stride=1)
self.conv5_3 = conv(od + dd[2], 64, kernel_size=3, stride=1)
self.conv5_4 = conv(od + dd[3], 32, kernel_size=3, stride=1)
self.predict_flow5 = predict_flow(od + dd[4])
self.deconv5 = deconv_pwc_article(2,
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
self.upfeat5 = deconv_pwc_article(od + int(dd[4]),
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
od = nd + 96 + 4
self.conv4_0 = conv(od, 128, kernel_size=3, stride=1)
self.conv4_1 = conv(od + dd[0], 128, kernel_size=3, stride=1)
self.conv4_2 = conv(od + dd[1], 96, kernel_size=3, stride=1)
self.conv4_3 = conv(od + dd[2], 64, kernel_size=3, stride=1)
self.conv4_4 = conv(od + dd[3], 32, kernel_size=3, stride=1)
self.predict_flow4 = predict_flow(od + dd[4])
self.deconv4 = deconv_pwc_article(2,
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
self.upfeat4 = deconv_pwc_article(od + int(dd[4]),
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
od = nd + 64 + 4
self.conv3_0 = conv(od, 128, kernel_size=3, stride=1)
self.conv3_1 = conv(od + dd[0], 128, kernel_size=3, stride=1)
self.conv3_2 = conv(od + dd[1], 96, kernel_size=3, stride=1)
self.conv3_3 = conv(od + dd[2], 64, kernel_size=3, stride=1)
self.conv3_4 = conv(od + dd[3], 32, kernel_size=3, stride=1)
self.predict_flow3 = predict_flow(od + dd[4])
self.deconv3 = deconv_pwc_article(2,
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
self.upfeat3 = deconv_pwc_article(od + int(dd[4]),
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
od = nd + 32 + 4
self.conv2_0 = conv(od, 128, kernel_size=3, stride=1)
self.conv2_1 = conv(od + dd[0], 128, kernel_size=3, stride=1)
self.conv2_2 = conv(od + dd[1], 96, kernel_size=3, stride=1)
self.conv2_3 = conv(od + dd[2], 64, kernel_size=3, stride=1)
self.conv2_4 = conv(od + dd[3], 32, kernel_size=3, stride=1)
self.predict_flow2 = predict_flow(od + dd[4])
self.deconv2 = deconv_pwc_article(2,
2,
kernel_size=4,
stride=2,
padding=1,
relu=self.deconv_relu)
self.dc_conv1 = conv(od + dd[4],
128,
kernel_size=3,
stride=1,
padding=1,
dilation=1)
self.dc_conv2 = conv(128,
128,
kernel_size=3,
stride=1,
padding=2,
dilation=2)
self.dc_conv3 = conv(128,
128,
kernel_size=3,
stride=1,
padding=4,
dilation=4)
self.dc_conv4 = conv(128,
96,
kernel_size=3,
stride=1,
padding=8,
dilation=8)
self.dc_conv5 = conv(96,
64,
kernel_size=3,
stride=1,
padding=16,
dilation=16)
self.dc_conv6 = conv(64,
32,
kernel_size=3,
stride=1,
padding=1,
dilation=1)
self.dc_conv7 = predict_flow(32)
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
if m.bias is not None:
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm2d):
init.uniform_(m.weight.data, 1.0, 0.02)
init.constant_(m.bias.data, 0.0)
def __init__(self, div_flow=20):
super(PwcNet, self).__init__()
self.div_flow = div_flow
self.py_conv1 = conv(False, 3, 16, stride=2, bias=True)
self.py_conv1_b = conv(False, 16, 16, bias=True)
self.py_conv2 = conv(False, 16, 32, stride=2, bias=True)
self.py_conv2_b = conv(False, 32, 32, bias=True)
self.py_conv3 = conv(False, 32, 64, stride=2, bias=True)
self.py_conv3_b = conv(False, 64, 64, bias=True)
self.py_conv4 = conv(False, 64, 96, stride=2, bias=True)
self.py_conv4_b = conv(False, 96, 96, bias=True)
self.py_conv5 = conv(False, 96, 128, stride=2, bias=True)
self.py_conv5_b = conv(False, 128, 128, bias=True)
self.py_conv6 = conv(False, 128, 196, stride=2, bias=True)
self.py_conv6_b = conv(False, 196, 196, bias=True)
# fixed size correlation layer
# the correlation layer output is always (4+1)*(4+1) = 81
self.corr = Correlation(pad_size=4,
kernel_size=1,
max_displacement=4,
stride1=1,
stride2=1,
corr_multiply=1)
self.corr_activation = nn.LeakyReLU(0.1, inplace=True)
# every dense block will have output that is
# inplanes + 128 + 128 + 96 + 64 + outplanes = 448
self.dense6 = DenseBlock(False, 81, 32)
# the input channel is 81+128+2+2=213
self.dense5 = DenseBlock(False, 213, 32)
# the input channle is 81+96+2+2 =181
self.dense4 = DenseBlock(False, 181, 32)
# the input channel is 81+64+2+2 =149
self.dense3 = DenseBlock(False, 149, 32)
# the input channel is 81+32+2+2 =117
self.dense2 = DenseBlock(False, 117, 32)
self.upsample_conv6to5 = nn.ConvTranspose2d(81 + 448,
2,
4,
2,
1,
bias=True)
self.upsample_conv5to4 = nn.ConvTranspose2d(213 + 448,
2,
4,
2,
1,
bias=True)
self.upsample_conv4to3 = nn.ConvTranspose2d(181 + 448,
2,
4,
2,
1,
bias=True)
self.upsample_conv3to2 = nn.ConvTranspose2d(149 + 448,
2,
4,
2,
1,
bias=True)
self.predict_flow6 = predict_flow(448 + 81)
self.predict_flow5 = predict_flow(448 + 81 + 128 + 2 + 2)
self.predict_flow4 = predict_flow(448 + 81 + 96 + 2 + 2)
self.predict_flow3 = predict_flow(448 + 81 + 64 + 2 + 2)
self.predict_flow2 = predict_flow(448 + 81 + 32 + 2 + 2)
# all the flow have different upsampling weights
self.upsample_flow6to5 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsample_flow5to4 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsample_flow4to3 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
self.upsample_flow3to2 = nn.ConvTranspose2d(2, 2, 4, 2, 1, bias=True)
# the warping layer
self.resample = Resample2d()
# the context network
self.context_net = nn.Sequential(
conv(False, 117 + 448, 128, bias=True),
conv(False, 128, 128, dilation=2, bias=True),
conv(False, 128, 128, dilation=4, bias=True),
conv(False, 128, 96, dilation=8, bias=True),
conv(False, 96, 64, dilation=16, bias=True),
conv(False, 64, 32, bias=True),
nn.Conv2d(32, 2, 3, 1, 1, bias=True),
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
if m.bias is not None:
m.bias.data.zero_()
init.kaiming_uniform(m.weight)
if isinstance(m, nn.ConvTranspose2d):
if m.bias is not None:
m.bias.data.zero_()
init.kaiming_uniform(m.weight)
self.upsample = nn.Upsample(scale_factor=4, mode='bilinear')