TF 2.0 @tf.function 示例
TF 2.0 @tf.function example
在 autograph 部分的 tensorflow 文档中,我们有以下代码片段
@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())
我有一个关于step
变量的小问题,它是一个整数而不是张量,autograph支持内置的python类型,例如整数。因此 tf.equal(step%10,0)
可以简单地更改为 step%10 == 0
对吗?
是的,你是对的。整数变量 step 仍然是一个 Python 变量,即使在转换为它的图形表示时也是如此。调用tf.autograph.to_code(train.python_function)
.
即可看到转换结果
不报告所有代码,只报告 step
变量相关部分,您会看到
def loop_body(loop_vars, loss_1, step_1):
with ag__.function_scope('loop_body'):
x, y = loop_vars
step_1 += 1
仍然是一个 python 操作(否则如果第 1 步是 tf.Tensor
,它将是 step_1.assign_add(1)
)。
有关 autograph 和 tf.function 的更多信息,我建议阅读文章 https://pgaleone.eu/tensorflow/tf.function/2019/03/21/dissecting-tf-function-part-1/,该文章轻松解释了函数转换时发生的情况。
虽然这在生成的代码中不可见,但 step 变量实际上会被正在转换为 TF 的 for 循环自动装箱到 Tensor while_loop.
您可以通过添加打印语句来验证:
loss = train_one_step(model, optimizer, x, y)
print(step)
if tf.equal(step % 10, 0):
在 autograph 部分的 tensorflow 文档中,我们有以下代码片段
@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())
我有一个关于step
变量的小问题,它是一个整数而不是张量,autograph支持内置的python类型,例如整数。因此 tf.equal(step%10,0)
可以简单地更改为 step%10 == 0
对吗?
是的,你是对的。整数变量 step 仍然是一个 Python 变量,即使在转换为它的图形表示时也是如此。调用tf.autograph.to_code(train.python_function)
.
不报告所有代码,只报告 step
变量相关部分,您会看到
def loop_body(loop_vars, loss_1, step_1):
with ag__.function_scope('loop_body'):
x, y = loop_vars
step_1 += 1
仍然是一个 python 操作(否则如果第 1 步是 tf.Tensor
,它将是 step_1.assign_add(1)
)。
有关 autograph 和 tf.function 的更多信息,我建议阅读文章 https://pgaleone.eu/tensorflow/tf.function/2019/03/21/dissecting-tf-function-part-1/,该文章轻松解释了函数转换时发生的情况。
虽然这在生成的代码中不可见,但 step 变量实际上会被正在转换为 TF 的 for 循环自动装箱到 Tensor while_loop.
您可以通过添加打印语句来验证:
loss = train_one_step(model, optimizer, x, y)
print(step)
if tf.equal(step % 10, 0):