1# Copyright 2020 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15"""learning rate generator""" 16 17import math 18import numpy as np 19 20 21def get_lr(global_step, lr_init, lr_end, lr_max, warmup_epochs, total_epochs, steps_per_epoch): 22 """ 23 generate learning rate array 24 25 Args: 26 global_step(int): total steps of the training 27 lr_init(float): init learning rate 28 lr_end(float): end learning rate 29 lr_max(float): max learning rate 30 warmup_epochs(int): number of warmup epochs 31 total_epochs(int): total epoch of training 32 steps_per_epoch(int): steps of one epoch 33 34 Returns: 35 np.array, learning rate array 36 """ 37 lr_each_step = [] 38 total_steps = steps_per_epoch * total_epochs 39 warmup_steps = steps_per_epoch * warmup_epochs 40 for i in range(total_steps): 41 if i < warmup_steps: 42 lr = lr_init + (lr_max - lr_init) * i / warmup_steps 43 else: 44 lr = lr_end + \ 45 (lr_max - lr_end) * \ 46 (1. + math.cos(math.pi * (i - warmup_steps) / 47 (total_steps - warmup_steps))) / 2. 48 if lr < 0.0: 49 lr = 0.0 50 lr_each_step.append(lr) 51 52 current_step = global_step 53 lr_each_step = np.array(lr_each_step).astype(np.float32) 54 learning_rate = lr_each_step[current_step:] 55 56 return learning_rate 57