Using batch size as variable in model creation

NicFit_88 0
Viewed 66

I face the following issue about tf.reshape.

Code:

class ReshapeLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(ReshapeLayer, self).__init__(**kwargs)

  def call(self, x):
    assert len(x.shape) == 3
    b, n, d = x.shape
    x = tf.transpose(x, [1, 0, 2])
    x = tf.reshape(x, (1, n, b*d))
    return x

x = tf.random.normal((20, 10, 5))
ReshapeLayer()(x).shape

TensorShape([1, 10, 100])

However:

inputs = tf.keras.layers.Input(shape=(10, 5))
output = ReshapeLayer()(inputs)
model = tf.keras.Model(input=inputs, output=output)

TypeError: Exception encountered when calling layer "reshape_layer_3" (type ReshapeLayer).

in user code:

    File "<ipython-input-5-13bfec4ca4b4>", line 9, in call  *
        x = tf.reshape(x, (1, n, b*d))

    TypeError: unsupported operand type(s) for *: 'NoneType' and 'int'


Call arguments received by layer "reshape_layer_3" (type ReshapeLayer):
  • x=tf.Tensor(shape=(None, 10, 5), dtype=float32)

I am aware that batch size during model creation is None. I tried tf.shape(x) to treat the dimensions as variable, but it didn't work. How can I get around that exception?

-------------- Addition -------------------

Finally, I want to be able to do something like this after changing the layer definition as suggested.

class ReshapeLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(ReshapeLayer, self).__init__(**kwargs)

  def call(self, x):
    assert len(x.shape) == 3
    n = tf.shape(x)[1]  # x.shape[1] should work as well
    x = tf.transpose(x, [1, 0, 2])
    x = tf.reshape(x, (1, n, -1))
    return x

x = tf.random.normal((20, 10, 5))
ReshapeLayer()(x).shape

Now, I want to use the result of the reshape layer as input to another, possibly custom, layer. Here in the example I took a Dense layer.

inputs = tf.keras.layers.Input(shape=(10, 5))
x = ReshapeLayer()(inputs)
x = tf.keras.layers.Dense(5)(x)
model = tf.keras.Model(inputs=inputs, outputs=x)

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

How can I handle this error?

Cheers

keras tensorflow python
1 Answers

This works

class ReshapeLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(ReshapeLayer, self).__init__(**kwargs)

  def call(self, x):
    assert len(x.shape) == 3
    shape = tf.shape(x)
    b, n, d = shape[0], shape[1], shape[2]
    x = tf.transpose(x, [1, 0, 2])
    x = tf.reshape(x, (1, n, b*d))
    return x

x = tf.random.normal((20, 10, 5))
ReshapeLayer()(x).shape

You just have to access the tf.shape result a bit differently as it returns a tensor.

Note that model = tf.keras.Model(input=inputs, output=output) crashes, you need to use model = tf.keras.Model(inputs=inputs, outputs=output) (the kwargs are plural).

Finally, you can also do this

class ReshapeLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(ReshapeLayer, self).__init__(**kwargs)

  def call(self, x):
    assert len(x.shape) == 3
    n = tf.shape(x)[1]  # x.shape[1] should work as well
    x = tf.transpose(x, [1, 0, 2])
    x = tf.reshape(x, (1, n, -1))
    return x

x = tf.random.normal((20, 10, 5))
ReshapeLayer()(x).shape

You can use -1 in reshape calls to basically tell Tensorflow "figure it out yourself".

xdurch00
Related