tensorflow 2.1 自定义训练
tensorflow
一个面向所有人的开源机器学习框架
项目地址:https://gitcode.com/gh_mirrors/te/tensorflow
免费下载资源
·
常常会遇到自定义网络结构的情况,自定结构后往往会有多个输入,或者还需要自定义loss或者accuracy函数,那么keras的fit就无法使用了,
因此需要自定义训练步骤
下面则自定义一次batch的训练步骤,包含了计算loss,accuracy和梯度下降。
tensorflow2.0 主推eager模式,那么tf.GradientTape则是eager模式下的利器,自动计算梯度并传递
最后别忘了加@tf.function进行封装,使train函数在tensorflow框架下加速运行
然后就可以把封装好的train函数应用到每个batch data上,开始训练
@tf.function
def train_step(input1, input2, target):
loss = 0
acc = 0
with tf.GradientTape() as tape: # 开启自动梯度
predictions = model(input1, input2) # 获取model中call函数的输出
loss += loss_func(target, predictions) # 计算loss
acc += train_acc(target, predictions) # 计算accuracy
variables = model.trainable_variables # 获取model的所有可训练参数,好进行梯度更新
gradients = tape.gradient(loss, variables) # 将loss函数及可训练参数传入得到梯度
optimizer.apply_gradients(zip(gradients, variables)) # 使用优化器进行梯度更新
return loss, acc
GitHub 加速计划 / te / tensorflow
184.55 K
74.12 K
下载
一个面向所有人的开源机器学习框架
最近提交(Master分支:2 个月前 )
a49e66f2
PiperOrigin-RevId: 663726708
2 个月前
91dac11a
This test overrides disabled_backends, dropping the default
value in the process.
PiperOrigin-RevId: 663711155
2 个月前
更多推荐
已为社区贡献2条内容
所有评论(0)