Использовать словарь в tf.function input_signature в Tensorflow 2.0

Я использую Tensorflow 2.0 и сталкиваюсь со следующей ситуацией:

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

Если items - это тензоры, например:

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

Есть ли способ использовать аргумент input_signature tf.function, чтобы я мог заставить tf2 избегать создания нескольких графиков, когда item1, например, tf.zeros([2,1])?


python tensorflow2.0
person m33n    schedule 24.03.2020    source источник

Ответы (1)


arrow_upward
6
arrow_downward

Входная сигнатура должна быть списком, но элементы в списке могут быть словарями или списками спецификаций тензор. В вашем случае я бы попробовал: (атрибуты name необязательны)

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)

Однако, если вы хотите вызвать конкретную конкретную функцию, созданную my_fn, вам нужно распаковать словарь. Вы также должны указать атрибут name в 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)

Это раздражает, но должно быть решено в TensorFlow 2.3. (см. конец Руководства TF по «Конкретным функциям»)

person Sven Meinhardt    schedule 26.06.2020