- 入门教程中,定义loss和grad的计算函数如下:
def loss(model, x, y):
y_ = model(x)
return tf.losses.sparse_softmax_cross_entropy(labels=y, logits=y_)
def grad(model, inputs, targets):
with tf.GradientTape() as tape:
loss_value = loss(model, inputs, targets)
return tape.gradient(loss_value, model.variables)
不写with tf.GradientTape() as tape:
时应该就是相当于pytorch里的关闭梯度。也就是pytorch默认打开梯度,TF Eager默认关闭梯度,这一点还是不错的。
注意到,loss重复计算了啊,不如用一个函数。
def loss_and_grad(model, inputs, targets):
with tf.GradientTape() as tape:
loss_value = loss(model, inputs, targets)
return loss_value, tape.gradient(loss_value, model.variables)
经过测试,确实可以节省时间,对结果没影响。不改变入门教程-鸢尾花分类中的其他参数,原例程时间训练总时间约7.2s,修改后约5.9s。