Method graph_structure

examples/OpticalFlow/flownet_models.py:379–442

Architecture of FlowNetCorr in Figure 2 of FlowNet 1.0. Args: x: 2CHW.

(self, x1x2)

class FlowNet2C(FlowNetBase):
    def graph_structure(self, x1x2):
        """
        Architecture of FlowNetCorr in Figure 2 of FlowNet 1.0.
        Args:
            x: 2CHW.
        """
        with argscope([tf.layers.conv2d], activation=lambda x: tf.nn.leaky_relu(x, 0.1),
                      padding='valid', strides=2, kernel_size=3,
                      data_format='channels_first'), \
            argscope([tf.layers.conv2d_transpose], padding='same', activation=tf.identity,
                     data_format='channels_first', strides=2, kernel_size=4):

            # extract features
            x = tf.layers.conv2d(pad(x1x2, 3), 64, kernel_size=7, name='conv1')
            conv2 = tf.layers.conv2d(pad(x, 2), 128, kernel_size=5, name='conv2')
            conv3 = tf.layers.conv2d(pad(conv2, 2), 256, kernel_size=5, name='conv3')

            conv2a, _ = tf.split(conv2, 2, axis=0)
            conv3a, conv3b = tf.split(conv3, 2, axis=0)

            corr = correlation(conv3a, conv3b,
                               kernel_size=1,
                               max_displacement=20,
                               stride_1=1,
                               stride_2=2,
                               pad=20, data_format='NCHW')
            corr = tf.nn.leaky_relu(corr, 0.1)

            conv_redir = tf.layers.conv2d(conv3a, 32, kernel_size=1, strides=1, name='conv_redir')

            in_conv3_1 = tf.concat([conv_redir, corr], axis=1, name='in_conv3_1')
            conv3_1 = tf.layers.conv2d(pad(in_conv3_1, 1), 256, name='conv3_1', strides=1)

            x = tf.layers.conv2d(pad(conv3_1, 1), 512, name='conv4')
            conv4 = tf.layers.conv2d(pad(x, 1), 512, name='conv4_1', strides=1)
            x = tf.layers.conv2d(pad(conv4, 1), 512, name='conv5')
            conv5 = tf.layers.conv2d(pad(x, 1), 512, name='conv5_1', strides=1)
            x = tf.layers.conv2d(pad(conv5, 1), 1024, name='conv6')
            conv6 = tf.layers.conv2d(pad(x, 1), 1024, name='conv6_1', strides=1)

            flow6 = tf.layers.conv2d(pad(conv6, 1), 2, name='predict_flow6', strides=1, activation=tf.identity)
            flow6_up = tf.layers.conv2d_transpose(flow6, 2, name='upsampled_flow6_to_5')
            x = tf.layers.conv2d_transpose(conv6, 512, name='deconv5', activation=lambda x: tf.nn.leaky_relu(x, 0.1))

            # return flow6
            concat5 = tf.concat([conv5, x, flow6_up], axis=1, name='concat5')
            flow5 = tf.layers.conv2d(pad(concat5, 1), 2, name='predict_flow5', strides=1, activation=tf.identity)
            flow5_up = tf.layers.conv2d_transpose(flow5, 2, name='upsampled_flow5_to_4')
            x = tf.layers.conv2d_transpose(concat5, 256, name='deconv4', activation=lambda x: tf.nn.leaky_relu(x, 0.1))

            concat4 = tf.concat([conv4, x, flow5_up], axis=1, name='concat4')
            flow4 = tf.layers.conv2d(pad(concat4, 1), 2, name='predict_flow4', strides=1, activation=tf.identity)
            flow4_up = tf.layers.conv2d_transpose(flow4, 2, name='upsampled_flow4_to_3')
            x = tf.layers.conv2d_transpose(concat4, 128, name='deconv3', activation=lambda x: tf.nn.leaky_relu(x, 0.1))

            concat3 = tf.concat([conv3_1, x, flow4_up], axis=1, name='concat3')
            flow3 = tf.layers.conv2d(pad(concat3, 1), 2, name='predict_flow3', strides=1, activation=tf.identity)
            flow3_up = tf.layers.conv2d_transpose(flow3, 2, name='upsampled_flow3_to_2')