Aquecimento da taxa de aprendizado com decaimento de cosseno em Keras/TensorFlow

A taxa de aprendizagem é um hiperparâmetro importante em redes de aprendizagem profunda – e dita diretamente a grau para o qual são realizadas atualizações nos pesos, que são estimados para minimizar alguma determinada função de perda. Em SGD:

$$
peso_{t+1} = peso_t – lr * frac{derror}{dpeso_t}
$$

Com uma taxa de aprendizagem de 0, o peso atualizado voltou a si mesmo – peso_t. A taxa de aprendizagem é efetivamente um botão que podemos girar para ativar ou desativar a aprendizagem e tem grande influência sobre a quantidade de aprendizagem que está acontecendo, controlando diretamente o grau de atualizações de peso.

Diferentes otimizadores utilizam taxas de aprendizagem de maneira diferente – mas o conceito subjacente permanece o mesmo. Escusado será dizer que as taxas de aprendizagem têm sido objeto de muitos estudos, artigos e referências de profissionais.

De modo geral, quase todo mundo concorda que uma taxa de aprendizagem estática não vai resolver isso, e algum tipo de redução da taxa de aprendizagem acontece na maioria das técnicas que ajustam a taxa de aprendizagem durante o treinamento - seja monotônico, cosseno, triangular ou outros tipos de redução.

Uma técnica que vem ganhando espaço nos últimos anos é a aquecimento da taxa de aprendizagem, que pode ser combinada com praticamente qualquer outra técnica de redução.

Aquecimento da Taxa de Aprendizagem

A ideia por trás do aquecimento da taxa de aprendizagem é simples. Nas primeiras fases do treino – os pesos estão longe dos seus estados ideais. Isto significa grandes atualizações em todos os aspectos, que podem ser vistas como “correções excessivas” para cada peso – onde a atualização drástica de outro pode anular a atualização de algum outro peso, tornando os estágios iniciais do treinamento mais instáveis.

Essas mudanças são resolvidas, mas podem ser evitadas com uma taxa de aprendizado pequena para começar, atingindo um estado abaixo do ideal mais estável e, em seguida, aplicando uma taxa de aprendizado maior. Você pode facilitar as atualizações da rede, em vez de acessá-las.

Isso é aquecimento da taxa de aprendizagem! Começando com uma taxa de aprendizado baixa (ou 0) e aumentando até uma taxa de aprendizado inicial (a qual você começaria de qualquer maneira). Esse aumento pode realmente seguir qualquer função, mas geralmente é linear.

Depois de atingir a taxa inicial, outros cronogramas, como decaimento de cosseno, redução linear, etc., podem ser aplicados para diminuir progressivamente a taxa até o final do treinamento. O aquecimento da taxa de aprendizagem geralmente faz parte de um cronograma de dois horários, onde o aquecimento LR é o primeiro, enquanto outro cronograma assume o controle depois que a taxa atinge um ponto inicial.

Neste guia, implementaremos um aquecimento da taxa de aprendizagem no Keras/TensorFlow como um keras.optimizers.schedules.LearningRateSchedule subclasse e keras.callbacks.Callback ligar de volta. A taxa de aprendizagem será aumentada de 0 para target_lr e aplique o decaimento do cosseno, pois este é um esquema secundário muito comum. Como sempre, Keras simplifica a implementação de soluções flexíveis de várias maneiras e o envio delas com sua rede.

Taxa de aprendizagem com retornos de chamada Keras

A maneira mais simples de implementar qualquer cronograma de taxa de aprendizagem é criar uma função que receba o lr parâmetro (float32), passa por alguma transformação e o retorna. Esta função é então passada para o LearningRateScheduler retorno de chamada, que aplica a função à taxa de aprendizado.

Agora, a tf.keras.callbacks.LearningRateScheduler() passa o número da época para a função usada para calcular a taxa de aprendizado, o que é bastante grosseiro. O aquecimento LR deve ser feito em cada passo (lote), não época, então teremos que derivar um global_step (em todas as épocas) para calcular a taxa de aprendizagem e subclassificar o Callback class para criar um retorno de chamada personalizado em vez de apenas passar a função, já que precisaremos passar argumentos em cada chamada, o que é impossível apenas passando a função:

def func():
    return ...
    
keras.callbacks.LearningRateScheduler(func)

Essa abordagem é favorável quando você não deseja um alto nível de personalização e não quer interferir na forma como Keras trata o cliente. lr, e especialmente se você quiser usar retornos de chamada como ReduceLROnPlateau() já que só pode funcionar com um sistema baseado em float lr. Vamos implementar um aquecimento da taxa de aprendizagem usando um retorno de chamada Keras, começando com uma função de conveniência:

def lr_warmup_cosine_decay(global_step,
                           warmup_steps,
                           hold = 0,
                           total_steps=0,
                           start_lr=0.0,
                           target_lr=1e-3):
    
    learning_rate = 0.5 * target_lr * (1 + np.cos(np.pi * (global_step - warmup_steps - hold) / float(total_steps - warmup_steps - hold)))

    
    warmup_lr = target_lr * (global_step / warmup_steps)

    
    
    if hold > 0:
        learning_rate = np.where(global_step > warmup_steps + hold,
                                 learning_rate, target_lr)
    
    learning_rate = np.where(global_step < warmup_steps, warmup_lr, learning_rate)
    return learning_rate

Em cada etapa, calculamos a taxa de aprendizagem e a taxa de aprendizagem de aquecimento (ambos elementos do cronograma), com relação ao start_lr e target_lr. start_lr normalmente começará às 0.0, Enquanto que o target_lr depende da sua rede e do otimizador – 1e-3 pode não ser um bom padrão, portanto, certifique-se de definir seu destino iniciando o LR ao chamar o método.

Se o global_step no treinamento é maior que o warmup_steps definimos – usamos o cronograma de decaimento do cosseno LR. Caso contrário, significa que ainda estamos aquecendo, então o aquecimento LR é usado. Se o hold argumento está definido, manteremos o target_lr para esse número de etapas após o aquecimento e antes do decaimento do cosseno. np.where() fornece uma ótima sintaxe para isso:

np.where(condition, value_if_true, value_if_false)

Você pode visualizar a função com:

steps = np.arange(0, 1000, 1)
lrs = []

for step in steps:
  lrs.append(lr_warmup_cosine_decay(step, total_steps=len(steps), warmup_steps=100, hold=10))
plt.plot(lrs)

Agora, queremos usar esta função como parte de um retorno de chamada e passar a etapa do otimizador como global_step em vez de um elemento de uma matriz arbitrária – ou você pode realizar o cálculo dentro da classe. Vamos subclassificar o Callback classe:

from keras import backend as K

class WarmupCosineDecay(keras.callbacks.Callback):
    def __init__(self, total_steps=0, warmup_steps=0, start_lr=0.0, target_lr=1e-3, hold=0):

        super(WarmupCosineDecay, self).__init__()
        self.start_lr = start_lr
        self.hold = hold
        self.total_steps = total_steps
        self.global_step = 0
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.lrs = []

    def on_batch_end(self, batch, logs=None):
        self.global_step = self.global_step + 1
        lr = model.optimizer.lr.numpy()
        self.lrs.append(lr)

    def on_batch_begin(self, batch, logs=None):
        lr = lr_warmup_cosine_decay(global_step=self.global_step,
                                    total_steps=self.total_steps,
                                    warmup_steps=self.warmup_steps,
                                    start_lr=self.start_lr,
                                    target_lr=self.target_lr,
                                    hold=self.hold)
        K.set_value(self.model.optimizer.lr, lr)

Primeiro, definimos o construtor da classe e monitoramos seus campos. A cada lote finalizado, aumentaremos o passo global, anotaremos o LR atual e o adicionaremos à lista de LRs até o momento. No início de cada lote – calcularemos o LR usando o lr_warmup_cosine_decay() função e defina esse LR como o LR atual do otimizador. Isso é feito com o backend set_value().

Feito isso – basta calcular o total de etapas (tamanho/tamanho do lote*épocas) e pegar uma parte desse número para o seu warmup_steps:

total_steps = len(train_set)*config['EPOCHS']



warmup_steps = int(0.05*total_steps)

callback = WarmupCosineDecay(total_steps=total_steps, 
                             warmup_steps=warmup_steps,
                             hold=int(warmup_steps/2), 
                             start_lr=0.0, 
                             target_lr=1e-3)

Por fim, construa seu modelo e forneça o retorno de chamada no fit() ligue para:

model = keras.applications.EfficientNetV2B0(weights=None, 
                                            classes=n_classes, 
                                            input_shape=[224, 224, 3])
  
model.compile(loss="sparse_categorical_crossentropy",
                  optimizer='adam',
                  jit_compile=True,
                  metrics=['accuracy'])

Ao final do treinamento, você poderá obter e visualizar os LRs alterados via:

lrs = callback.lrs 
plt.plot(lrs)

Se você traçar o histórico de um modelo treinado com e sem aquecimento LR – você verá uma diferença distinta na estabilidade do treinamento:

Taxa de aprendizagem com subclasse LearningRateSchedule

Uma alternativa para criar um retorno de chamada é criar um LearningRateSchedule subclasse, que não manipula o LR – ela o substitui. Essa abordagem permite que você investigue um pouco mais no back-end do Keras/TensorFlow, mas, quando usada, não pode ser combinada com outros retornos de chamada relacionados ao LR, como ReduceLROnPlateau(), que trata de LRs como números de ponto flutuante.

Além disso, o uso da subclasse exigirá que você a torne serializável (sobrecarga get_config()) à medida que se torna parte do modelo, se você quiser salvar os pesos do modelo. Outra coisa a notar é que a turma espera trabalhar exclusivamente com tf.TensorS. Felizmente, a única diferença na forma como trabalhamos será ligar tf.func() em vez de np.func() já que as APIs TensorFlow e NumPy são incrivelmente semelhantes e compatíveis.

Vamos reescrever a conveniência lr_warmup_cosine_decay() função para usar operações do TensorFlow:

def lr_warmup_cosine_decay(global_step,
                           warmup_steps,
                           hold = 0,
                           total_steps=0,
                           start_lr=0.0,
                           target_lr=1e-3):
    
    
    learning_rate = 0.5 * target_lr * (1 + tf.cos(tf.constant(np.pi) * (global_step - warmup_steps - hold) / float(total_steps - warmup_steps - hold)))

    
    warmup_lr = target_lr * (global_step / warmup_steps)

    
    
    if hold > 0:
        learning_rate = tf.where(global_step > warmup_steps + hold,
                                 learning_rate, target_lr)
    
    learning_rate = tf.where(global_step < warmup_steps, warmup_lr, learning_rate)
    return learning_rate

Com a função de conveniência, podemos subclassificar o LearningRateSchedule aula. Em cada __call__() (lote), calcularemos o LR usando a função e o retornaremos. Naturalmente, você também pode empacotar o cálculo dentro da classe de subclasse.

A sintaxe é mais limpa que a Callback sublcass, principalmente porque temos acesso ao step campo, em vez de rastreá-lo por conta própria, mas também torna um pouco mais difícil trabalhar com propriedades de classe - particularmente, torna difícil extrair o lr a partir de um tf.Tensor() em qualquer outro tipo para acompanhar em uma lista. Isso pode ser contornado tecnicamente executando-se no modo ansioso, mas apresenta um incômodo para acompanhar o LR para fins de depuração e é melhor evitar:

class WarmUpCosineDecay(keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, start_lr, target_lr, warmup_steps, total_steps, hold):
        super().__init__()
        self.start_lr = start_lr
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.hold = hold

    def __call__(self, step):
        lr = lr_warmup_cosine_decay(global_step=step,
                                    total_steps=self.total_steps,
                                    warmup_steps=self.warmup_steps,
                                    start_lr=self.start_lr,
                                    target_lr=self.target_lr,
                                    hold=self.hold)

        return tf.where(
            step > self.total_steps, 0.0, lr, name="learning_rate"
        )

Os parâmetros são os mesmos e podem ser calculados da mesma maneira que antes:

total_steps = len(train_set)*config['EPOCHS']



warmup_steps = int(0.05*total_steps)

schedule = WarmUpCosineDecay(start_lr=0.0, target_lr=1e-3, warmup_steps=warmup_steps, total_steps=total_steps, hold=warmup_steps)

E o pipeline de treinamento só difere porque definimos o LR do otimizador para o schedule:

model = keras.applications.EfficientNetV2B0(weights=None, 
                                            classes=n_classes, 
                                            input_shape=[224, 224, 3])
  
model.compile(loss="sparse_categorical_crossentropy",
                  optimizer=tf.keras.optimizers.Adam(learning_rate=schedule),
                  jit_compile=True,
                  metrics=['accuracy'])

history3 = model.fit(train_set,
                    epochs = config['EPOCHS'],
                    validation_data=valid_set)

Se desejar salvar o modelo, o WarmupCosineDecay programação terá que substituir o get_config() método:

    def get_config(self):
        config = {
          'start_lr': self.start_lr,
          'target_lr': self.target_lr,
          'warmup_steps': self.warmup_steps,
          'total_steps': self.total_steps,
          'hold': self.hold
        }
        return config

Finalmente, ao carregar o modelo, você terá que passar um WarmupCosineDecay como um objeto personalizado:

model = keras.models.load_model('weights.h5', 
                                custom_objects={'WarmupCosineDecay', WarmupCosineDecay})

Conclusão

Neste guia, demos uma olhada na intuição por trás do aquecimento da taxa de aprendizagem – uma técnica comum para manipular a taxa de aprendizagem durante o treinamento de redes neurais.

Implementamos um aquecimento da taxa de aprendizagem com decaimento de cosseno, o tipo mais comum de redução de LR emparelhado com aquecimento. Você pode implementar qualquer outra função para redução ou não reduzir a taxa de aprendizado – deixando isso para outros retornos de chamada, como ReduceLROnPlateau(). Implementamos o aquecimento da taxa de aprendizagem como um retorno de chamada Keras, bem como um cronograma Keras Optimizer e traçamos a taxa de aprendizagem ao longo das épocas.