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