• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2020 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #ifndef MINDSPORE_LITE_INCLUDE_TRAIN_LR_SCHEDULER_H_
17 #define MINDSPORE_LITE_INCLUDE_TRAIN_LR_SCHEDULER_H_
18 #include <vector>
19 #include <string>
20 #include <utility>
21 #include <functional>
22 #include <unordered_map>
23 #include "include/train/train_loop_callback.h"
24 
25 namespace mindspore {
26 namespace lite {
27 
28 constexpr int DONT_UPDATE_LR = 0;
29 constexpr int UPDATE_LR = 1;
30 
31 using LR_Lambda = std::function<int(float *lr, int epoch, void *cb_data)>;
32 
33 /// \brief Multiply the LR by a factor of gamma every epoch
34 int MultiplicativeLRLambda(float *lr, int epoch, void *multiplication);
35 
36 /// \brief Multiply the LR by a factor of gamma every step_size
37 int StepLRLambda(float *lr, int epoch, void *step_size);
38 struct StepLRLambda {
StepLRLambdaStepLRLambda39   StepLRLambda(int step, float g) : step_size(step), gamma(g) {}
40 
41   int step_size;  // period of LR decay
42   float gamma;    // LR decay factor
43 };
44 
45 class LRScheduler : public session::TrainLoopCallBack {
46  public:
47   explicit LRScheduler(LR_Lambda lambda_func, void *lr_cb_data = nullptr, int step_ = 1);
48   virtual ~LRScheduler() = default;
49   int EpochEnd(const session::TrainLoopCallBackData &cb_data) override;
50 
51  private:
52   LR_Lambda lambda_func_;
53   void *lr_data_ = nullptr;
54   int step_ = 1;
55 };
56 
57 }  // namespace lite
58 }  // namespace mindspore
59 #endif  // MINDSPORE_LITE_INCLUDE_TRAIN_LR_SCHEDULER_H_
60