Find the maximum number of consecutive ones in Tensorflow

Question

I have found solutions in

Matlab: https://uk.mathworks.com/matlabcentral/answers/85143-find-maximum-number-of-consecutive-negative-values

Numpy: Find consecutive ones in numpy array

But not in Tensorflow. The closest question I've found is this: Reduce sum with condition in tensorflow

However this only counts the first group of ones, rather than finding the largest group. Basically, the problem would be easily solved if there was a Tensorflow equivalent of Matlab's RCUMSUMC

https://uk.mathworks.com/matlabcentral/fileexchange/28685-rcumsumc

My input is a binary tensor of shape NxHxW, and the output is expected to be NxW, where each column represents the maximum number of consecutive ones:

Input = [[1,1,0,0,1,0,1,1,1,1,0,0,1],
         [1,0,0,1,1,1,1,1,1,0,0,1,0],
         [0,0,0,1,1,1,0,0,1,0,1,1,0]]

Output = [4,6,3]

Can you provide a simple example of the input you have and the expected output? — thushv89, Nov 26 '19 at 10:06

score 1 · Answer 1 · answered Nov 26 '19 at 11:18

Alright, this is a bit convoluted. But it was fun coming up with this. (Tested both on TF 2.0 and TF 1.15)

# Let's assume a simpler example for demonstration [0, 1, 1, 1, 0, 1]
tf_a = tf.constant([[1,1,0,0,1,0,1,1,1,1,0,0,1],
         [1,0,0,1,1,1,1,1,1,0,0,1,0],
         [0,0,0,1,1,1,0,0,1,0,1,1,0]])

# You get the cumsum [0, 1, 2, 3, 3, 4]
tf_a_sum = tf.cumsum(tf_a, axis=1)
# You get the range  [0, 1, 2, 3, 4, 5]
tf_range = tf.range(tf_a.shape[1])

# You get the difference [0, 0, 0, 0, 1, 1]
tf_diff = tf_range - tf_a_sum
# To make sure it's starting with 0
tf_diff = tf_diff - tf.reduce_min(tf_diff, axis=1, keepdims=True)

# Now comes the tricky bit. We are using segment_sum
# I would have liked to achieve this with tf.map_fn but segment_sum didn't play nicely with map_fn
# So we are first unstacking the arrays on rows
# Applying segment_sum to individual rows
# And restacking them

tf_list_diff = tf.unstack(tf_diff)
tf_list_a = tf.unstack(tf_a)

# [0, 1, 1, 1, 0, 1] => summing segment-wise with [0, 0, 0, 0, 1, 1]
# Gives [3,1]
# And you get max of that which is 3
tf_res = tf.stack([tf.reduce_max(tf.math.segment_sum(a, diff)) for a, diff in zip(tf_list_a, tf_list_diff)])

Warning: This will work as long as you're interested in getting the number of 1s in an array containing 0s and 1s only. This won't work if you have other numbers in the array. So this way of solving is very specific to the problem you want to solve.

jdehesa · Accepted Answer · 2019-11-26T17:10:34.863

EDIT: Here is how you could do the same for a 3-dimensional tensor searching for groups of ones across the second dimension:

import tensorflow as tf

inp = tf.random.stateless_uniform((3, 4, 5), [0, 0], 0, 2, tf.int32)
tf.print(inp)
# [[[0 1 0 0 0]
#   [0 0 0 1 1]
#   [0 1 1 0 0]
#   [0 0 0 0 0]]
# 
#  [[1 1 0 0 0]
#   [0 1 1 0 0]
#   [0 1 0 1 0]
#   [0 0 0 0 0]]
# 
#  [[1 1 1 0 0]
#   [0 0 1 1 1]
#   [0 1 1 0 1]
#   [0 1 0 0 0]]]
# Pad with ones to avoid all-zeros
inp = tf.pad(inp, [(0, 1), (0, 0), (0, 0)], constant_values=1)
s = tf.shape(inp)
# Transpose and reshape
inp_t = tf.reshape(tf.transpose(inp, [0, 2, 1]), [-1, s[1]])
# Surround with zeros
inp_t = tf.pad(inp_t, [(0, 0), (1, 1)])
# Find bounds of groups of ones
groups = tf.reshape(tf.where(tf.not_equal(inp_t[:, 1:], inp_t[:, :-1])), [-1, 2, 2])
# Compute group sizes
group_sizes = groups[:, 1, 1] - groups[:, 0, 1]
# Find maximum group sizes
max_ones_group = tf.math.segment_max(group_sizes, groups[:, 0, 0], name=None)
# Reshape back
out = tf.reshape(max_ones_group, [s[0], s[2]])[:-1]
tf.print(out)
# [[0 1 1 1 1]
#  [1 3 1 1 0]
#  [1 2 3 1 2]]

Assuming that the input is a binary tensor (only zeros and ones), this is one way to do that:

import tensorflow as tf

inp = tf.constant([[1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 1],
                   [1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0],
                   [0, 0, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 0]])
# Surround with zeros
inp = tf.pad(inp, [(0, 0), (1, 1)])
# Find bounds of groups of ones
groups = tf.reshape(tf.where(tf.not_equal(inp[:, 1:], inp[:, :-1])), [-1, 2, 2])
# Compute group sizes
group_sizes = groups[:, 1, 1] - groups[:, 0, 1]
# Find maximum group sizes
max_ones_group = tf.math.segment_max(group_sizes, groups[:, 0, 0], name=None)
tf.print(max_ones_group)
# [4 6 3]

It works for N=1, but what if we have many samples? For input of shape NxHxW, the output needs to be NxW. — metalruka, Nov 26 '19 at 11:54
Just a side note, this fails if there are no ones in the input array. Padding the input with a single row of ones is an easy workaround. — metalruka, Nov 26 '19 at 16:56
@metalruka You're right, thanks for the feedback. Yes, I think padding is the easiest solution, I updated the answer. — jdehesa, Nov 26 '19 at 17:10

Find the maximum number of consecutive ones in Tensorflow

2 Answers2