python - 在 Keras 到 TPU 模型中使用 tensorflow 学习率衰减

Question

我按照“如何使用 TPU 免费训练 Keras 模型速度提高 20 倍”指南 ( click here ) 在 Google 的 colab TPU 上运行 keras 模型。它工作完美。但是......当我适合我的模型时，我喜欢使用余弦重新启动学习率衰减。我已经将自己的代码编写为 keras 回调，但它无法在此框架中工作，因为 TensorFlow TFOptimizer 类没有可以重置的学习率变量。我看到tensorflow本身在tf.train中有一堆衰减函数，比如tf.train.cosine_decay，但我不知道如何将它嵌入到我的模型中。

这是该博客文章中的基本代码。有人解决了吗？

import tensorflow as tf
import os
from tensorflow.python.keras.layers import Input, LSTM, Bidirectional, Dense, Embedding

def make_model(batch_size=None):
    source = Input(shape=(maxlen,), batch_size=batch_size,
                   dtype=tf.int32, name='Input')
    embedding = Embedding(input_dim=max_features,
                          output_dim=128, name='Embedding')(source)
    lstm = LSTM(32, name='LSTM')(embedding)
    predicted_var = Dense(1, activation='sigmoid', name='Output')(lstm)
    model = tf.keras.Model(inputs=[source], outputs=[predicted_var])
    model.compile(
        optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),
        loss='binary_crossentropy',
        metrics=['acc'])
    return model

training_model = make_model(batch_size=128)

# This address identifies the TPU we'll use when configuring TensorFlow.
TPU_WORKER = 'grpc://' + os.environ['COLAB_TPU_ADDR']
tf.logging.set_verbosity(tf.logging.INFO)

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    training_model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

history = tpu_model.fit(x_train, y_train,
                    epochs=20,
                    batch_size=128 * 8,
                    validation_split=0.2)

Answer 1

一种选择是手动设置学习率 - 这里有一个带有回调的 Keras+TPU 示例:https://github.com/tensorflow/tpu/blob/master/models/experimental/resnet50_keras/resnet50.py#L197-L201