TF 2.0 @ tf.function示例

Question

In the tensorflow documentation at the autograph section we have the following code snippet

@tf.function
def train(model, optimizer):
  train_ds = mnist_dataset()
  step = 0
  loss = 0.0
  accuracy = 0.0
  for x, y in train_ds:
    step += 1
    loss = train_one_step(model, optimizer, x, y)
    if tf.equal(step % 10, 0):
      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
  return step, loss, accuracy

step, loss, accuracy = train(model, optimizer)
print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())

I have a small question concerning the step variable, it's an integer and not a tensor, autograph supports built-in python type such as integer. Therefore the tf.equal(step%10,0) could be changed to simply step%10 == 0 right ?

Answer 1

Yes, you're right. The integer variable step remains a Python variable even when converted to its graph representation. You can see the conversion result by calling tf.autograph.to_code(train.python_function) .

Without reporting all the code but only the step variable related part, you'll see that

  def loop_body(loop_vars, loss_1, step_1):
    with ag__.function_scope('loop_body'):
      x, y = loop_vars
      step_1 += 1

Is still a python operation (otherwise it will be step_1.assign_add(1) if step 1 was a tf.Tensor ).

For more information about autograph and tf.function I suggest reading the article https://pgaleone.eu/tensorflow/tf.function/2019/03/21/dissecting-tf-function-part-1/ that explains easily what happens when a function is converted.

Answer 2

Although this is not visible in the generated code, the step variable will actually be autoboxed to a Tensor by the for loop which is being converted to a TF while_loop.

You can verify that by adding a print statement:

    loss = train_one_step(model, optimizer, x, y)
    print(step)
    if tf.equal(step % 10, 0):