Warmup-decay learning rate for PanguAlpha network.
| 230 | |
| 231 | |
| 232 | class LearningRate(LearningRateSchedule): |
| 233 | """ |
| 234 | Warmup-decay learning rate for PanguAlpha network. |
| 235 | """ |
| 236 | |
| 237 | def __init__(self, |
| 238 | learning_rate, |
| 239 | end_learning_rate, |
| 240 | warmup_steps, |
| 241 | decay_steps, |
| 242 | power=1.0, |
| 243 | use_cosine=True): |
| 244 | super(LearningRate, self).__init__() |
| 245 | self.warmup_flag = False |
| 246 | if warmup_steps > 0: |
| 247 | self.warmup_flag = True |
| 248 | self.warmup_lr = WarmUpLR(learning_rate, warmup_steps) |
| 249 | self.decay_lr = PolynomialDecayLR(learning_rate, end_learning_rate, |
| 250 | decay_steps, power) |
| 251 | self.cosine_decay_lr = CosineDecayLR(end_learning_rate, learning_rate, |
| 252 | decay_steps) |
| 253 | self.warmup_steps = Tensor(np.array([warmup_steps]).astype(np.float32)) |
| 254 | |
| 255 | self.greater = P.Greater() |
| 256 | self.one = Tensor(np.array([1.0]).astype(np.float32)) |
| 257 | self.cast = P.Cast() |
| 258 | self.print = P.Print() |
| 259 | self.use_cosine = use_cosine |
| 260 | |
| 261 | def construct(self, global_step): |
| 262 | """dynamic learning rate""" |
| 263 | if not self.use_cosine: |
| 264 | decay_lr = self.decay_lr(global_step) |
| 265 | else: |
| 266 | decay_lr = self.cosine_decay_lr(global_step) |
| 267 | if self.warmup_flag: |
| 268 | is_warmup = self.cast(self.greater(self.warmup_steps, global_step), |
| 269 | mstype.float32) |
| 270 | warmup_lr = self.warmup_lr(global_step) |
| 271 | lr = (self.one - is_warmup) * decay_lr + is_warmup * warmup_lr |
| 272 | else: |
| 273 | lr = decay_lr |
| 274 | # self.print(f"Learning rate: {lr.asnumpy().tolist()}") |
| 275 | return lr |
| 276 | |
| 277 | |
| 278 | def add_inference_params(opt): |
no outgoing calls
no test coverage detected