tf.keras.backend.function for transforming embeddings inside tf.data.dataset

Question

I am trying to use the output of a neural network to transform data inside tf.data.dataset. Specifically, I am using a Delta-Encoder to manipulate embeddings inside the tf.data pipeline. In so doing, however, I get the following error:

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

I have searched the dataset pipeline page and stack overflow, but I could not find something that addresses my question. In the code below I am using an Autoencoder, as it yields an identical error with more concise code.

The offending part seems to be [[x,]] = tf.py_function(Auto_Func, [x], [tf.float32]) inside tf_auto_transform.

num_embeddings = 100
input_dims = 1000
embeddings = np.random.normal(size = (num_embeddings, input_dims)).astype(np.float32)
target = np.zeros(num_embeddings)

#creating Autoencoder
inp = Input(shape = (input_dims,), name ='input')
hidden = Dense(10, activation = 'relu', name = 'hidden')(inp)
out = Dense(input_dims, activation = 'relu', name='output')(hidden)
auto_encoder = tf.keras.models.Model(inputs =inp, outputs=out)

Auto_Func = tf.keras.backend.function(inputs = Autoencoder.get_layer(name='input').input, 
                                      outputs = Autoencoder.get_layer(name='output').input )

#Autoencoder transform for dataset.map
def tf_auto_transform(x, target):
    x_shape = x.shape
    #@tf.function
    #def func(x):
    #    return tf.py_function(Auto_Func, [x], [tf.float32])
    #[[x,]] = func(x)
    [[x,]] = tf.py_function(Auto_Func, [x], [tf.float32]) 
    x.set_shape(x_shape) 
    return x, target

def get_dataset(X,y, batch_size = 32):
    train_ds = tf.data.Dataset.from_tensor_slices((X, y))
    train_ds = train_ds.map(tf_auto_transform)
    train_ds = train_ds.batch(batch_size)
    return train_ds

dataset = get_dataset(embeddings, target, 2)

The above code yields the following error:

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

I tried to eliminate the error by running the commented out section of the tf_auto_transform function, but the error persisted.

SideNote: While it is true that the Delta encoder paper has code, it is written in tf 1.x. I am trying to use tf 2.x with the tf functional API instead. Thank you for your help!

Answer 1

At the risk of outing myself as a n00b, the answer is to switch the order of the map and batch functions. I am trying to apply a neural network to make some changes on data. tf.keras models take batches as input, not individual samples. By batching the data first, I can run batches through my nn.

def get_dataset(X,y, batch_size = 32):
    train_ds = tf.data.Dataset.from_tensor_slices((X, y))
    
    #The changed order
    train_ds = train_ds.batch(batch_size)
    train_ds = train_ds.map(tf_auto_transform)**strong text**
    return train_ds

It really is that simple.