Gpu dev

common/forward_warp/forward_warp.py

@@ -1,10 +1,46 @@
 import tensorflow as tf
+from common.utils.tf import load_op_library
+from tensorflow.python.framework import ops
 
 
-DISOCC_THRESH = 0.8
+DISOCC_THRESH = 0.5
 
 
-def forward_warp(features, flow):
+# Load op library.
+mod = load_op_library('forward_warp_op', 'build')
+
+
+def is_forward_warp_cuda():
+    """
+    :return: Bool. Whether the forward_warp is using a custom CUDA op.
+    """
+    return mod is not None
+
+
+def forward_warp(features, flow, splat_variance=0.5):
+    """
+    For an algorithm that gives the same end result, see section 3 in https://arxiv.org/pdf/1711.05890.pdf.
+    Note that the actual implementation here is not n^2, and should be linear in GPU memory.
+    :param features: A Tensor. Features to be warped, of shape [batch_size, H, W, C].
+    :param flow: A Tensor. Un-normalized flow in image pixel units, of shape [batch_size, H, W, 2].
+                 Flow vectors should have (x, y) ordering.
+    :param splat_variance: Float. Variance of the splat. Only used for the CUDA op.
+    """
+    if is_forward_warp_cuda():
+        return mod.forward_warp(features, flow, variance=splat_variance)
+    else:
+        return forward_warp_tf(features, flow)
+
+
+if is_forward_warp_cuda():
+    @ops.RegisterGradient('ForwardWarp')
+    def _ForwardWarpGrad(op, grad):
+        image_grad, flow_grad = mod.forward_warp_grad(
+            grad, op.inputs[0], op.inputs[1], variance=op.get_attr('variance'))
+        return [image_grad, flow_grad]
+
+
+def forward_warp_tf(features, flow):
     """
     For an algorithm that gives the same end result, see section 3 in https://arxiv.org/pdf/1711.05890.pdf.
     Note that the actual implementation here is not n^2, and should be linear in GPU memory.
@@ -88,7 +124,7 @@ def _get_translated_pixels(features, translations):
         return all_indices, all_vals
 
 
-def create_disocclusion_mask(flow):
+def create_disocclusion_mask(flow, splat_variance=1.0):
     """
     Creates a disocclusion mask representing areas that were previously occluded and will become visible.
     This is done by forward warping some ones and thresholding them for visibility.
@@ -97,10 +133,11 @@ def create_disocclusion_mask(flow):
     https://github.com/simonmeister/UnFlow/blob/8bff4939963c7d0adb9435880dc506fb3f988080/src/e2eflow/core/losses.py#L28
     This isn't mentioned in the paper anywhere, but clearly enough, it is in the code.
     :param flow: Tensor of shape [B, H, W, 2].
+    :param splat_variance: Float. Variance of the splat. Only used for the CUDA op.
     :return: Tensor of shape [B, H, W, 1].
     """
     with tf.name_scope('disocclusion_mask'):
         batch, height, width, _ = tf.unstack(tf.shape(flow))
         prewarp_mask = tf.ones([batch, height, width, 1], dtype=tf.float32)
-        forward_warped_mask = forward_warp(prewarp_mask, flow)
+        forward_warped_mask = forward_warp(prewarp_mask, flow, splat_variance=splat_variance)
         return tf.cast(forward_warped_mask < DISOCC_THRESH, dtype=tf.float32)

common/forward_warp/forward_warp_profile.py

@@ -0,0 +1,29 @@
+import numpy as np
+import tensorflow as tf
+from common.forward_warp.forward_warp import forward_warp
+from common.utils.profile import run_profiler
+
+if __name__ == '__main__':
+    height = 512
+    width = 512
+    im_channels = 3
+    batch_size = 8
+
+    # Create the graph.
+    image_shape = [batch_size, height, width, im_channels]
+    flow_shape = [batch_size, height, width, 2]
+    image_placeholder = tf.placeholder(shape=image_shape, dtype=tf.float32)
+    flow_placeholder = tf.placeholder(shape=flow_shape, dtype=tf.float32)
+    warped = forward_warp(image_placeholder, flow_placeholder, splat_variance=0.5)
+    grads = tf.gradients(warped, [image_placeholder, flow_placeholder])
+
+    # Create dummy images.
+    image = np.zeros(shape=[batch_size, height, width, im_channels], dtype=np.float32)
+    flow = np.zeros(shape=[batch_size, height, width, 2], dtype=np.float32)
+    image[:, 2:height - 2, 2:width - 2, :] = 1.0
+    flow[:, 4:height - 4, 5:width - 5, :] = 1.0
+
+    query = [warped, grads]
+    feed_dict = {image_placeholder: image,
+                 flow_placeholder: flow}
+    run_profiler(query, feed_dict, name='forward-warp')

common/forward_warp/forward_warp_test.py

@@ -3,14 +3,16 @@
 import numpy as np
 import tensorflow as tf
 from common.utils.img import read_image, show_image
-from common.forward_warp.forward_warp import forward_warp, create_disocclusion_mask
+from common.forward_warp.forward_warp import forward_warp, create_disocclusion_mask, is_forward_warp_cuda, forward_warp_tf
 from common.utils.flow import read_flow_file
+from tensorflow.python.ops import gradient_checker
+
 
 VISUALIZE = False
 WRITE_TO_VIDEO = False
 
 
-class TestForwardWarp(unittest.TestCase):
+class TestForwardWarpTF(unittest.TestCase):
     def setUp(self):
         config = tf.ConfigProto()
         config.gpu_options.allow_growth = True
@@ -37,7 +39,7 @@ def test_forward_warp_whole_1(self):
 
         flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
         features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
-        warp_tensor = forward_warp(features_tensor, flow_tensor)
+        warp_tensor = forward_warp_tf(features_tensor, flow_tensor)
         warp = self.sess.run(warp_tensor, feed_dict={flow_tensor: flow, features_tensor: features})
         self.assertEqual(warp.tolist(), expected_warp)
 
@@ -62,7 +64,7 @@ def test_forward_warp_partial_1(self):
 
         flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
         features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
-        warp_tensor = forward_warp(features_tensor, flow_tensor)
+        warp_tensor = forward_warp_tf(features_tensor, flow_tensor)
         warp = self.sess.run(warp_tensor, feed_dict={flow_tensor: flow, features_tensor: features})
         self.assertEqual(warp.tolist(), expected_warp)
 
@@ -90,7 +92,7 @@ def test_forward_warp_partial_2(self):
 
         flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
         features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
-        warp_tensor = forward_warp(features_tensor, flow_tensor)
+        warp_tensor = forward_warp_tf(features_tensor, flow_tensor)
         warp = self.sess.run(warp_tensor, feed_dict={flow_tensor: flow, features_tensor: features})
         self.assertEqual(warp.tolist(), expected_warp)
 
@@ -131,7 +133,7 @@ def test_forward_warp_oob(self):
 
         flow_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
         features_tensor = tf.placeholder(tf.float32, (1, height, width, 2))
-        warp_tensor = forward_warp(features_tensor, flow_tensor)
+        warp_tensor = forward_warp_tf(features_tensor, flow_tensor)
         warp = self.sess.run(warp_tensor, feed_dict={flow_tensor: flow, features_tensor: features})
         self.assertEqual(warp.tolist(), expected_warp)
 
@@ -174,7 +176,7 @@ def test_forward_warp_batch(self):
 
         flow_tensor = tf.placeholder(tf.float32, (2, height, width, 2))
         features_tensor = tf.placeholder(tf.float32, (2, height, width, 2))
-        warp_tensor = forward_warp(features_tensor, flow_tensor)
+        warp_tensor = forward_warp_tf(features_tensor, flow_tensor)
         warp = self.sess.run(warp_tensor, feed_dict={flow_tensor: flow, features_tensor: features})
         self.assertEqual(warp[0].tolist(), expected_warp[0])
         self.assertEqual(warp[1].tolist(), expected_warp[1])
@@ -189,33 +191,47 @@ def test_forward_warp_batch(self):
         self.assertNotEqual(np.sum(flow_grads), 0.0)
         self.assertNotEqual(np.sum(feature_grads), 0.0)
 
+
+class TestForwardWarpCommon(unittest.TestCase):
+    def setUp(self):
+        config = tf.ConfigProto()
+        config.gpu_options.allow_growth = True
+        self.sess = tf.Session(config=config)
+
+        self.flow_path = os.path.join('pwcnet', 'warp', 'test_data', 'flow_ab.flo')
+        self.image_path_a = os.path.join('pwcnet', 'warp', 'test_data', 'image_a.png')
+        self.image_path_b = os.path.join('pwcnet', 'warp', 'test_data', 'image_b.png')
+
+        self.max_allowable_grad_err = 5e-4
+
     def test_visualization(self):
         if not VISUALIZE:
             return
 
-        cur_dir = os.path.dirname(os.path.abspath(__file__))
-        root_dir = os.path.join(cur_dir, '..')
-        flow_path = os.path.join(root_dir, 'pwcnet', 'warp', 'test_data', 'flow_ab.flo')
-        image_path = os.path.join(root_dir, 'pwcnet', 'warp', 'test_data', 'image_a.png')
-
-        flow_ab = [read_flow_file(flow_path)]
-        img_a = [read_image(image_path, as_float=True)]
+        flow_ab = [read_flow_file(self.flow_path)]
+        img_a = [read_image(self.image_path_a, as_float=True)]
         t_tensor = tf.placeholder(tf.float32, None)
         flow_ab_tensor = tf.placeholder(tf.float32, np.shape(flow_ab))
         img_a_tensor = tf.placeholder(tf.float32, np.shape(img_a))
         warp_tensor = forward_warp(img_a_tensor, t_tensor * flow_ab_tensor)
 
         warp = self.sess.run(warp_tensor, feed_dict={flow_ab_tensor: flow_ab, img_a_tensor: img_a, t_tensor: 1.0})
         warp = np.clip(warp[0], 0.0, 1.0)
-        show_image(warp)
+        try:
+            show_image(warp)
+        except:
+            print('show_image(warp) failed.')
 
         # For writing to video.
         if WRITE_TO_VIDEO:
+            if not os.path.exists('outputs'):
+                os.makedirs('outputs')
+
             import cv2
-            import mpimg
+            import matplotlib.image as mpimg
             height = img_a[0].shape[0]
             width = img_a[0].shape[1]
-            writer = cv2.VideoWriter(cur_dir + '/outputs/warped.avi',
+            writer = cv2.VideoWriter('outputs/warped.avi',
                                      cv2.VideoWriter_fourcc(*'MJPG'), 20, (width, height))
             steps = 60
             for i in range(steps):
@@ -225,17 +241,31 @@ def test_visualization(self):
                                        feed_dict={flow_ab_tensor: flow_ab, img_a_tensor: img_a, t_tensor: t})
                 warped = warped[0]
                 warped = np.clip(warped, 0.0, 1.0)
-                output_path = cur_dir + '/outputs/out-%.2f.png' % t
+                output_path = 'outputs/out-%.2f.png' % t
                 mpimg.imsave(output_path, warped)
                 writer.write(cv2.imread(output_path))
             writer.release()
 
+    def test_warp_error(self):
+        flow_ab = [read_flow_file(self.flow_path)]
+        img_a = [read_image(self.image_path_a, as_float=True)]
+        img_b = read_image(self.image_path_b, as_float=True)
+        flow_ab_tensor = tf.placeholder(tf.float32, np.shape(flow_ab))
+        img_a_tensor = tf.placeholder(tf.float32, np.shape(img_a))
+        warp_tensor = forward_warp(img_a_tensor, flow_ab_tensor, splat_variance=0.3)
+        mask = 1.0 - create_disocclusion_mask(flow_ab_tensor)
+
+        warp, mask = self.sess.run([warp_tensor, mask], feed_dict={flow_ab_tensor: flow_ab, img_a_tensor: img_a})
+        warp = np.clip(warp[0], 0.0, 1.0)
+
+        self.assertLess(np.average(np.abs(warp - img_b) * mask[0]), 0.0212)
+
     def test_create_disocclusion_map(self):
         height = 3
         width = 3
 
         flow_tensor = tf.placeholder(shape=(None, height, width, 2), dtype=tf.float32)
-        mask_tensor = create_disocclusion_mask(flow_tensor)
+        mask_tensor = create_disocclusion_mask(flow_tensor, splat_variance=0.2)
 
         flow = np.asarray([
             [
@@ -261,7 +291,7 @@ def test_create_disocclusion_map_batched(self):
         width = 3
 
         flow_tensor = tf.placeholder(shape=(None, height, width, 2), dtype=tf.float32)
-        mask_tensor = create_disocclusion_mask(flow_tensor)
+        mask_tensor = create_disocclusion_mask(flow_tensor, splat_variance=0.2)
 
         flow = np.asarray([
             [
@@ -302,6 +332,64 @@ def test_create_disocclusion_map_no_gradient(self):
         grad = tf.gradients(mask_tensor, flow_tensor)[0]
         self.assertEqual(None, grad)
 
+    def test_gradients_errors(self):
+        self.gradient_errors_helper(splat_variance=1.0)
+
+    def test_gradients_errors_low_splat(self):
+        if not is_forward_warp_cuda():
+            return
+        self.gradient_errors_helper(splat_variance=0.2)
+
+    def test_gradients_errors_high_splat(self):
+        if not is_forward_warp_cuda():
+            return
+        self.gradient_errors_helper(splat_variance=0.4)
+
+    def gradient_errors_helper(self, splat_variance):
+        with self.sess:
+            # This test is flaky, so retry if fail.
+            num_tries = 2 if is_forward_warp_cuda() else 4
+            error1 = 0
+            error2 = 0
+            for i in range(num_tries):
+                img_shape = (16, 3, 4, 4)
+                flow_shape = (16, 3, 4, 2)
+                img_a = np.random.rand(*img_shape)
+                flow_ab = (np.random.rand(*flow_shape) - 0.5) * 3
+                input = tf.placeholder(shape=img_a.shape, dtype=tf.float32)
+                flow_tensor = tf.placeholder(shape=flow_ab.shape, dtype=tf.float32)
+                warped_tensor = forward_warp(input, flow_tensor, splat_variance=splat_variance)
+
+                error1 = gradient_checker.compute_gradient_error(input, img_a.shape, warped_tensor, img_a.shape,
+                                                                 extra_feed_dict={flow_tensor: flow_ab},
+                                                                 x_init_value=img_a)
+                error2 = gradient_checker.compute_gradient_error(flow_tensor, flow_ab.shape, warped_tensor,
+                                                                 img_a.shape, extra_feed_dict={input: img_a},
+                                                                 x_init_value=flow_ab)
+                if error1 <= self.max_allowable_grad_err and error2 <= self.max_allowable_grad_err:
+                    return
+            self.assertLessEqual(max(error1, error2), self.max_allowable_grad_err,
+                                 'Exceeded the error threshold. Note that this test may be flaky.')
+
+    def test_gradient_errors_simultaneous(self):
+        with self.sess:
+            # This test is flaky, so retry if fail.
+            num_tries = 2 if is_forward_warp_cuda() else 4
+            error = 0
+            for i in range(num_tries):
+                img_shape = (16, 3, 4, 4)
+                flow_shape = (16, 3, 4, 2)
+                input = tf.ones(shape=img_shape, dtype=tf.float32)
+                flow_tensor = tf.ones(shape=flow_shape, dtype=tf.float32)
+                warped_tensor = forward_warp(input, flow_tensor, splat_variance=0.2)
+
+                error = gradient_checker.compute_gradient_error([input, flow_tensor], [img_shape, flow_shape],
+                                                                warped_tensor, img_shape)
+                if error <= self.max_allowable_grad_err:
+                    return
+            self.assertLessEqual(error, self.max_allowable_grad_err,
+                                 'Exceeded the error threshold. Note that this test may be flaky.')
+
 
 if __name__ == '__main__':
     unittest.main()

common/forward_warp/native/CMakeLists.txt

@@ -0,0 +1,6 @@
+cmake_minimum_required(VERSION 3.5)
+
+add_op_library(NAME forward_warp_op SOURCES
+    "forward_warp_op.cc"
+    "forward_warp_op.cc.cu"
+)