def build_model(weights):
net = dict()
# T.nnet.abstract_conv.bilinear_upsampling doesn't work properly if not to
# specify a batch size
batch_size = 1
net['input_1'] = InputLayer([batch_size, 3, 384, 512])
net['input_2'] = InputLayer([batch_size, 3, 384, 512])
net['conv1'] = leaky_conv(
net['input_1'], num_filters=64, filter_size=7, stride=2)
net['conv1b'] = leaky_conv(
net['input_2'], num_filters=64, filter_size=7, stride=2,
W=net['conv1'].W, b=net['conv1'].b)
net['conv2'] = leaky_conv(
net['conv1'], num_filters=128, filter_size=5, stride=2)
net['conv2b'] = leaky_conv(
net['conv1b'], num_filters=128, filter_size=5, stride=2,
W=net['conv2'].W, b=net['conv2'].b)
net['conv3'] = leaky_conv(
net['conv2'], num_filters=256, filter_size=5, stride=2)
net['conv3b'] = leaky_conv(
net['conv2b'], num_filters=256, filter_size=5, stride=2,
W=net['conv3'].W, b=net['conv3'].b)
net['corr'] = CorrelationLayer(net['conv3'], net['conv3b'])
net['corr'] = ExpressionLayer(net['corr'], leaky_rectify)
net['conv_redir'] = leaky_conv(
net['conv3'], num_filters=32, filter_size=1, stride=1, pad=0)
net['concat'] = ConcatLayer([net['conv_redir'], net['corr']])
net['conv3_1'] = leaky_conv(net['concat'], num_filters=256, filter_size=3, stride=1)
net['conv4'] = leaky_conv(net['conv3_1'], num_filters=512, filter_size=3, stride=2)
net['conv4_1'] = leaky_conv(net['conv4'], num_filters=512, filter_size=3, stride=1)
net['conv5'] = leaky_conv(net['conv4_1'], num_filters=512, filter_size=3, stride=2)
net['conv5_1'] = leaky_conv(net['conv5'], num_filters=512, filter_size=3, stride=1)
net['conv6'] = leaky_conv(net['conv5_1'], num_filters=1024, filter_size=3, stride=2)
net['conv6_1'] = leaky_conv(net['conv6'], num_filters=1024, filter_size=3, stride=1)
for layer_id in ['1', '2', '3', '_redir', '3_1', '4', '4_1', '5', '5_1', '6', '6_1']:
layer_name = 'conv' + layer_id
print(layer_name, net[layer_name].W.shape.eval(), weights[layer_name][0].shape)
print(layer_name, net[layer_name].b.shape.eval(), weights[layer_name][1].shape)
net[layer_name].W.set_value(weights[layer_name][0])
net[layer_name].b.set_value(weights[layer_name][1])
refine_flow(net, weights)
return net
评论列表
文章目录