Dataset.maps causes model.fit to fail in 2.10 but the same works fine 2.9.x

TensorFlow General Discussion
,
1

The below sample works fine in 2.9.x but fails in 2.10:

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node _wrapped__Reshape_device/job:localhost/replica:0/task:0/device:GPU:0}} Input to reshape is a tensor with 10 values, but the requested shape has 3072 [Op:Reshape]

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def augment(rgb, label):          
    x_vals = tf.reshape(rgb, [32, 32, 3])
    x_vals = rgb / tf.constant(255, dtype=tf.uint8)
    y_vals = tf.one_hot(label, depth=10)
    y_vals = tf.reshape(y_vals, [10])
    return x_vals, y_vals

def fixup_shape(rgb, label):
    rgb.set_shape([32, 32, 3])
    label.set_shape([10])
    return rgb, label

def make_datasets(batch_size=None):

    datasets = {}

    def prepare_dataset(dataset, type):
        dataset = dataset.map(lambda x, y: tf.py_function(augment, [x, y], (np.float32, np.float32)), num_parallel_calls=tf.data.experimental.AUTOTUNE)
        dataset = dataset.map(fixup_shape)
        return dataset

    cifar10 = tf.keras.datasets.cifar10
    (x_train, y_train), (x_test, y_test) = cifar10.load_data()

    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))

    datasets["train"] = prepare_dataset(train_dataset, "train").batch(batch_size)
    datasets["test"] = prepare_dataset(test_dataset, "test").batch(batch_size)

    return datasets

datasets = make_datasets(batch_size=128)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def generate_images(test_input, labels):
  plt.figure(figsize=(20, 8))
  for i in range(16):
    plt.subplot(4, 16, i*2+1)
    plt.imshow(test_input[i].numpy())
    # plt.title(classes[tf.argmax(labels[i], axis=0)])
  plt.axis('off')
  plt.show() 

for imgs, labels in datasets["train"].take(1):
  generate_images(imgs, labels)

for imgs, labels in datasets["test"].take(1):
  generate_images(imgs, labels)

# Model

model = tf.keras.models.Sequential([
    # CNN
    tf.keras.layers.Conv2D(32, kernel_size=(3,3), padding='same', activation='relu', input_shape=(32,32,3)),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(2,2)),
    tf.keras.layers.Dropout(0.25),

    tf.keras.layers.Conv2D(64, kernel_size=(3,3), padding='same', activation='relu',),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(2,2)),
    tf.keras.layers.Dropout(0.25),

    tf.keras.layers.Conv2D(128, kernel_size=(3,3), padding='same', activation='relu',),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.MaxPool2D(pool_size=(2,2)),
    tf.keras.layers.Dropout(0.25),

    # Deep Layers
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(512, activation="relu"),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(10, activation="softmax"),

])

model.compile(optimizer="Adam", loss="categorical_crossentropy", metrics=["accuracy"])

history = model.fit(datasets["train"], batch_size=128, epochs=2, validation_data=datasets["test"], callbacks=[])
3

Hi @T_Palit,

I was able to reproduce the issue in this gist. Thanks for reporting the issue we will try to look into this issue.

Thanks.

5

Hi @T_Palit,

I have tried executing code in CPU and GPU run times, code succefully executed in CPU runtime but not on GPU. When it comes to tf.py_functions it has some limitations which can be found here. Also find this gist, which I have executed using gpu and only making tf.py_function operations in CPU. The operations which are being done have some problem from GPU side. If you just want to do one hot encoding in augment function, I have included sparse_categorical_crossentropy model training, which doesn’t require you to do one hot encoding and much more robust and reliable. I hope this will help you.

Thanks & Regards,
Sravana Neeli.