TensorFlow在model.fit自定义学习率变化

import tensorflow.keras.backend as K
# 自定义学习率的变化
def scheduler(epoch):
    # 每隔1个epoch，学习率减小为原来的1/10
    # if epoch % 1 == 0 and epoch != 0:
    #     lr = K.get_value(model.optimizer.lr)
    #     K.set_value(model.optimizer.lr, lr * 0.1)
    #     print("lr changed to {}".format(lr * 0.1))
    if epoch % 20 == 19 and epoch != 0:             # 每测试20轮，学习率降低一半
        lr = K.get_value(model.optimizer.lr)
        K.set_value(model.optimizer.lr, lr * 0.5)
        print("lr changed to {}".format(lr * 0.5))
    return K.get_value(model.optimizer.lr)

model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.summary()
model.compile(optimizer=tf.keras.optimizers.Adam(lr=3e-5, beta_1=0.5, beta_2=0.99),
              loss=v_loss,
              metrics='accuracy')
# 自定义学习率
reduce_lr = LearningRateScheduler(scheduler)
# reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', patience=50, mode='auto')
history = model.fit(x_train, y_train,
                    batch_size=1024,
                    epochs=800,
                    initial_epoch=1,
                    validation_data=(x_test, y_test),
                    callbacks=[reduce_lr])