Utiliser un dictionnaire dans la fonction tf.input_signature dans Tensorflow 2.0

Question

J'utilise Tensorflow 2.0 et je suis confronté à la situation suivante :

@tf.function
def my_fn(items):
    .... #do stuff
    return

Si items est un dict de tenseurs comme par exemple :

item1 = tf.zeros([1, 1])
item2 = tf.zeros(1)
items = {"item1": item1, "item2": item2}

Existe-t-il un moyen d'utiliser l'argument input_signature de tf.function afin de forcer tf2 à ne pas créer de graphiques multiples lorsque l'élément 1 est par exemple tf.zeros([2,1]) ?

Answer 1

La signature d'entrée doit être une liste, mais les éléments de la liste peuvent être des dictionnaires ou des listes de Tensor Specs. Dans votre cas, j'essaierais : (le name sont facultatifs)

signature_dict = { "item1": tf.TensorSpec(shape=[2], dtype=tf.int32, name="item1"),
                   "item2": tf.TensorSpec(shape=[], dtype=tf.int32, name="item2") } 

# don't forget the brackets around the 'signature_dict'
@tf.function(input_signature = [signature_dict])
def my_fn(items):
    .... # do stuff
    return

# calling the TensorFlow function
my_fun(items)

Toutefois, si vous souhaitez appeler une fonction concrète particulière créée par my_fn Il faut déballer le dictionnaire. Vous devez également fournir le name dans l'attribut tf.TensorSpec .

# creating a concrete function with an input signature as before but without
# brackets and with mandatory 'name' attributes in the TensorSpecs 
my_concrete_fn = my_fn.get_concrete_function(signature_dict)

# calling the concrete function with the unpacking operator
my_concrete_fn(**items)

Ceci est ennuyeux mais devrait être résolu dans TensorFlow 2.3 (voir la fin du TF Guide to 'Concrete functions').