Warm tip: This article is reproduced from serverfault.com, please click

Temporarily merge the batch dimension in Keras

发布于 2020-11-28 13:43:05

I have a Keras model that has in input shape of [None, 500, 500, 3] and output shape of [None, 1]. Now, I want to make a wrapper model that has an input shape of [None, 48, 500, 500, 3], and output shape of [None, 48].

To do that, the naive way is to iterate 48 times on the second axis and apply the first model, then use Keras' Concatenate layer to get the desired shape.

model_outputs = []
for i in range(inputs.shape[1]):
    im_block = inputs[:, i]
    model_outputs += [self.model(im_block)]
return Concatenate()(model_outputs)

However, this makes the graph quite complicated. So I would like to do the following instead:

        [None, 48, 500, 500, 3]
     -> [None*48,  500, 500, 3]
           (apply the model)
     -> [None*48,  1]
     -> [None, 48, 1]

My attempt at that is

outputs = tf.reshape(inputs, (inputs[0] * inputs[1], *inputs[2:]))
outputs = self.model(outputs)
outputs = tf.reshape(outputs, (inputs[0], inputs[1]))
return outputs

But this gives me

TypeError: Cannot iterate over a tensor with unknown first dimension.

Is there a way to do that ?

Questioner
Samuel Prevost
Viewed
0
Andrey 2020-11-29 01:29:39

This should work:

inp = tf.reshape(inp, (-1, 500, 500, 3))
res = model(inp)
res = tf.reshape(res, (-1, 48))