using a `tf.Tensor` as a Python `bool` is not allowed in Graph execution with `tf.data.Dataset`

Question

I'm trying to make this function accept a single element tensor.

def classi(i):
    out = np.zeros((1, 49), np.uint8)
    for j in range(len(classcount)):
        i -= classcount[j]
        if i<0:
            break
    out[0][j] += 1
    return tf.convert_to_tensor(out)
    #basically the error seems to be related to the if i<0 line

This function will be called by another function here

def formatcars(elem):
    return (elem['image'], tf.function(classi(elem['label'])))
    #note elem['label'] is just a single element tensor of integer.

Which in turn is mapped to the cars dataset.

dataset.map(formatcars)

And I keep getting the error:

OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

I've tried enabling eager execution. I've tried using tf.function, using tf.cond, tf.greater, .tonumpy(), .eval(), etc. to no avail. It keeps giving the same error. I'm out of ideas now.

The classcount list is defined as follow:

classcount = [ 1,  6,  4, 14, 13,  6,  2,  4,  3, 22,  6,  1, 15,  1,  2,  4,  1,
       12,  5,  1,  2,  4, 11,  2,  1,  1,  5,  4,  2,  1,  1,  1,  1,  1,
        6,  1,  4,  1,  1,  1,  3,  1,  2,  4,  1,  4,  3,  3,  1]

it's simply a list of integers created from

import scipy
import tensorflow_datasets as tfds

dataset = tfds.load('cars196', split = 'train')
mat = scipy.io.loadmat('cars_annos.mat')
classcount = []
starti = 0
curmake = ''
for i in range(len(mat['class_names'][0])):
    print(mat['class_names'][0][i][0].split(' ', 1)[0])
    if mat['class_names'][0][i][0].split(' ', 1)[0] != curmake:
        print(i-starti)
        if i-starti != 0:
            classcount.append(i-starti)
        starti = i
    curmake = mat['class_names'][0][i][0].split(' ', 1)[0]
classcount.append(1)

cars_annos.mat is from http://imagenet.stanford.edu/internal/car196/cars_annos.mat

Answer 1

As the error states, you can't use a tensor as a boolean value in a python conditional statement outside of eager execution, and the tf.data.Dataset will force you to use graph mode (for performances reason). You can't simply decorate the function with the @tf.function decorator either, because Autograph is not able to convert a code where a tensor is used as a python bool (in a condition statement for example).

The best way of handling that is to rewrite the function using TF ops. One way of doing it could be the following:

def graphmode_classi(i):
    """ 
    performs a cumulative sum on classcount and substract that sum from i
    the first negative value will give us the value to one hot encode
    """  
    cumsum = tf.math.cumsum(classcount)
    first_neg = tf.where((i - cumsum)<0)[0]
    return tf.one_hot(first_neg, 49)

We can compare if my rewritten function is equivalent:

# random data
data = tf.cast(tf.random.normal((200, 1))**2 * 10, tf.int32)
for d in data:
    assert (classi(d).numpy() == graphmode_classi(d).numpy()).all()

Now you should be able to use the function written with only tf ops with the tf.data.Dataset API:

data = tf.cast(tf.random.normal((200, 1))**2 * 10, tf.int32)
ds = tf.data.Dataset.from_tensor_slices(data)
ds.map(graphmode_classi)