Batch-size dimension not honored when composing models

Question

I have a Keras model that I defined during training as:

   img = keras.Input(shape=[65, 65, 2])
   bnorm = keras.layers.BatchNormalization()(img)
   ...
   model = keras.Model(img, outputprob)

During serving, though, my inputs come differently. Therefore, I defined an input layer (verifying that the to_img shape comes out to be also (65, 65, 2)) and tried to do model composition using:

  to_img = keras.layers.Lambda(...)(json_input)
  model_output = model(to_img)
  serving_model = keras.Model(json_input, model_output)

However, I get this error:

tensorflow.python.framework.errors_impl.InvalidArgumentError:
 Shape must be rank 4 but is rank 3 for 
'model/batch_normalization/cond/FusedBatchNorm' (op:
 'FusedBatchNorm') with input shapes: [65,65,2],
 [2], [2], [0], [0].

This seems to indicate the batch-dimension didn't make it through. Why?

EDIT: Things I have tried:

(1) Explicitly set trainable=False in all the layers but that didn't seem to make any difference:

  model_core = model
  for layer in model_core.layers:
    layer.trainable = False
  model_output = model_core(to_img)

(2) Tried expanding the result of the preproc:

   to_img = keras.layers.Lambda(
      lambda x : preproc(x))(json_input)
   to_img = keras.layers.Lambda( 
      lambda x : tf.expand_dims(x, axis=0) )(to_img)

This results in an error: AttributeError: 'Model' object has no attribute '_name' on the line serving_model = keras.Model(json_input, model_output)

(3) Changed the lambda layer to do a map_fn to process the data individually:

to_img = keras.layers.Lambda(
    lambda items: K.map_fn(lambda x: preproc, items))(json_input)

This caused a shape error that indicates that the preproc function is getting [65,2] items instead of [65,65,2]. This suggests that the Lambda layer applies the function to examples one at a time.

(4) Here's the full code for the model:

  img = keras.Input(shape=[height, width, 2])

  # convolutional part of model
  cnn = keras.layers.BatchNormalization()(img)
  for layer in range(nlayers):
    nfilters = nfil * (layer + 1)
    cnn = keras.layers.Conv2D(nfilters, (ksize, ksize), padding='same')(cnn)
    cnn = keras.layers.Activation('elu')(cnn)
    cnn = keras.layers.BatchNormalization()(cnn)
    cnn = keras.layers.MaxPooling2D(pool_size=(2, 2))(cnn)
  cnn = keras.layers.Flatten()(cnn)
  cnn = keras.layers.Dropout(dprob)(cnn)
  cnn = keras.layers.Dense(10, activation='relu')(cnn)

  # feature engineering part of model
  engfeat = keras.layers.Lambda(
    lambda x: engineered_features(x, height//2))(img)

  # concatenate the two parts
  both = keras.layers.concatenate([cnn, engfeat])
  ltgprob = keras.layers.Dense(1, activation='sigmoid')(both)

  # create a model
  model = keras.Model(img, ltgprob)
  def rmse(y_true, y_pred):
    import tensorflow.keras.backend as K
    return K.sqrt(K.mean(K.square(y_pred - y_true), axis=-1))
  optimizer = tf.keras.optimizers.Adam(lr=params['learning_rate'],
                                       clipnorm=1.)
  model.compile(optimizer=optimizer,
                loss='binary_crossentropy',
                metrics=['accuracy', 'mse', rmse])

and the code for the preprocessing function:

def reshape_into_image(features, params):
  # stack the inputs to form a 2-channel input
  # features['ref'] is [-1, height*width]
  # stacked image is [-1, height*width, n_channels]
  n_channels = 2
  stacked = tf.concat([features['ref'], features['ltg']], axis=1)
  height = width = PATCH_SIZE(params)
  return tf.reshape(stacked, [height, width, n_channels])

and the serving layer:

  # 1. layer that extracts multiple inputs from JSON
  height = width = PATCH_SIZE(hparams)
  json_input = keras.layers.concatenate([
    keras.layers.Input(name='ref', dtype=tf.float32, shape=(height * width,)),
    keras.layers.Input(name='ltg', dtype=tf.float32, shape=(height * width,)),
  ], axis=0)

  # 2. convert json_input to image (what model wants)
  to_img = keras.layers.Lambda(
    lambda x: reshape_into_image(features={
      'ref': tf.reshape(x[0], [height * width, 1]),
      'ltg': tf.reshape(x[1], [height * width, 1])
    }, params=hparams),
    name='serving_reshape')(json_input)


  # 3. now, use trained model to predict
  model_output = model(to_img)

  # 4. create serving model
  serving_model = keras.Model(json_input, model_output)

Answer 1

The input shape of your model, considering the samples axis, is (?, 65, 65, 2) where ? can be one or more than one. So you need to modify the Lambda layer (actually the function wrapped inside it) such that its output to be (?, 65, 65, 2) as well. One way of doing it is to use K.expand_dims(out, axis=0) in the wrapped function so that the output would have a shape of (1, 65, 65, 2).

By the way, K refers to backend: from keras import backend as K.

Further, note that you must define the function wrapped by the Lambda such that it preserves the batch axis; otherwise, it is very likely that you are doing something wrong in the definition of that function.

Update:

The error AttributeError: 'Model' object has no attribute '_name' is raised because you are passing json_input as the input of the model. However, it is not an input layer. Rather, it is the output of concatenation layer. To resolve this, first define the input layers and then pass them to concatenation layer and Model class, like this:

inputs = [keras.layers.Input(name='ref', dtype=tf.float32, shape=(height * width,)), 
          keras.layers.Input(name='ltg', dtype=tf.float32, shape=(height * width,))]

json_input = keras.layers.concatenate(inputs, axis=0)

# ...
serving_model = keras.Model(inputs, model_output)

Update 2: I think you can write this much more simpler and without going into so much unnecessary trouble. You want to go from two tensors of shape (?, h*w) to a tensor of shape (?, h, w, 2). You can use Reshape layer, so that would be:

from keras.layers import Reshape

inputs = [keras.layers.Input(name='ref', dtype=tf.float32, shape=(height * width,)), 
          keras.layers.Input(name='ltg', dtype=tf.float32, shape=(height * width,))]

reshape_layer = Reshape((height, width, 1))
r_in1 = reshape_layer(inputs[0])
r_in2 = reshape_layer(inputs[1])
img = concatenate([r_in1, r_in2])

output = model(img)

serving_model = keras.Model(inputs, output)

No need for any custom function or Lambda layer.

And by the way, if you are interested to know, the trouble with batch axis removal is caused by this line:

return tf.reshape(stacked, [height, width, n_channels])

You are not considering the batch axis when reshaping.