Program Listing for File lr_scheduler.h

↰ Return to documentation for file (include/runtime/include/api/callback/lr_scheduler.h)

#ifndef MINDSPORE_INCLUDE_API_CALLBACK_LR_SCHEDULER_H
#define MINDSPORE_INCLUDE_API_CALLBACK_LR_SCHEDULER_H

#include <cstddef>
#include <string>
#include <vector>
#include <memory>
#include "include/api/callback/callback.h"

namespace mindspore {

constexpr int DONT_UPDATE_LR = 0;
constexpr int UPDATE_LR = 1;

using LR_Lambda = std::function<int(float *lr, int epoch, void *cb_data)>;

MS_API int MultiplicativeLRLambda(float *lr, int epoch, void *multiplication);

MS_API int StepLRLambda(float *lr, int epoch, void *step_size);
struct MS_API StepLRLambda {
  StepLRLambda(int step, float g) : step_size(step), gamma(g) {}

  int step_size;  // period of LR decay
  float gamma;    // LR decay factor
};

class MS_API LRScheduler : public TrainCallBack {
 public:
  explicit LRScheduler(LR_Lambda lambda_func, void *lr_cb_data = nullptr, int step = 1);
  virtual ~LRScheduler();
};

}  // namespace mindspore
#endif  // MINDSPORE_INCLUDE_API_CALLBACK_LR_SCHEDULER_H