Merge branch 'master' into BLD/add_cuda_py_test

Author: facaiy

examples/optimizers_lazyadam.ipynb

@@ -90,7 +90,7 @@
       "source": [
         "# LazyAdam\n",
         "\n",
-        "> LazyAdam is a variant of the Adam optimizer that handles sparse updates moreefficiently.\n",
+        "> LazyAdam is a variant of the Adam optimizer that handles sparse updates more efficiently.\n",
         "    The original Adam algorithm maintains two moving-average accumulators for\n",
         "    each trainable variable; the accumulators are updated at every step.\n",
         "    This class provides lazier handling of gradient updates for sparse\n",

tensorflow_addons/image/distance_transform.py

@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import tensorflow as tf
+from tensorflow_addons.image import utils as img_utils
 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
 
 _image_ops_so = tf.load_op_library(
@@ -32,8 +33,7 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None):
 
     Args:
       images: A tensor of shape (num_images, num_rows, num_columns, 1) (NHWC),
-        or (num_rows, num_columns, 1) (HWC). The rank must be statically known
-        (the shape is not `TensorShape(None)`.
+        or (num_rows, num_columns, 1) (HWC) or (num_rows, num_columns) (HW).
       dtype: DType of the output tensor.
       name: The name of the op.
 
@@ -45,7 +45,7 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None):
     Raises:
       TypeError: If `image` is not tf.uint8, or `dtype` is not floating point.
       ValueError: If `image` more than one channel, or `image` is not of
-        rank 3 or 4.
+        rank between 2 and 4.
     """
 
     with tf.name_scope(name or "euclidean_distance_transform"):
@@ -54,14 +54,9 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None):
         if image_or_images.dtype.base_dtype != tf.uint8:
             raise TypeError(
                 "Invalid dtype %s. Expected uint8." % image_or_images.dtype)
-        if image_or_images.get_shape().ndims is None:
-            raise ValueError("`images` rank must be statically known")
-        elif len(image_or_images.get_shape()) == 3:
-            images = image_or_images[None, :, :, :]
-        elif len(image_or_images.get_shape()) == 4:
-            images = image_or_images
-        else:
-            raise ValueError("`images` should have rank between 3 and 4")
+
+        images = img_utils.to_4D_image(image_or_images)
+        original_ndims = img_utils.get_ndims(image_or_images)
 
         if images.get_shape()[3] != 1 and images.get_shape()[3] is not None:
             raise ValueError("`images` must have only one channel")
@@ -72,6 +67,4 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None):
         images = tf.cast(images, dtype)
         output = _image_ops_so.euclidean_distance_transform(images)
 
-        if len(image_or_images.get_shape()) == 3:
-            return output[0, :, :, :]
-        return output
+        return img_utils.from_4D_image(output, original_ndims)

tensorflow_addons/image/distance_transform_test.py

@@ -21,7 +21,7 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow_addons.image import distance_transform as distance_tranform_ops
+from tensorflow_addons.image import distance_transform as dist_ops
 from tensorflow_addons.utils import test_utils
 
 
@@ -45,7 +45,7 @@ def test_single_binary_image(self):
         image = tf.constant(image, dtype=tf.uint8)
 
         for output_dtype in [tf.float16, tf.float32, tf.float64]:
-            output = distance_tranform_ops.euclidean_dist_transform(
+            output = dist_ops.euclidean_dist_transform(
                 image, dtype=output_dtype)
             output_flat = tf.reshape(output, [-1])
 
@@ -73,7 +73,7 @@ def test_batch_binary_images(self):
         # yapf: enable
         images = tf.constant([image] * batch_size, dtype=tf.uint8)
         for output_dtype in [tf.float16, tf.float32, tf.float64]:
-            output = distance_tranform_ops.euclidean_dist_transform(
+            output = dist_ops.euclidean_dist_transform(
                 images, dtype=output_dtype)
             output_flat = tf.reshape(output, [-1])
 
@@ -97,38 +97,37 @@ def test_image_with_invalid_dtype(self):
             # pylint: disable=bad-continuation
             with self.assertRaisesRegex(
                     TypeError, "`dtype` must be float16, float32 or float64"):
-                _ = distance_tranform_ops.euclidean_dist_transform(
+                _ = dist_ops.euclidean_dist_transform(
                     image, dtype=output_dtype)
 
     def test_image_with_invalid_shape(self):
-        for invalid_shape in ([1], [2, 1], [2, 4, 4, 4, 1]):
-            image = tf.zeros(invalid_shape, tf.uint8)
-
-            # pylint: disable=bad-continuation
-            with self.assertRaisesRegex(
-                    ValueError, "`images` should have rank between 3 and 4"):
-                _ = distance_tranform_ops.euclidean_dist_transform(image)
-
         image = tf.zeros([2, 4, 3], tf.uint8)
         with self.assertRaisesRegex(ValueError,
                                     "`images` must have only one channel"):
-            _ = distance_tranform_ops.euclidean_dist_transform(image)
+            _ = dist_ops.euclidean_dist_transform(image)
 
     def test_all_zeros(self):
-        image = tf.zeros([10, 10, 1], tf.uint8)
-        expected_output = np.zeros([10, 10, 1])
+        image = tf.zeros([10, 10], tf.uint8)
+        expected_output = np.zeros([10, 10])
 
         for output_dtype in [tf.float16, tf.float32, tf.float64]:
-            output = distance_tranform_ops.euclidean_dist_transform(
+            output = dist_ops.euclidean_dist_transform(
                 image, dtype=output_dtype)
             self.assertAllClose(output, expected_output)
 
     def test_all_ones(self):
         image = tf.ones([10, 10, 1], tf.uint8)
-        output = distance_tranform_ops.euclidean_dist_transform(image)
+        output = dist_ops.euclidean_dist_transform(image)
         expected_output = np.full([10, 10, 1], tf.float32.max)
         self.assertAllClose(output, expected_output)
 
+    def test_unknown_shape(self):
+        fn = dist_ops.euclidean_dist_transform.get_concrete_function(
+            tf.TensorSpec(None, tf.uint8))
+        for shape in [[5, 10], [10, 7, 1], [4, 10, 10, 1]]:
+            image = tf.zeros(shape, dtype=tf.uint8)
+            self.assertAllClose(image, fn(image))
+
 
 if __name__ == "__main__":
     tf.test.main()

tensorflow_addons/image/filters.py

@@ -18,6 +18,7 @@
 from __future__ import print_function
 
 import tensorflow as tf
+from tensorflow_addons.image import utils as img_utils
 from tensorflow_addons.utils import keras_utils
 
 
@@ -59,7 +60,8 @@ def mean_filter2d(image,
     """Perform mean filtering on image(s).
 
     Args:
-      image: Either a 3-D `Tensor` of shape `[height, width, channels]`,
+      image: Either a 2-D `Tensor` of shape `[height, width]`,
+        a 3-D `Tensor` of shape `[height, width, channels]`,
         or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
       filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
         the height and width of the 2-D mean filter. Can be a single integer
@@ -74,16 +76,14 @@ def mean_filter2d(image,
     Returns:
       3-D or 4-D `Tensor` of the same dtype as input.
     Raises:
-      ValueError: If `image` is not 3 or 4-dimensional,
+      ValueError: If `image` is not 2, 3 or 4-dimensional,
         if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
         or if `filter_shape` is invalid.
     """
     with tf.name_scope(name or "mean_filter2d"):
         image = tf.convert_to_tensor(image, name="image")
-
-        rank = image.shape.rank
-        if rank != 3 and rank != 4:
-            raise ValueError("image should be either 3 or 4-dimensional.")
+        original_ndims = img_utils.get_ndims(image)
+        image = img_utils.to_4D_image(image)
 
         if padding not in ["REFLECT", "CONSTANT", "SYMMETRIC"]:
             raise ValueError(
@@ -93,10 +93,6 @@ def mean_filter2d(image,
         filter_shape = keras_utils.normalize_tuple(filter_shape, 2,
                                                    "filter_shape")
 
-        # Expand to a 4-D tensor
-        if rank == 3:
-            image = tf.expand_dims(image, axis=0)
-
         # Keep the precision if it's float;
         # otherwise, convert to float32 for computing.
         orig_dtype = image.dtype
@@ -119,11 +115,7 @@ def mean_filter2d(image,
 
         output /= area
 
-        # Squeeze out the first axis to make sure
-        # output has the same dimension with image.
-        if rank == 3:
-            output = tf.squeeze(output, axis=0)
-
+        output = img_utils.from_4D_image(output, original_ndims)
         return tf.dtypes.cast(output, orig_dtype)
 
 
@@ -136,7 +128,8 @@ def median_filter2d(image,
     """Perform median filtering on image(s).
 
     Args:
-      image: Either a 3-D `Tensor` of shape `[height, width, channels]`,
+      image: Either a 2-D `Tensor` of shape `[height, width]`,
+        a 3-D `Tensor` of shape `[height, width, channels]`,
         or a 4-D `Tensor` of shape `[batch_size, height, width, channels]`.
       filter_shape: An `integer` or `tuple`/`list` of 2 integers, specifying
         the height and width of the 2-D median filter. Can be a single integer
@@ -151,16 +144,14 @@ def median_filter2d(image,
     Returns:
       3-D or 4-D `Tensor` of the same dtype as input.
     Raises:
-      ValueError: If `image` is not 3 or 4-dimensional,
+      ValueError: If `image` is not 2, 3 or 4-dimensional,
         if `padding` is other than "REFLECT", "CONSTANT" or "SYMMETRIC",
         or if `filter_shape` is invalid.
     """
     with tf.name_scope(name or "median_filter2d"):
         image = tf.convert_to_tensor(image, name="image")
-
-        rank = image.shape.rank
-        if rank != 3 and rank != 4:
-            raise ValueError("image should be either 3 or 4-dimensional.")
+        original_ndims = img_utils.get_ndims(image)
+        image = img_utils.to_4D_image(image)
 
         if padding not in ["REFLECT", "CONSTANT", "SYMMETRIC"]:
             raise ValueError(
@@ -170,10 +161,6 @@ def median_filter2d(image,
         filter_shape = keras_utils.normalize_tuple(filter_shape, 2,
                                                    "filter_shape")
 
-        # Expand to a 4-D tensor
-        if rank == 3:
-            image = tf.expand_dims(image, axis=0)
-
         image_shape = tf.shape(image)
         batch_size = image_shape[0]
         height = image_shape[1]
@@ -212,10 +199,5 @@ def median_filter2d(image,
                 top[:, :, :, :, floor - 1] + top[:, :, :, :, ceil - 1]) / 2
 
         output = tf.cast(median, image.dtype)
-
-        # Squeeze out the first axis to make sure
-        # output has the same dimension with image.
-        if rank == 3:
-            output = tf.squeeze(output, axis=0)
-
+        output = img_utils.from_4D_image(output, original_ndims)
         return output

tensorflow_addons/image/filters_test.py

@@ -82,12 +82,13 @@ def setUp(self):
         super(MeanFilter2dTest, self).setUp()
 
     def test_invalid_image(self):
-        msg = "image should be either 3 or 4-dimensional."
-
-        for image_shape in [(28, 28), (16, 28, 28, 1, 1)]:
+        msg = "`image` must be 2/3/4D tensor"
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        for image_shape in [(1,), (16, 28, 28, 1, 1)]:
             with self.subTest(dim=len(image_shape)):
-                with self.assertRaisesRegexp(ValueError, msg):
-                    mean_filter2d(tf.ones(shape=image_shape))
+                with self.assertRaisesRegexp(errors, msg):
+                    image = tf.ones(shape=image_shape)
+                    self.evaluate(mean_filter2d(image))
 
     def test_invalid_filter_shape(self):
         msg = ("The `filter_shape` argument must be a tuple of 2 integers.")
@@ -119,6 +120,16 @@ def test_none_channels(self):
         fn(tf.ones(shape=(1, 3, 3, 1)))
         fn(tf.ones(shape=(1, 3, 3, 3)))
 
+    def test_unknown_shape(self):
+        fn = mean_filter2d.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.dtypes.float32),
+            padding="CONSTANT",
+            constant_values=1.)
+
+        for shape in [(3, 3), (3, 3, 3), (1, 3, 3, 3)]:
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
     def test_reflect_padding_with_3x3_filter(self):
         expected_plane = tf.constant([[33. / 9., 36. / 9., 39. / 9.],
                                       [42. / 9., 45. / 9., 48. / 9.],
@@ -191,12 +202,13 @@ def setUp(self):
         super(MedianFilter2dTest, self).setUp()
 
     def test_invalid_image(self):
-        msg = "image should be either 3 or 4-dimensional."
-
-        for image_shape in [(28, 28), (16, 28, 28, 1, 1)]:
+        msg = "`image` must be 2/3/4D tensor"
+        errors = (ValueError, tf.errors.InvalidArgumentError)
+        for image_shape in [(1,), (16, 28, 28, 1, 1)]:
             with self.subTest(dim=len(image_shape)):
-                with self.assertRaisesRegexp(ValueError, msg):
-                    median_filter2d(tf.ones(shape=image_shape))
+                with self.assertRaisesRegexp(errors, msg):
+                    image = tf.ones(shape=image_shape)
+                    self.evaluate(median_filter2d(image))
 
     def test_invalid_filter_shape(self):
         msg = ("The `filter_shape` argument must be a tuple of 2 integers.")
@@ -228,6 +240,16 @@ def test_none_channels(self):
         fn(tf.ones(shape=(1, 3, 3, 1)))
         fn(tf.ones(shape=(1, 3, 3, 3)))
 
+    def test_unknown_shape(self):
+        fn = median_filter2d.get_concrete_function(
+            tf.TensorSpec(shape=None, dtype=tf.dtypes.float32),
+            padding="CONSTANT",
+            constant_values=1.)
+
+        for shape in [(3, 3), (3, 3, 3), (1, 3, 3, 3)]:
+            image = tf.ones(shape=shape)
+            self.assertAllEqual(self.evaluate(image), self.evaluate(fn(image)))
+
     def test_reflect_padding_with_3x3_filter(self):
         expected_plane = tf.constant([[4, 4, 5], [5, 5, 5], [5, 6, 6]])
 

tools/ci_testing/addons_gpu.sh

@@ -25,7 +25,7 @@ fi
 
 set -x
 
-N_JOBS=$(grep -c ^processor /proc/cpuinfo)
+N_JOBS=1 # Must limit GPU testing to single job to prevent OOM error.
 
 echo ""
 echo "Bazel will use ${N_JOBS} concurrent job(s)."