| 35 | |
| 36 | |
| 37 | class LearningRate: |
| 38 | def __init__(self, cfg): |
| 39 | self.steps = cfg["multi_step"] |
| 40 | self.current_step = 0 |
| 41 | |
| 42 | def get_lr(self, iteration): |
| 43 | lr = self.steps[self.current_step][0] |
| 44 | if iteration == self.steps[self.current_step][1]: |
| 45 | self.current_step += 1 |
| 46 | |
| 47 | return lr |
| 48 | |
| 49 | |
| 50 | def get_batch_spec(cfg): |