Merge pull request #476 from richardsliu/hp_tuning

Fix xgboost example for hyperparameter tuning
kubeflow · Jan 15, 2019 · 64c3889 · 64c3889
2 parents c83ed09 + 3859564
commit 64c3889
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/xgboost_ames_housing/housing.py b/xgboost_ames_housing/housing.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import logging
 import joblib
 import pandas as pd
 from sklearn.metrics import mean_absolute_error
@@ -30,7 +31,8 @@ def read_input(file_name, test_size=0.25):
 
   train_X, test_X, train_y, test_y = train_test_split(X.values,
                                                       y.values,
-                                                      test_size=test_size)
+                                                      test_size=test_size,
+                                                      shuffle=False)
 
   imputer = Imputer()
   train_X = imputer.fit_transform(train_X)
@@ -53,20 +55,20 @@ def train_model(train_X,
             early_stopping_rounds=40,
             eval_set=[(test_X, test_y)])
 
-  print("Best RMSE on eval: {:.2f} with {} rounds".format(
-                 model.best_score,
-                 model.best_iteration+1))
+  logging.info("Best RMSE on eval: %.2f with %d rounds",
+               model.best_score,
+               model.best_iteration+1)
   return model
 
 def eval_model(model, test_X, test_y):
   """Evaluate the model performance."""
   predictions = model.predict(test_X)
-  print("MAE on test: {:.2f}".format(mean_absolute_error(predictions, test_y)))
+  logging.info("mean_absolute_error=%.2f", mean_absolute_error(predictions, test_y))
 
 def save_model(model, model_file):
   """Save XGBoost model for serving."""
   joblib.dump(model, model_file)
-  print("Model export success {}".format(model_file))
+  logging.info("Model export success: %s", model_file)
 
 def main(args):
   (train_X, train_y), (test_X, test_y) = read_input(args.train_input)
@@ -115,5 +117,7 @@ def main(args):
           default=50
           )
 
+  logging.basicConfig(format='%(message)s')
+  logging.getLogger().setLevel(logging.INFO)
   main_args = parser.parse_args()
   main(main_args)