Использовать затухание скорости обучения тензорного потока в модели Keras-to-TPU

Question

Я следую руководству «Как обучить модель Keras в x20 раз быстрее с помощью TPU бесплатно» ( нажмите здесь ), чтобы запустить модель keras на Google TPAB colab.Работает отлично.Но ... Мне нравится использовать спад скорости обучения при косинусном перезапуске, когда я подхожу к своим моделям.Я закодировал свой собственный как обратный вызов keras, но он не будет работать в этой среде, потому что класс tenorflow TFOptimizer не имеет переменной скорости обучения, которую можно сбросить.Я вижу, что сам тензорный поток имеет функцию распада в 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

Кажется, что работает следующее, где lr - начальная скорость обучения, которую вы выбираете, а M - количество начальных шагов, над которыми вы хотите работать, чтобы косинусный распад работал.

def make_model(batch_size=None,lr=1.e-3,n_steps=2000):
    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])
    # implement cosine decay or other learning rate decay here
    global_step = tf.Variable(0)    
    global_step=1
    learning_rate = tf.train.cosine_decay_restarts(
        learning_rate=lr,
        global_step=global_step,
        first_decay_steps=n_steps,
        t_mul= 1.5,
        m_mul= 1.,
        alpha=0.1
    )
    # now feed this into the optimizer as shown below
    model.compile(
        optimizer=tf.train.RMSPropOptimizer(learning_rate=learning_rate),
        loss='binary_crossentropy',
        metrics=['acc'])
    return model

Answer 2

Один из вариантов - установить скорость обучения вручную - здесь есть пример Keras + TPU с обратным вызовом: https://github.com/tensorflow/tpu/blob/master/models/experimental/resnet50_keras/resnet50.py#L197-L201