常常会遇到自定义网络结构的情况,自定结构后往往会有多个输入,或者还需要自定义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 个月前
Logo

旨在为数千万中国开发者提供一个无缝且高效的云端环境,以支持学习、使用和贡献开源项目。

更多推荐