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(lr_init, lr_end, lr_max, warmup_epochs, total_epochs, steps_per_epoch, lr_decay_mode): 22 """ 23 generate learning rate array 24 25 Args: 26 lr_init(float): init learning rate 27 lr_end(float): end learning rate 28 lr_max(float): max learning rate 29 warmup_epochs(int): number of warmup epochs 30 total_epochs(int): total epoch of training 31 steps_per_epoch(int): steps of one epoch 32 lr_decay_mode(string): learning rate decay mode, including steps, poly, cosine or default 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 if lr_decay_mode == 'steps': 41 decay_epoch_index = [0.3 * total_steps, 42 0.6 * total_steps, 0.8 * total_steps] 43 for i in range(total_steps): 44 if i < decay_epoch_index[0]: 45 lr = lr_max 46 elif i < decay_epoch_index[1]: 47 lr = lr_max * 0.1 48 elif i < decay_epoch_index[2]: 49 lr = lr_max * 0.01 50 else: 51 lr = lr_max * 0.001 52 lr_each_step.append(lr) 53 elif lr_decay_mode == 'poly': 54 if warmup_steps != 0: 55 inc_each_step = (float(lr_max) - float(lr_init)) / \ 56 float(warmup_steps) 57 else: 58 inc_each_step = 0 59 for i in range(total_steps): 60 if i < warmup_steps: 61 lr = float(lr_init) + inc_each_step * float(i) 62 else: 63 base = (1.0 - (float(i) - float(warmup_steps)) / 64 (float(total_steps) - float(warmup_steps))) 65 lr = float(lr_max) * base * base 66 if lr < 0.0: 67 lr = 0.0 68 lr_each_step.append(lr) 69 elif lr_decay_mode == 'cosine': 70 decay_steps = total_steps - warmup_steps 71 for i in range(total_steps): 72 if i < warmup_steps: 73 lr_inc = (float(lr_max) - float(lr_init)) / float(warmup_steps) 74 lr = float(lr_init) + lr_inc * (i + 1) 75 else: 76 linear_decay = (total_steps - i) / decay_steps 77 cosine_decay = 0.5 * \ 78 (1 + math.cos(math.pi * 2 * 0.47 * i / decay_steps)) 79 decayed = linear_decay * cosine_decay + 0.00001 80 lr = lr_max * decayed 81 lr_each_step.append(lr) 82 else: 83 for i in range(total_steps): 84 if i < warmup_steps: 85 lr = lr_init + (lr_max - lr_init) * i / warmup_steps 86 else: 87 lr = lr_max - (lr_max - lr_end) * \ 88 (i - warmup_steps) / (total_steps - warmup_steps) 89 lr_each_step.append(lr) 90 91 learning_rate = np.array(lr_each_step).astype(np.float32) 92 93 return learning_rate 94