InvalidArgumentError: ConcatOp : Dimensions of inputs should match when predicting on X

I'm learning Tensorflow and am trying to build a classifier on the Fashion MNIST dataset. I can fit the model, but when I try to predict on my test set I get the following error:

y_pred = model.predict(X_test).argmax(axis=1)

InvalidArgumentError: ConcatOp : Dimensions of inputs should match: shape[0] = [1,32,10] vs. shape[312] = [1,16,10] [Op:ConcatV2] name: concat

I don't get an error if I predict on X_test in batches, for example:

y_pred = []
step_size = 10
for i in trange(0, len(X_test), step_size):
  y_pred  = model.predict(X_test[i:i step_size]).argmax(axis=1).tolist()[0]

I've spent some time googling and looking at other examples of the same error but still can't figure out what I'm doing wrong. I've tried a few different things, such as applying the scale and expand dimensions steps manually to X_train and X_test before building the model, but get the same result.

This is my full code (using Python 3.7.12 and Tensorflow 2.7.0):

import tensorflow as tf # 2.7.0
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# load data 
mnist = tf.keras.datasets.fashion_mnist
(X_train, y_train), (X_test, y_test) = mnist.load_data()

# Build model 

# Input 
inputs = tf.keras.Input(shape=X_train[0].shape)

# # Scale 
x = tf.keras.layers.Rescaling(scale=1.0/255)(inputs)

# Add extra dimension for use in conv2d
x = tf.expand_dims(x, -1)

# Conv2D
x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), activation="relu", strides=2)(x)
x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), activation="relu", strides=2)(x)
x = tf.keras.layers.Conv2D(filters=128, kernel_size=(3, 3), activation="relu", strides=2)(x)

# Flatten
x = tf.keras.layers.Flatten()(x),
x = tf.keras.layers.Dropout(rate=.2)(x)  # 20% chance of dropout  
x = tf.keras.layers.Dense(512, activation='relu')(x)
x = tf.keras.layers.Dropout(rate=.2)(x)   
x = tf.keras.layers.Dense(K, activation='softmax')(x)

model = tf.keras.Model(inputs=inputs, outputs=x)

# Compile
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Fit
r = model.fit(X_train, y_train, validation_data=[X_test, y_test], epochs=10)

# Throws an error
y_pred = model.predict(X_test).argmax(axis=1)

Which gives

InvalidArgumentError: ConcatOp : Dimensions of inputs should match: shape[0] = [1,32,10] vs. shape[312] = [1,16,10] [Op:ConcatV2] name: concat

CodePudding user response：

With model.predict you are making predictions on batches as stated here:

Computation is done in batches. This method is designed for batch processing of large numbers of inputs. It is not intended for use inside of loops that iterate over your data and process small numbers of inputs at a time.

But the size of X_test is not evenly divisible by the default batch_size=32. I think this might be the cause of your problem. You could change your batch_size to 16 for example and it will work:

y_pred = model.predict(X_test, batch_size=16).argmax(axis=1)
print(y_pred)

[[ 8  0  2 ... 14  8  2]
 [15 15  8 ... 10  8 14]
 [ 5 13  4 ...  4  5  6]
 ...
 [11 11 12 ...  7  2  3]
 [ 3  8  0 ... 15  3 14]
 [ 3 13  1 ...  1 15  0]]

You could also use model.predict_on_batch(X_test) to make predictions for a single batch of samples. However, you are most flexible if you use the call function of your model directly:

y_pred = model(X_test[:10])
tf.print(tf.argmax(y_pred, axis=1), summarize=-1)

[[2 8 0 1 1 1 8 2 2 6]]