在签名部分的 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
,它是一个整数而不是张量,签名支持内置的 python 类型,例如整数。因此,tf.equal(step%10,0)
可以将其更改为简单step%10 == 0
对吗?