Блог

документ tensorflow приводит этот пример

 def train_one_step():
  pass

@tf.function
def train(num_steps):
  print("Tracing with num_steps = ", num_steps)
  tf.print("Executing with num_steps = ", num_steps)
  for _ in tf.range(num_steps):
    train_one_step()

print("Retracing occurs for different Python arguments.")
train(num_steps=10)
train(num_steps=20)

print()
print("Traces are reused for Tensor arguments.")
train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))
 

Вот результат

 Retracing occurs for different Python arguments.
Tracing with num_steps =  10
Executing with num_steps =  10
Tracing with num_steps =  20
Executing with num_steps =  20

Traces are reused for Tensor arguments.
Tracing with num_steps =  Tensor("num_steps:0", shape=(), dtype=int32)
Executing with num_steps =  10
Executing with num_steps =  20
 

Я думаю, я понимаю первую часть, где train(num_steps) определяется как Function объект. train(num_steps=10) и train(num_steps=20) используйте разные значения Python, что приведет к повторению.

Во второй части train(num_steps=tf.constant(10)) и train(num_steps=tf.constant(20)) вызывает повторное использование трассировок, где print("Tracing with num_steps = ", num_steps) выполняется нормально, в то время как часть параметров num_steps=tf.constant(10) захватывается tf.Graph и не выполняется на этапе трассировки.

Правильно ли я понимаю?

Украсьте свою функцию tf.function подсказками tensorflow , чтобы сгенерировать из нее вычислительный график. По умолчанию новый график создается при передаче тензора с новой формой

это создает новый график, так как это первый вызов:

 my_tf_func(tf.constant(10)) 
 

это не приводит к созданию нового графика, так как форма аргумента совпадает с входной формой графика:

 my_tf_func(tf.constant(20)) 
 

это создает новый график, так как форма отличается:

 my_tf_func(tf.constant([10, 20])) 
 

При передаче значений python (не tf.тензоров) каждый раз, когда выполняется вызов с определенным значением, создается новый график. Все эти вызовы генерируют новый график:

 my_tf_func(10) 
my_tf_func(20) 
my_tf_func(30)