tensorflow 2.0 自定义训练

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