syntax = "proto2";
package oneflow;

import "oneflow/core/common/shape.proto";
import "oneflow/core/common/data_type.proto";

message ConstantFillConf {
  optional float value = 1 [default = 0];
}

message UniformFillConf {
  optional float min = 1 [default = 0];
  optional float max = 2 [default = 1];
}

message GaussianFillConf {
  optional float mean = 1 [default = 0];
  optional float std = 2 [default = 1];
}

message FillConf {
  oneof type {
    ConstantFillConf constant_conf = 1;
    UniformFillConf uniform_conf = 2;
    GaussianFillConf gaussian_conf = 3;
  }
}


message ConvolutionOpConf {
  required string in = 1;
  required string out = 2;

  optional int32 out_num = 3 [default = 1];
  optional bool has_bias_term = 4 [default = true];

  optional int32 pad_h = 5 [default = 0];
  optional int32 pad_w = 6 [default = 0];
  required int32 kernel_h = 7;
  required int32 kernel_w = 8;
  optional int32 stride_h = 9 [default = 1];
  optional int32 stride_w = 10 [default = 1];
  optional int32 dilation_h = 11 [default = 1];
  optional int32 dilation_w = 12 [default = 1];

  optional FillConf weight_fill = 14;
  optional FillConf bias_fill = 15;
}

message InnerProductOpConf {
  required string in = 1;
  required string out = 2;

  required int32 out_num = 3;
  optional bool has_bias_term = 4 [default = true];
  optional FillConf weight_fill = 5; 
  optional FillConf bias_fill = 6; 
}

message DataLoaderOpConf {
  required string out = 1;
  required DataType data_type = 2;
  required ShapeProto shape = 3;
  required string data_dir = 4;
}

message PoolingOpConf {
  required string in = 1;
  required string out = 2;

  enum PoolMethod {
    kMax = 0;
    kAve = 1;
    kStochastic = 2;
  }
  required PoolMethod pool = 3;
  optional int32 pad_h = 4 [default = 0];
  optional int32 pad_w = 5 [default = 0];
  required int32 kernel_h = 6;
  required int32 kernel_w = 7;
  optional int32 stride_h = 8 [default = 1];
  optional int32 stride_w = 9 [default = 1];
}

message ReluOpConf {
  required string in = 1;
  required string out = 2;
}

message SoftmaxOpConf {
  required string in = 1;
  required string out = 2;
}

message SoftmaxLossOpConf {
  required string prediction = 1;
  required string label = 3;
  required string loss = 4;
}

message MultinomialLogisticLossOpConf {
  required string prediction = 1;
  required string label = 2;
  required string loss = 3;
}

message ConcatOpConf {
  repeated string in = 1;
  required string out = 2;
  required int32 axis = 4;
}

message CopyCommNetOpConf {
}

message CopyHdOpConf {
  enum Type {
    H2D = 0;
    D2H = 1;
  }
  required Type type = 1;
}

message CloneOpConf {
  required int32 out_num = 1;
  required string lbn = 2;
  required DataType data_type = 3;
}

message BoxConcatConf {
  required int32 axis = 1;
}

message BoxAddConf {
}

message BoxSplitConf {
  required int32 axis = 1;
  // left_bound_size, (base_part_size + 1) * bigger_part_num, base_part_size * base_part_num, right_bound_size
  required int32 left_bound_size = 2;
  required int32 bigger_part_num = 3;
  required int32 base_part_size = 4;
  required int32 base_part_num = 5;
  required int32 right_bound_size = 6;
}

message BoxCloneConf {
}

message BoxingOpConf {
  required string lbn = 2;
  required int32 in_num = 3;
  required int32 out_num = 4;

  oneof in_box {
    BoxConcatConf concat_box = 5;
    BoxAddConf add_box = 6;
  }
  oneof out_box {
    BoxSplitConf split_box = 7;
    BoxCloneConf clone_box = 8;
  }
}

message NormalModelUpdateOpConf {
  optional float learning_rate = 1 [default = 0.01];
}

message MomentumModelUpdateOpConf {
  optional float learning_rate = 1 [default = 0.01];
  optional float beta = 2 [default = 0.9];
}

message RMSPropModelUpdateOpConf {
  optional float learning_rate = 1 [default = 0.01];
  optional float decay_rate = 2 [default = 0.99];
  optional float epsilon = 3 [default = 1e-8];
}

message AccumulateOpConf {
}

message ModelSaveOpConf {
  repeated string lbns = 1;
}

message RecordOpConf {
  repeated string lbn = 1;
  required string record_path = 2;
}

message LossRecordOpConf {
}

message OperatorConf {
  required string name = 1;
  optional string model_load_dir = 2;
  oneof op_type {
    ConvolutionOpConf convolution_conf = 100;
    InnerProductOpConf innerproduct_conf = 101;
    DataLoaderOpConf data_loader_conf = 102;
    PoolingOpConf pooling_conf = 103;
    ReluOpConf relu_conf = 104;
    SoftmaxOpConf softmax_conf = 105;
    MultinomialLogisticLossOpConf multinomial_logistic_loss_conf = 106;
    CopyHdOpConf copy_hd_conf = 107;
    CloneOpConf clone_conf = 108;
    BoxingOpConf boxing_conf = 109;
    NormalModelUpdateOpConf normal_mdupdt_conf = 110;
    ModelSaveOpConf model_save_conf = 111;
    AccumulateOpConf accumulate_conf = 112;
    ConcatOpConf concat_conf = 113;
    CopyCommNetOpConf copy_comm_net_conf = 114;
    MomentumModelUpdateOpConf momentum_mdupdt_conf = 115;
    RMSPropModelUpdateOpConf rmsprop_mdupdt_conf = 116;
    SoftmaxLossOpConf softmax_loss_conf = 117;
    RecordOpConf record_conf = 118;
    LossRecordOpConf loss_record_conf = 119;
  }
}