TensorFlow
machine learning
input size
training data
testing data

Tensorflow Is it possible to use different train input size and test input size?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Yes, TensorFlow can handle different train and test input sizes, but only if the model architecture allows it. The real question is not whether TensorFlow supports it in principle, but whether your layers expect fixed dimensions or can operate on variable-sized inputs. This article explains when different sizes work, when they fail, and how to design a model that stays compatible.

Fixed-Shape Models Usually Need the Same Size

If the model contains layers that depend on a fixed input dimension, training and testing must use the same shape for those dimensions.

Typical examples:

  • dense networks on tabular data with a fixed feature count
  • image models that Flatten feature maps before a dense classifier
  • models exported with a fixed input signature

Example of a fixed-shape image model:

python
1import tensorflow as tf
2from tensorflow import keras
3
4model = keras.Sequential([
5    keras.layers.Input(shape=(128, 128, 3)),
6    keras.layers.Conv2D(16, 3, activation="relu"),
7    keras.layers.Flatten(),
8    keras.layers.Dense(10, activation="softmax"),
9])

This model expects 128 x 128 images. Feeding 256 x 256 images at test time will fail because the flattened vector size changes.

Variable Spatial Size Works with Fully Convolutional Designs

If the architecture avoids fixed-size flattening and uses pooling that adapts to the current spatial dimensions, different image sizes can work.

python
1import tensorflow as tf
2from tensorflow import keras
3
4model = keras.Sequential([
5    keras.layers.Input(shape=(None, None, 3)),
6    keras.layers.Conv2D(16, 3, activation="relu"),
7    keras.layers.Conv2D(32, 3, activation="relu"),
8    keras.layers.GlobalAveragePooling2D(),
9    keras.layers.Dense(10, activation="softmax"),
10])
11
12model.build((None, None, None, 3))
13model.summary()

Because the model uses GlobalAveragePooling2D, it can accept different image heights and widths during inference as long as the channel count remains 3.

Sequence Models Often Support Variable Length

Variable-length sequence inputs are also common. Recurrent, transformer, and masking-based models often allow variable sequence length while keeping the feature dimension fixed.

python
1import tensorflow as tf
2from tensorflow import keras
3
4inputs = keras.Input(shape=(None, 32))
5x = keras.layers.Masking()(inputs)
6x = keras.layers.LSTM(64)(x)
7outputs = keras.layers.Dense(1)(x)
8
9model = keras.Model(inputs, outputs)
10model.summary()

Here the time dimension can vary, but each time step must still have 32 features.

Training on One Size and Testing on Another Is Not Automatically Safe

Even if the model accepts a different size, performance can still degrade. A network trained only on 128 x 128 images may technically accept 256 x 256 inputs at inference time, but the learned scale statistics may not generalize well.

That is why there are two separate questions:

  • can the model run on the new size
  • will the predictions still be good

The second question can only be answered empirically.

Common Strategies

The usual approaches are:

  • resize all data to one fixed size
  • train with augmentation across multiple scales
  • use a variable-shape architecture and validate on the intended inference sizes

For many production systems, resizing remains the simplest and most predictable option.

Example: Multi-Scale Training

If robustness across sizes matters, training on varied input sizes can help.

python
1import tensorflow as tf
2
3def preprocess(image, label):
4    target_size = tf.random.shuffle(
5        tf.constant([[128, 128], [160, 160], [192, 192]])
6    )[0]
7    image = tf.image.resize(image, target_size)
8    return image, label

This does not make every model size-agnostic, but it can reduce sensitivity to scale changes.

Tabular Data Is Different

For tabular models, "input size" usually means feature count, not height and width. In that case, different train and test sizes are generally not valid unless you explicitly redesign the feature representation.

For example, a dense layer trained on 20 input features cannot normally accept 25 features at inference time.

Saved Models and Signatures

Even if the architecture is flexible, exported models can still constrain shape. If you save a model with a fixed signature, serving infrastructure may reject different sizes even when the internal layers could handle them.

Always verify:

  • layer-level compatibility
  • model-level input signature
  • serving-time expectations

Those are separate constraints.

Common Pitfalls

  • Assuming TensorFlow flexibility means every trained model can accept any input size.
  • Using Flatten in an image model and then expecting arbitrary spatial dimensions.
  • Changing feature count in tabular models and calling it the same problem as variable image size.
  • Ignoring distribution shift when inference sizes differ from training sizes.
  • Forgetting that exported model signatures may be stricter than the raw layer graph.

Summary

  • Different train and test input sizes are possible only when the architecture supports them.
  • Fixed-size layers such as Flatten plus dense heads usually require the same input size.
  • Fully convolutional models with global pooling can often accept variable image sizes.
  • Variable sequence length is common, but feature dimension usually stays fixed.
  • Compatibility at runtime does not guarantee good accuracy, so validate on the actual test sizes you plan to use.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design