Tensorflow-datasets: Cannot batch tensors of different shapes error even after resize?

Viewed 1435

Having some trouble with the tensorflow-datasets module. Using the stanford_dogs dataset, I resize images to [180,180], yet when the model is trained, from the error message, it appears tensorflow is trying to load the images in their original sizes.

What am I doing wrong?

Code to replicate error (and the error) below. Dataset is around 750mb. Can copy paste it into google colab and run to replicate.

import io
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds

def _normalize_img(img, label):
    img = tf.cast(img, tf.float32) / 255.
    img = tf.image.resize(img,[180,180])
    return (img, label)
    

train_dataset, test_dataset = tfds.load(name="stanford_dogs", split=['train', 'test'], as_supervised=True)

train_dataset = train_dataset.shuffle(1024).batch(32)
train_dataset = train_dataset.map(_normalize_img)

test_dataset = test_dataset.batch(32)
test_dataset = test_dataset.map(_normalize_img)

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64,2,padding='same',activation='relu',input_shape=(180,180,3)),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(32,2,padding='same',activation='relu'),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(120,activation='softmax')
])


model.compile(
    optimizer=tf.keras.optimizers.Adam(0.001),
    loss='sparse_categorical_crossentropy')


history = model.fit(
    train_dataset,
    epochs=5)

Fails with the error:

InvalidArgumentError:  Cannot batch tensors with different shapes in component 0. First element had shape [278,300,3] and element 1 had shape [375,500,3].
     [[node IteratorGetNext (defined at <ipython-input-29-15023f95f627>:39) ]] [Op:__inference_train_function_4908]
1 Answers

You encounter this error because the tf.data.Dataset API cannot create a batch of tensors with different shapes. As the batch function will return Tensors of shape (batch, height, width, channels), the height, width and channels values must be constant throughout the dataset. You can read more about why in the Introduction to Tensors guide.

Batching after resizing will solve your issue :

train_dataset = train_dataset.shuffle(1024)
train_dataset = train_dataset.map(_normalize_img)
# we batch once every image is the same size
train_dataset = train_dataset.batch(32)
Related