0

I'm trying to write a model that extracts 10 regions of interest out of 128 proposals and feeds them into a Dense layer:

# x is an input tensor of size [None, 128, 4].
# scores is the corresponding [None, 128] score vector.

indices = tf.image.non_max_suppression(x, scores, 10) 
x = x[indices]
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(20)(x)

When I wrap this in a keras Model and try to compile it, the graph building fails because tf.image.non_max_suppression returns an index tensor of shape (None,) rather than (10,). This makes the (non-batch) size of the flattened x[indices] unknown and hence the Dense layer borks on compiling with the error:

ValueError: The last dimension of the inputs to a Dense layer should be defined. Found None. Full input shape received: (None, None)

My understanding is that NMS returns n <= 10 items rather than 10 items, which is why it doesn't specify a fixed output shape to the graph. The input number of RoI is chosen to be high so that the chance of getting less than 10 RoI out of NMS is basically zero. How can I tell my model that it should always expect NMS to return exactly 10 items so that the input size to the Dense layer is fixed to 10*4?

user28400
  • 217
  • 2
  • 8

1 Answers1

0

I managed to fix this problem by padding the indices to a fixed length like so:

fixed_size_indices = tf.zeros(10, tf.int32)
indices = tf.image.non_max_suppression(x, scores, 10) 
if tf.less_equal(tf.size(indices), 10):
   indices = tf.concat([indices, tf.zeros(10 - tf.size(indices), dtype="int32")], 0)
fixed_size_indices += indices
x = x[fixed_size_indices]

It wouldn't work if I just did the padding on indices without adding it to a fixed size tensor.

Not sure if this is the best or most elegant solution but it works for now. Ideally if NMS returns less than 10 regions it would pad with extra regions in order of ascending badness, but that's outside the scope of this question.

user28400
  • 217
  • 2
  • 8