Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

histogram: implement buckets_v3 #5356

Merged
merged 12 commits into from Oct 7, 2021
Merged

histogram: implement buckets_v3 #5356

Show file tree
Hide file tree
Changes from 10 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
05d8c78
rename functions to avoid confusion
yatbear Oct 1, 2021
3ea37da
push the single_value test cases down to all test cases
yatbear Oct 1, 2021
f0815dc
add v3 implementation for single value input
yatbear Oct 1, 2021
6774bea
add test got v3 graph ops
yatbear Oct 1, 2021
fffb21b
Revert "push the single_value test cases down to all test cases"
yatbear Oct 1, 2021
5bf44ad
check if bucket_count <= 0 and fix tf ops
yatbear Oct 1, 2021
3f4cfd3
Make SummaryV3OpGraphTest inherit from V2 test case and add test for …
yatbear Oct 1, 2021
9e18664
use an alternative (tf.fill) to be consistent
yatbear Oct 4, 2021
3e72bde
distinguish zero bucket count case v.s. empty input data case
yatbear Oct 5, 2021
9d32842
add SummaryV3OpGraphTest test case and update tests
yatbear Oct 5, 2021
72cb90b
make bucket_count a variable before tf.cond op
yatbear Oct 6, 2021
2d5fafe
remove debug lines added by accident
yatbear Oct 7, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
3 changes: 3 additions & 0 deletions tensorboard/plugins/histogram/summary.py
View file
Open in desktop
Expand Up @@ -39,6 +39,9 @@
histogram = summary_v2.histogram
histogram_pb = summary_v2.histogram_pb

# Export V3 versions.
histogram_v3 = summary_v2.histogram_v3


def _buckets(data, bucket_count=None):
"""Create a TensorFlow op to group data into histogram buckets.
Expand Down
98 changes: 95 additions & 3 deletions tensorboard/plugins/histogram/summary_test.py
View file
Open in desktop
Expand Up @@ -201,9 +201,12 @@ def write_histogram_event(self, *args, **kwargs):
kwargs.setdefault("step", 1)
writer = tf2.summary.create_file_writer(self.get_temp_dir())
with writer.as_default():
summary.histogram(*args, **kwargs)
self.call_histogram_op(*args, **kwargs)
writer.close()

def call_histogram_op(self, *args, **kwargs):
summary.histogram(*args, **kwargs)

def test_scoped_tag(self):
with tf.name_scope("scope"):
self.assertEqual("scope/a", self.histogram("a", []).value[0].tag)
Expand Down Expand Up @@ -238,7 +241,86 @@ def write_histogram_event(self, *args, **kwargs):
def graph_fn():
# Recreate the active scope inside the defun since it won't propagate.
with tf.name_scope(scope):
summary.histogram(*args, **kwargs)
self.call_histogram_op(*args, **kwargs)

writer = tf2.summary.create_file_writer(self.get_temp_dir())
with writer.as_default():
graph_fn()
writer.close()

def test_no_gradient_error_xla(self):
@tf2.function(jit_compile=True)
def graph_fn():
x = tf.constant(1.0)
with tf2.GradientTape() as tape1:
with tf2.GradientTape() as tape2:
tape1.watch(x)
tape2.watch(x)
self.call_histogram_op(
name="loss", step=0, data=x, buckets=10
)

# Note that XLA CPU/GPU has no outside compilation support, so summaries
# won't actually run in a jit_compiled function. TPUs do, and follow
# some similar codepaths, so this test stops at graph building to
# exercise those paths without a TPU available.
writer = tf2.summary.create_file_writer(self.get_temp_dir())
with writer.as_default():
graph_fn.get_concrete_function()


class SummaryV3OpTest(SummaryV2OpTest, tf.test.TestCase):
def call_histogram_op(self, *args, **kwargs):
yatbear marked this conversation as resolved.
Show resolved Hide resolved
summary.histogram_v3(*args, **kwargs)

def test_singleton_input(self):
pb = self.histogram("twelve", [12])
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
# By default there will be 30 buckets.
expected_buckets = np.array(
[[12, 12, 0] for _ in range(29)] + [[12, 12, 1]]
)
np.testing.assert_allclose(buckets, expected_buckets)

def test_input_with_all_same_values(self):
pb = self.histogram("twelven", [12, 12, 12])
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
# By default there will be 30 buckets.
expected_buckets = np.array(
[[12, 12, 0] for _ in range(29)] + [[12, 12, 3]]
)
np.testing.assert_allclose(buckets, expected_buckets)

def test_empty_input(self):
pb = self.histogram("empty", [])
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
# By default there will be 30 buckets.
np.testing.assert_allclose(buckets, np.zeros((30, 3)))

def test_empty_input_of_high_rank(self):
pb = self.histogram("empty_but_fancy", [[[], []], [[], []]])
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
# By default there will be 30 buckets.
np.testing.assert_allclose(buckets, np.zeros((30, 3)))

def test_zero_bucket_count(self):
pb = self.histogram("zero_bucket_count", [1, 1, 1], buckets=0)
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
np.testing.assert_array_equal(buckets, np.array([]).reshape((0, 3)))


class SummaryV3OpGraphTest(SummaryV3OpTest, tf.test.TestCase):
def write_histogram_event(self, *args, **kwargs):
kwargs.setdefault("step", 1)
# Hack to extract current scope since there's no direct API for it.
with tf.name_scope("_") as temp_scope:
scope = temp_scope.rstrip("/_")

@tf2.function
def graph_fn():
# Recreate the active scope inside the defun since it won't propagate.
with tf.name_scope(scope):
self.call_histogram_op(*args, **kwargs)

writer = tf2.summary.create_file_writer(self.get_temp_dir())
with writer.as_default():
Expand All @@ -253,7 +335,9 @@ def graph_fn():
with tf2.GradientTape() as tape2:
tape1.watch(x)
tape2.watch(x)
summary.histogram(name="loss", step=0, data=x, buckets=10)
self.call_histogram_op(
name="loss", step=0, data=x, buckets=10
)

# Note that XLA CPU/GPU has no outside compilation support, so summaries
# won't actually run in a jit_compiled function. TPUs do, and follow
Expand All @@ -263,6 +347,14 @@ def graph_fn():
with writer.as_default():
graph_fn.get_concrete_function()

def test_zero_bucket_count(self):
self.skipTest(
"TODO: figure out why this doesn't work in graph test case"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I think the reason this was failing is that in graph mode, when it compiles the TensorFlow code into a graph, all of the branches of the tf.cond()s are inserted as parts of the graph (even if those branches won't be taken when the graph is actually executed). When doing that, TensorFlow does shape inference, described a little bit here: https://www.tensorflow.org/guide/create_op#shape_functions_in_c

Basically, it tries to compute what the resulting shape will be for all of the ops in all of the branches. At the line where we currently do tf.fill([bucket_count - 1], 0), because bucket_count is a constant, it is able to determine at compile time that this will result in tf.fill([-1], 0) and the shape function for tf.fill emits an error, because this would result in an output shape that doesn't make sense.

Hence the error (at least, the one that I was seeing when trying this PR locally):

ValueError: Fill dimensions must be >= 0 for '{{node zero_bucket_count/write_summary/buckets/cond/cond/Fill_1}} = Fill[T=DT_INT32, index_type=DT_INT32](zero_bucket_count/write_summary/buckets/cond/cond/Fill_1/dims, zero_bucket_count/write_summary/buckets/cond/cond/Fill_1/value)' with input shapes: [1], [] and with input tensors computed as partial shapes: input[0] = [?].

In other words, it doesn't actually mean it's picking the wrong branch of the conditional to execute; it's that during graph compilation, it visits all the branches regardless of what will be ultimately executed.

(Aside: you could argue that when bucket_count is a constant 0, we could actually optimize away all the tf.cond ops and all the branches besides the empty case, at which point we wouldn't hit this error. I'm guessing that the compilation's constant-folding isn't quite smart enough to do that. Also, if you change the test to pass bucket_count=tf.Variable(0) and replace the use of max() with tf.math.maximum() as mentioned in the other comment, it will also pass, since now that the bucket count is a variable, it doesn't do as much shape inference and can't raise the false positive error.)

Ultimately, the best way to fix this is to just ensure that all the branches are still creating the correctly shaped tensor of dimensions (bucket_count, 3), even when we know that branch will never be executed. Here's a patch that gets the test to pass:

diff --git i/tensorboard/plugins/histogram/summary_v2.py w/tensorboard/plugins/histogram/summary_v2.py
index bff980492..e1cd01846 100644
--- i/tensorboard/plugins/histogram/summary_v2.py
+++ w/tensorboard/plugins/histogram/summary_v2.py
@@ -425,6 +425,8 @@ def _buckets_v3(data, bucket_count=None):
     with tf.name_scope("buckets"):
         tf.debugging.assert_scalar(bucket_count)
         tf.debugging.assert_type(bucket_count, tf.int32)
+        # Treat a negative bucket count as zero.
+        bucket_count = tf.math.maximum(0, bucket_count)
         data = tf.reshape(data, shape=[-1])  # flatten
         data = tf.cast(data, tf.float64)
         data_size = tf.size(input=data)
@@ -440,7 +442,7 @@ def _buckets_v3(data, bucket_count=None):
             2. If the input data is empty, a tensor of shape (bucket_count, 3)
               of all zero values will be returned.
             """
-            return tf.zeros((max(0, bucket_count), 3), dtype=tf.float64)
+            return tf.zeros((bucket_count, 3), dtype=tf.float64)

         def when_nonempty():
             min_ = tf.reduce_min(input_tensor=data)
@@ -480,11 +482,12 @@ def _buckets_v3(data, bucket_count=None):
                 """When input data contains a single unique value."""
                 # Left and right edges are the same for single value input.
                 edges = tf.fill([bucket_count], max_)
-                # Counts for the first {bucket_count - 1} buckets [v, v) are 0.
-                zero_bucket_counts = tf.fill([bucket_count - 1], 0)
-                # Count for last bucket [v, v] is {data_size}.
+                # Bucket counts are 0 except the last bucket (if bucket_count > 0),
+                # which is `data_size`. Ensure that the resulting counts vector has
+                # length `bucket_count` always, including the bucket_count==0 case.
+                zeroes = tf.fill([bucket_count], 0)
                 bucket_counts = tf.cast(
-                    tf.concat([zero_bucket_counts, [data_size]], 0),
+                    tf.concat([zeroes[:-1], [data_size]], 0)[:bucket_count],
                     dtype=tf.float64,
                 )
                 return tf.transpose(a=tf.stack([edges, edges, bucket_counts]))

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks so much for the clear explanation and the patch!

)
pb = self.histogram("zero_bucket_count", [1, 1, 1], buckets=0)
buckets = tensor_util.make_ndarray(pb.value[0].tensor)
np.testing.assert_array_equal(buckets, np.array([]).reshape((0, 3)))


if __name__ == "__main__":
tf.test.main()
47 changes: 31 additions & 16 deletions tensorboard/plugins/histogram/summary_v2.py
View file
Open in desktop
Expand Up @@ -412,7 +412,8 @@ def _buckets_v3(data, bucket_count=None):

Arguments:
data: A `Tensor` of any shape. Must be castable to `float64`.
bucket_count: Optional positive `int` or scalar `int32` `Tensor`.
bucket_count: Optional non-negative `int` or scalar `int32` `Tensor`,
defaults to 30.
Returns:
A `Tensor` of shape `[k, 3]` and type `float64`. The `i`th row is
a triple `[left_edge, right_edge, count]` for a single bucket.
Expand All @@ -426,19 +427,29 @@ def _buckets_v3(data, bucket_count=None):
tf.debugging.assert_type(bucket_count, tf.int32)
data = tf.reshape(data, shape=[-1]) # flatten
data = tf.cast(data, tf.float64)
is_empty = tf.equal(tf.size(input=data), 0)
data_size = tf.size(input=data)
is_empty = tf.logical_or(
tf.equal(data_size, 0), tf.less_equal(bucket_count, 0)
)

def when_empty():
return tf.constant([], shape=(0, 3), dtype=tf.float64)
"""When input data is empty or bucket_count is zero.

1. If bucket_count is specified as zero, an empty tensor of shape
(0, 3) will be returned.
2. If the input data is empty, a tensor of shape (bucket_count, 3)
of all zero values will be returned.
"""
return tf.zeros((max(0, bucket_count), 3), dtype=tf.float64)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should use tf.math.maximum() here, rather than Python's max(), in order to be compatible with bucket_count values that aren't native Python numbers (e.g. a tf.constant() instead).

Also I'd recommend just adding a conversion up at the top that does bucket_count = tf.math.maximum(0, bucket_count) so we don't have to worry about negative bucket counts in any of the later logic.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. Thanks!


# TODO(ytjing): Make the nonempty case handling TPU compatible.
def when_nonempty():
min_ = tf.reduce_min(input_tensor=data)
max_ = tf.reduce_max(input_tensor=data)
range_ = max_ - min_
is_singular = tf.equal(range_, 0)
has_single_value = tf.equal(range_, 0)

def when_nonsingular():
def when_multiple_values():
"""When input data contains multiple values."""
bucket_width = range_ / tf.cast(bucket_count, tf.float64)
offsets = data - min_
bucket_indices = tf.cast(
Expand All @@ -465,17 +476,21 @@ def when_nonsingular():
a=tf.stack([left_edges, right_edges, bucket_counts])
)

def when_singular():
center = min_
bucket_starts = tf.stack([center - 0.5])
bucket_ends = tf.stack([center + 0.5])
bucket_counts = tf.stack(
[tf.cast(tf.size(input=data), tf.float64)]
)
return tf.transpose(
a=tf.stack([bucket_starts, bucket_ends, bucket_counts])
def when_single_value():
"""When input data contains a single unique value."""
# Left and right edges are the same for single value input.
edges = tf.fill([bucket_count], max_)
# Counts for the first {bucket_count - 1} buckets [v, v) are 0.
zero_bucket_counts = tf.fill([bucket_count - 1], 0)
# Count for last bucket [v, v] is {data_size}.
bucket_counts = tf.cast(
tf.concat([zero_bucket_counts, [data_size]], 0),
dtype=tf.float64,
)
return tf.transpose(a=tf.stack([edges, edges, bucket_counts]))

return tf.cond(is_singular, when_singular, when_nonsingular)
return tf.cond(
has_single_value, when_single_value, when_multiple_values
)

return tf.cond(is_empty, when_empty, when_nonempty)