syntax = "proto2";
package oneflow;

message ExponentialDecayConf {
  required int64 decay_batches = 1;
  required double decay_rate = 2;
  optional bool staircase = 3 [default = false];
}

message InverseTimeDecayConf {
  required int64 decay_batches = 1;
  required double decay_rate = 2;
  optional bool staircase = 3 [default = false];
}

message NaturalExpDecayConf {
  required int64 decay_batches = 1;
  required double decay_rate = 2;
  optional bool staircase = 3 [default = false];
}

message PiecewiseConstantConf {
  repeated int64 boundaries = 1;
  repeated double values = 2;
}

message PolynomialDecayConf {
  required int64 decay_batches = 1;
  optional double end_learning_rate = 2 [default = 0.0001];
  optional double power = 3 [default = 1.0];
  optional bool cycle = 4 [default = false];
}

message CosineDecayConf {
  required int64 decay_batches = 1;
  optional double alpha = 2 [default = 0.0];
}

message CosineAnnealingDecayConf {
  required int64 t_max = 1;
  optional double eta_min = 2 [default = 0.0];
}

message LinearCosineDecayConf {
  required int64 decay_batches = 1;
  optional double num_periods = 2 [default = 0.5];
  optional double alpha = 3 [default = 0.0];
  optional double beta = 4 [default = 0.001];
}

message PiecewiseScalingConf {
  repeated int64 boundaries = 1;
  repeated double scales = 2;
}

message StepConf {
  required int64 step_size = 1;
  optional double gamma = 2 [default = 0.1];
}

message MultiStepConf {
  repeated int64 milestones = 1;
  optional double gamma = 2 [default = 0.1];
}

message LinearLRConf {
  required double start_factor = 1;
  required double end_factor = 2;
  required int64 total_iters = 3;
}

message ConstantLRConf {
  required double factor = 1;
  required int64 total_iters = 2;
}

message CosineAnnealingWarmRestartsConf {
  required int64 t_initial = 1;
  required int64 t_mult = 2;
  required double eta_min = 3;
  required double decay_rate = 4;
  required int64 restart_limit = 5;
}

message SequentialSchedulerConf {
  repeated LearningRateDecayConf schedulers = 1;
  repeated int64 milestones = 2;
  // NOTE(zwx): should be repeated bool, however it has bug in cfg
  repeated int32 interval_rescaling = 3;
}

// TODO(zwx): ChainedSchedulerConf 

message LearningRateDecayConf {
  oneof type {
    ExponentialDecayConf exponential_conf = 2000;
    InverseTimeDecayConf inverse_time_conf = 2001;
    NaturalExpDecayConf natural_exp_conf = 2002;
    PiecewiseConstantConf piecewise_constant_conf = 2003;
    PolynomialDecayConf polynomial_conf = 2004;
    CosineDecayConf cosine_conf = 2005;
    LinearCosineDecayConf linear_cosine_conf = 2006;
    PiecewiseScalingConf piecewise_scaling_conf = 2007;
    MultiStepConf multi_step_conf = 2008;
    StepConf step_conf = 2009;
    CosineAnnealingDecayConf cosine_annealing_conf = 2010;
    LinearLRConf linear_lr_conf = 2011;
    ConstantLRConf constant_lr_conf = 2012;
    CosineAnnealingWarmRestartsConf cosine_annealing_warm_restarts_conf = 2013;
    SequentialSchedulerConf sequential_scheduler_conf = 2014;
  }
}
