在TensorFlow2中使用学习率计划和学习率预热

import tensorflow as tf
from tensorflow.keras import layers

# 定义模型
model = tf.keras.Sequential([
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# 定义学习率计划
initial_learning_rate = 0.1
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate,
    decay_steps=1000,
    decay_rate=0.96,
    staircase=True
)

# 定义学习率预热
warmup_steps = 100
warmup_lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
    initial_learning_rate / 10,
    warmup_steps,
    end_learning_rate=initial_learning_rate,
    power=1.0
)

# 定义优化器
optimizer = tf.keras.optimizers.SGD(learning_rate=warmup_lr_schedule)

# 定义损失函数和评估指标
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
accuracy_metric = tf.keras.metrics.SparseCategoricalAccuracy()

# 定义训练过程
@tf.function
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        logits = model(inputs, training=True)
        loss_value = loss_fn(labels, logits)
    grads = tape.gradient(loss_value, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    accuracy_metric(labels, logits)
    return loss_value

# 进行训练
for epoch in range(num_epochs):
    for step, (inputs, labels) in enumerate(train_dataset):
        loss_value = train_step(inputs, labels)
    
    # 更新学习率
    if epoch < warmup_steps:
        optimizer.learning_rate.assign(warmup_lr_schedule(epoch))
    else:
        optimizer.learning_rate.assign(lr_schedule(epoch - warmup_steps))
    
    # 打印训练结果
    print('Epoch {}: loss = {}, accuracy = {}'.format(
        epoch+1, loss_value, accuracy_metric.result()
    ))
    accuracy_metric.reset_states()