Dealing with none dimension of symbolic tensor when creating a layer

Question

I want to implement my own Max Unpooling layer as explained in here. For that, I need the argmax output of tf.nn.max_pool_with_argmax. I successfully applied it using the layer outside any model, but when I want to add it to a model (and therefore, TensorFlow uses symbolic/static tensors, I get errors. I have tried many implementations and I came out with 2 options in summary:

Option 1: My favorite. Give the expected output shape to the constructor (without batch size, meaning in a 2D example would be WxHxC.
Option 2: Pass a reference tensor and imitate that shape.

In both options I have the same problem on tf.scatter_nd:

*** ValueError: Tried to convert 'shape' to a tensor and failed. Error: Cannot convert a partially known TensorShape to a Tensor

I understand why this error happens, with dynamic tensors, the batch is unknown and therefore the shape is tf.TensorShape([None, ...]). But how do I deal with it?

Here my implementation.

class UnPooling2D(Layer):

    def __init__(self, desired_output_shape, name=None, dtype=DEFAULT_COMPLEX_TYPE, dynamic=False, **kwargs):
        self.desired_output_shape = desired_output_shape    # If option 1
        super(ComplexUnPooling2D, self).__init__(trainable=False, name=name, dtype=self.my_dtype.real_dtype,
                                               dynamic=dynamic, **kwargs)

    def call(self, inputs, unpool_mat, reference_tensor_for_shape: Optional[Only if option 2], **kwargs):
        """
        Performs unpooling
        :param inputs: Input Tensor.
        :param unpool_mat: Result argmax from tf.nn.max_pool_with_argmax
            https://www.tensorflow.org/api_docs/python/tf/nn/max_pool_with_argmax
        """       
        updates = tf.reshape(inputs, [-1])
        indices = tf.expand_dims(tf.reshape(unpool_mat, [-1]), axis=-1)

        ####################
        # Option 1:
        flat_output_shape = tf.reduce_prod(self.desired_output_shape)
        ret = tf.scatter_nd(indices, updates, shape=(inputs.get_shape()[0]*flat_output_shape,))
        desired_output_shape_with_batch = tf.concat([[inputs.get_shape()[0]], self.desired_output_shape], axis=0)
        ret = tf.reshape(ret, shape=desired_output_shape_with_batch)

        # Option 2 (untested, speudo code):
        flatten_reference_tensor = tf.reshape(reference_tensor_for_shape, [-1])
        ret = tf.scatter_nd(indices, updates, shape=flatten_reference_tensor.get_shape())
        ret = tf.reshape(ret, reference_tensor_for_shape.get_shape())
        #################
        return ret

I tried an ugly fix but have the error:

WARNING:tensorflow:AutoGraph could not transform > and will run it as-is.

What I tried is to do:

    def call(self, inputs, unpool_mat, **kwargs):
        if inputs.get_shape()[0]:
           ... other solutions
        else:   # Dynamic tensors
            ret = tf.reshape(inputs, (-1,) + self.desired_output_shape)
        return ret

user_5931672 · Answer

I partially solved the issue, but I think my new bug merits a new question actually. In the end I had to change inputs.get_shape() for tf.shape(inputs). This is my final code that "works". (at least it gets the shape of the model right without errors in, for example, model.summary().

    def call(self, inputs, unpool_mat, **kwargs):
        flat_output_shape = tf.reduce_prod(self.desired_output_shape)

        updates = tf.reshape(inputs, [-1])
        indices = tf.expand_dims(tf.reshape(unpool_mat, [-1]), axis=-1)

        ret = tf.scatter_nd(indices, updates, shape=(tf.shape(inputs)[0]*flat_output_shape,))
        desired_output_shape_with_batch = tf.concat([[tf.shape(inputs)[0]], self.desired_output_shape], axis=0)
        ret = tf.reshape(ret, shape=desired_output_shape_with_batch)
        return ret

Dealing with none dimension of symbolic tensor when creating a layer

1 Answers