add comments and modify variable names

src/brain/modelgen/AugmentedNN.py

@@ -29,13 +29,13 @@ def decorator(self):
 
 class AugmentedNN:
 
-    def __init__(self, ncol, order=1, nneuron=16, lr=0.1, **kwargs):
+    def __init__(self, column_num, order=1, neuron_num=16, lr=0.1, **kwargs):
         tf.reset_default_graph()
-        self.data = tf.placeholder(tf.float32, [None, ncol*2], name="data_")
+        self.data = tf.placeholder(tf.float32, [None, column_num*2], name="data_")
         self.target =  tf.placeholder(tf.float32, [None, 1], name="target_")
-        self._ncol = ncol
+        self._column_num = column_num
         self._order = order
-        self._nneuron = nneuron
+        self._neuron_num = neuron_num
         self._lr = tf.placeholder_with_default(lr, shape=None,
                                                name="learn_rate_")
         self.tf_init = tf.global_variables_initializer
@@ -44,10 +44,14 @@ def __init__(self, ncol, order=1, nneuron=16, lr=0.1, **kwargs):
         self.optimize
 
     @staticmethod
-    def jumpActivation(k):
-        def jumpActivationk(x):
+    def jump_activation(k):
+        """
+        This is an activation function used to learn discontinuous functions.
+        Reference: https://dl.acm.org/citation.cfm?id=2326898
+        """
+        def jump_activation_k(x):
             return tf.pow(tf.maximum(0.0, 1-tf.exp(-x)), k)
-        return jumpActivationk
+        return jump_activation_k
 
     @lazy_property
     def prediction(self):
@@ -59,14 +63,14 @@ def prediction(self):
 
             h1_layers = []
             for i in range(1, self._order+1):
-                h1 = tf.layers.dense(net, self._nneuron, 
-                                     activation=self.jumpActivation(i),
+                h1 = tf.layers.dense(net, self._neuron_num, 
+                                     activation=self.jump_activation(i),
                                      kernel_initializer=kernel_init)
                 h1_layers.append(h1)
             h1_layers = tf.concat(h1_layers, 1)        
         with tf.name_scope("output_layer"):
             net = tf.layers.dense(h1_layers, 1, 
-                                  activation=self.jumpActivation(1),
+                                  activation=self.jump_activation(1),
                                   kernel_initializer=kernel_init)
         net = tf.reshape(net, [bsz, -1], name="pred_")
         return net
@@ -91,20 +95,19 @@ def write_graph(self, dir):
             tf.train.write_graph(tf.get_default_graph(),
                                  dir, fname, False)
 
-
     def __repr__(self):
         return "AugmentedNN"
 
 def main():
     parser = argparse.ArgumentParser(description='AugmentedNN Model Generator')
 
-    parser.add_argument('--ncol', type=int, default=1, help='Number of augmentedNN Hidden units')
-    parser.add_argument('--order', type=int, default=4, help='Order of activation function')
-    parser.add_argument('--nneuron', type=int, default=20, help='Number of neurons in hidden layer')
+    parser.add_argument('--column_num', type=int, default=1, help='Number of augmentedNN Hidden units')
+    parser.add_argument('--order', type=int, default=3, help='Max order of activation function')
+    parser.add_argument('--neuron_num', type=int, default=20, help='Number of neurons in hidden layer')
     parser.add_argument('--lr', type=float, default=0.001, help='Learning rate')
     parser.add_argument('graph_out_path', type=str, help='Path to write graph output', nargs='+')
     args = parser.parse_args()
-    model = AugmentedNN(args.ncol, args.order, args.nneuron, args.lr)
+    model = AugmentedNN(args.column_num, args.order, args.neuron_num, args.lr)
     model.tf_init()
     model.write_graph(' '.join(args.graph_out_path))
 

src/brain/workload/augmentedNN.cpp

@@ -20,13 +20,13 @@
 namespace peloton {
 namespace brain {
 
-AugmentedNN::AugmentedNN(int ncol, int order, int nneuron,
+AugmentedNN::AugmentedNN(int column_num, int order, int neuron_num,
                             float learn_rate, int batch_size, int epochs)
     : BaseTFModel("src/brain/modelgen", "src/brain/modelgen/AugmentedNN.py",
                   "src/brain/modelgen/AugmentedNN.pb"),
-      ncol_(ncol),
+      column_num_(column_num),
       order_(order),
-      nneuron_(nneuron),
+      neuron_num_(neuron_num),
       learn_rate_(learn_rate),
       batch_size_(batch_size),
       epochs_(epochs) {
@@ -39,9 +39,9 @@ AugmentedNN::AugmentedNN(int ncol, int order, int nneuron,
 
 std::string AugmentedNN::ConstructModelArgsString() const {
   std::stringstream args_str_builder;
-  args_str_builder << " --ncol " << ncol_;
+  args_str_builder << " --column_num " << column_num_;
   args_str_builder << " --order " << order_;
-  args_str_builder << " --nneuron " << nneuron_;
+  args_str_builder << " --neuron_num " << neuron_num_;
   args_str_builder << " --lr " << learn_rate_;
   args_str_builder << " " << this->modelgen_path_;
   return args_str_builder.str();
@@ -50,9 +50,9 @@ std::string AugmentedNN::ConstructModelArgsString() const {
 std::string AugmentedNN::ToString() const {
   std::stringstream model_str_builder;
   model_str_builder << "AugmentedNN(";
-  model_str_builder << "ncol = " << ncol_;
+  model_str_builder << "column_num = " << column_num_;
   model_str_builder << ", order = " << order_;
-  model_str_builder << ", nneuron = " << nneuron_;
+  model_str_builder << ", neuron_num = " << neuron_num_;
   model_str_builder << ", lr = " << learn_rate_;
   model_str_builder << ", batch_size = " << batch_size_;
   model_str_builder << ")";
@@ -86,7 +86,7 @@ float AugmentedNN::TrainEpoch(const matrix_eig &mat) {
   std::vector<float> losses;
   // Obtain relevant metadata
   int min_allowed_bsz = 1;
-  int bsz = std::max(batch_size_, min_allowed_bsz);
+  int bsz = std::min((int)mat.rows(), std::max(batch_size_, min_allowed_bsz));
   int number_of_batches = mat.rows() / bsz;
   int num_cols = mat.cols() - 1;
 
@@ -138,7 +138,7 @@ matrix_eig AugmentedNN::Predict(const matrix_eig &X, int bsz) const {
 float AugmentedNN::ValidateEpoch(const matrix_eig &mat) {
   // Obtain relevant metadata
   int min_allowed_bsz = 1;
-  int bsz = std::max(batch_size_, min_allowed_bsz);
+  int bsz = std::min((int)mat.rows(), std::max(batch_size_, min_allowed_bsz));
   int number_of_batches = mat.rows() / bsz;
   int num_cols = mat.cols() - 1;
 

src/include/brain/workload/augmentedNN.h

@@ -26,15 +26,24 @@ class TfSessionEntityOutput;
 using TfFloatIn = TfSessionEntityInput<float>;
 using TfFloatOut = TfSessionEntityOutput<float>;
 
+/**
+ * AugmentedNN is used to predict selectivity or cardinality
+ * for range predicates, e.g, 
+ *   SELECT * FROM table WHERE c1 >= l1 AND c1 <= u1
+ *                         AND c2 >= l2 AND c2 <= u2
+ *                         AND ... 
+ * Input is [l1, u1, l2, u2, ...]
+ * Output is selectivity or cardinality
+ */
 class AugmentedNN : public BaseTFModel {
  public:
   AugmentedNN(int ncol, int order, int nneuron,
               float learn_rate, int batch_size, int epochs);
   /**
    * Train the Tensorflow model
-   * @param mat: Contiguous time-series data
+   * @param mat: Training data
    * @return: Average Training loss
-   * Given a continuous sequence of data,
+   * Given a matrix of training data,
    * this function:
    * 1. breaks the data into batches('Batchify')
    * 2. prepares tensorflow-entity inputs/outputs
@@ -45,19 +54,27 @@ class AugmentedNN : public BaseTFModel {
   float TrainEpoch(const matrix_eig &mat);
 
   /**
-   * @param mat: Contiguous time-series data
-   * @param test_true: reference variable to which true values are returned
-   * @param test_pred: reference variable to which predicted values are returned
-   * @param return_preds: Whether to return predicted/true values
+   * @param mat: Validating data
    * @return: Average Validation Loss
    * This applies the same set of steps as TrainEpoch.
    * However instead of applying backprop it obtains predicted values.
-   * Then the validation loss is calculated for the relevant sequence
-   * - this is a function of segment and horizon.
+   * Then the validation loss is calculated.
    */
   float ValidateEpoch(const matrix_eig &mat); 
 
+  /**
+   * @param X: Input for the model
+   * @param y: Ground truth output corresponding to X
+   * @param bsz: Batchsize
+   * This function applies backprop once.
+   */ 
   void Fit(const matrix_eig &X, const matrix_eig &y, int bsz) override;
+
+  /**
+   * @param X: Input for the model
+   * @param bsz: Batchsize
+   * This function gets prediction for the input from the model.
+   */
   matrix_eig Predict(const matrix_eig &X, int bsz) const override;
 
   /**
@@ -75,10 +92,14 @@ class AugmentedNN : public BaseTFModel {
                 matrix_eig &data, matrix_eig &target);
   // Function to generate the args string to feed the python model
   std::string ConstructModelArgsString() const;
-  // Attributes needed for the Seq2Seq LSTM model(set by the user/settings.json)
-  int ncol_;
+
+  // Attributes needed for the AugmentedNN model
+  // number of columns used as input
+  int column_num_;
+  // max order for activation function
   int order_;
-  int nneuron_;
+  // number of neurons in hidden layer
+  int neuron_num_;
   float learn_rate_;
   int batch_size_;
   int epochs_;