Keras
predict function
incorrect shape
machine learning
neural networks

Keras predict getting incorrect shape?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

When model.predict() returns a shape you did not expect, the problem is usually not in predict() itself. It is almost always a mismatch between your mental model of the batch dimension and the actual model output shape. The fastest way to debug it is to compare model.input_shape, model.output_shape, and the shape of the array you passed in.

Start with the Batch Dimension

Keras model input shapes are usually declared without the batch dimension. For example:

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.Input(shape=(4,)),
5    tf.keras.layers.Dense(3, activation="softmax"),
6])
7
8print(model.input_shape)   # (None, 4)
9print(model.output_shape)  # (None, 3)

The leading None means "batch size goes here."

So if you pass one sample shaped (4,), Keras may complain because the model expects batches shaped (batch_size, 4). A single sample still needs a batch axis:

python
1import numpy as np
2
3x = np.array([1.0, 2.0, 3.0, 4.0], dtype="float32")
4x = np.expand_dims(x, axis=0)
5
6pred = model.predict(x, verbose=0)
7print(pred.shape)  # (1, 3)

That (1, 3) result is correct: one sample in the batch, three output units.

The Output Shape Comes from the Last Layer

predict() returns the model's output shape with the batch dimension included. If the last layer has Dense(3), the output per sample has length 3. If you predict one sample, the result is (1, 3), not (3,).

Example:

python
1import tensorflow as tf
2import numpy as np
3
4model = tf.keras.Sequential([
5    tf.keras.Input(shape=(2,)),
6    tf.keras.layers.Dense(5, activation="relu"),
7    tf.keras.layers.Dense(1),
8])
9
10x = np.array([[10.0, 20.0]], dtype="float32")
11pred = model.predict(x, verbose=0)
12
13print(model.output_shape)  # (None, 1)
14print(pred.shape)          # (1, 1)

If you want a scalar-like value from a single-sample prediction, index it deliberately:

python
value = pred[0, 0]

Do not treat the batch dimension as an error unless it truly should not be there.

Image Models Often Need One More Axis

Shape confusion is especially common with images. A single image often starts as:

text
(height, width, channels)

But a Keras image model usually expects:

text
(batch, height, width, channels)

Example:

python
1import numpy as np
2import tensorflow as tf
3
4model = tf.keras.Sequential([
5    tf.keras.Input(shape=(28, 28, 1)),
6    tf.keras.layers.Conv2D(8, 3, activation="relu"),
7    tf.keras.layers.Flatten(),
8    tf.keras.layers.Dense(10, activation="softmax"),
9])
10
11image = np.random.rand(28, 28, 1).astype("float32")
12batch = np.expand_dims(image, axis=0)
13
14pred = model.predict(batch, verbose=0)
15print(pred.shape)  # (1, 10)

If you pass (28, 28, 1) directly, you are missing the batch axis.

Sequence Models Add Another Dimension

RNNs and LSTMs expect shapes like:

text
(batch_size, timesteps, features)

That means one sequence sample must still be wrapped in a batch:

python
1import numpy as np
2import tensorflow as tf
3
4model = tf.keras.Sequential([
5    tf.keras.Input(shape=(5, 2)),
6    tf.keras.layers.LSTM(4),
7    tf.keras.layers.Dense(1),
8])
9
10sequence = np.random.rand(5, 2).astype("float32")
11batch = np.expand_dims(sequence, axis=0)
12
13pred = model.predict(batch, verbose=0)
14print(pred.shape)  # (1, 1)

If return_sequences=True is enabled inside the LSTM, the output shape changes again because the model now emits one output per timestep instead of only the final state.

Multiple Outputs Return Multiple Arrays

If your model has multiple outputs, predict() returns a list or nested structure of arrays, not one array.

python
1import tensorflow as tf
2import numpy as np
3
4inputs = tf.keras.Input(shape=(4,))
5x = tf.keras.layers.Dense(8, activation="relu")(inputs)
6class_output = tf.keras.layers.Dense(3, activation="softmax", name="class")(x)
7score_output = tf.keras.layers.Dense(1, name="score")(x)
8model = tf.keras.Model(inputs, [class_output, score_output])
9
10x_test = np.random.rand(2, 4).astype("float32")
11class_pred, score_pred = model.predict(x_test, verbose=0)
12
13print(class_pred.shape)  # (2, 3)
14print(score_pred.shape)  # (2, 1)

If you were expecting a single array, that can look like a shape bug even though it is the intended structure.

Use the Model Metadata First

Before guessing, inspect the model:

python
print(model.input_shape)
print(model.output_shape)
model.summary()

Those three checks answer most shape questions faster than trial-and-error reshaping.

If the shape mismatch still feels confusing, print the exact input shape right before prediction:

python
print(x.shape)
pred = model.predict(x, verbose=0)
print(pred.shape)

That narrows the issue quickly.

Common Pitfalls

The most common mistake is forgetting the batch dimension and passing a single sample shaped like the per-example input instead of the full batched input.

Another issue is expecting a single-sample classifier output shaped (classes,) when Keras correctly returns (1, classes).

Developers also get tripped up by model architectures that intentionally change output rank, such as convolutional models, sequence models with return_sequences=True, or multi-output models.

Finally, do not debug prediction shape in isolation. Always compare it with model.output_shape, because that is the contract predict() is following.

Summary

  • 'predict() returns the model output shape with the batch dimension included.'
  • Keras Input(shape=...) does not include the batch size, so single examples often need np.expand_dims(..., axis=0).
  • The last layer determines the per-sample output shape.
  • Image, sequence, and multi-output models naturally produce shapes that differ from simple dense classifiers.
  • Check model.input_shape, model.output_shape, and the actual input array shape before changing the model.

Course illustration
Course illustration

All Rights Reserved.