JSON configuration IO. (#5111)

* Add saving/loading JSON configuration. * Implement Python pickle interface with new IO routines. * Basic tests for training continuation.
2019-12-15 17:31:53 +08:00
parent 5aa007d7b2
commit 3136185bc5
24 changed files with 761 additions and 390 deletions
--- a/src/c_api/c_api.cc
+++ b/src/c_api/c_api.cc
@@ -458,8 +458,8 @@ XGB_DLL int XGDMatrixNumCol(const DMatrixHandle handle,

 // xgboost implementation
 XGB_DLL int XGBoosterCreate(const DMatrixHandle dmats[],
-                    xgboost::bst_ulong len,
-                    BoosterHandle *out) {
+                            xgboost::bst_ulong len,
+                            BoosterHandle *out) {
  API_BEGIN();
  std::vector<std::shared_ptr<DMatrix> > mats;
  for (xgboost::bst_ulong i = 0; i < len; ++i) {
@@ -485,6 +485,31 @@ XGB_DLL int XGBoosterSetParam(BoosterHandle handle,
  API_END();
 }

+XGB_DLL int XGBoosterLoadJsonConfig(BoosterHandle handle, char const* json_parameters) {
+  API_BEGIN();
+  CHECK_HANDLE();
+  std::string str {json_parameters};
+  Json config { Json::Load(StringView{str.c_str(), str.size()}) };
+  static_cast<Learner*>(handle)->LoadConfig(config);
+  API_END();
+}
+
+XGB_DLL int XGBoosterSaveJsonConfig(BoosterHandle handle,
+                                    xgboost::bst_ulong *out_len,
+                                    char const** out_str) {
+  API_BEGIN();
+  CHECK_HANDLE();
+  Json config { Object() };
+  auto* learner = static_cast<Learner*>(handle);
+  learner->Configure();
+  learner->SaveConfig(&config);
+  std::string& raw_str = XGBAPIThreadLocalStore::Get()->ret_str;
+  Json::Dump(config, &raw_str);
+  *out_str = raw_str.c_str();
+  *out_len = static_cast<xgboost::bst_ulong>(raw_str.length());
+  API_END();
+}
+
 XGB_DLL int XGBoosterUpdateOneIter(BoosterHandle handle,
                                   int iter,
                                   DMatrixHandle dtrain) {
@@ -579,7 +604,7 @@ XGB_DLL int XGBoosterLoadModel(BoosterHandle handle, const char* fname) {
    static_cast<Learner*>(handle)->LoadModel(in);
  } else {
    std::unique_ptr<dmlc::Stream> fi(dmlc::Stream::Create(fname, "r"));
-    static_cast<Learner*>(handle)->Load(fi.get());
+    static_cast<Learner*>(handle)->LoadModel(fi.get());
  }
  API_END();
 }
@@ -598,20 +623,18 @@ XGB_DLL int XGBoosterSaveModel(BoosterHandle handle, const char* c_fname) {
    fo->Write(str.c_str(), str.size());
  } else {
    auto *bst = static_cast<Learner*>(handle);
-    bst->Save(fo.get());
+    bst->SaveModel(fo.get());
  }
  API_END();
 }

-// The following two functions are `Load` and `Save` for memory based serialization
-// methods. E.g. Python pickle.
 XGB_DLL int XGBoosterLoadModelFromBuffer(BoosterHandle handle,
                                         const void* buf,
                                         xgboost::bst_ulong len) {
  API_BEGIN();
  CHECK_HANDLE();
  common::MemoryFixSizeBuffer fs((void*)buf, len);  // NOLINT(*)
-  static_cast<Learner*>(handle)->Load(&fs);
+  static_cast<Learner*>(handle)->LoadModel(&fs);
  API_END();
 }

@@ -621,6 +644,25 @@ XGB_DLL int XGBoosterGetModelRaw(BoosterHandle handle,
  std::string& raw_str = XGBAPIThreadLocalStore::Get()->ret_str;
  raw_str.resize(0);

+  API_BEGIN();
+  CHECK_HANDLE();
+  common::MemoryBufferStream fo(&raw_str);
+  auto *learner = static_cast<Learner*>(handle);
+  learner->Configure();
+  learner->SaveModel(&fo);
+  *out_dptr = dmlc::BeginPtr(raw_str);
+  *out_len = static_cast<xgboost::bst_ulong>(raw_str.length());
+  API_END();
+}
+
+// The following two functions are `Load` and `Save` for memory based
+// serialization methods. E.g. Python pickle.
+XGB_DLL int XGBoosterSerializeToBuffer(BoosterHandle handle,
+                                       xgboost::bst_ulong *out_len,
+                                       const char **out_dptr) {
+  std::string &raw_str = XGBAPIThreadLocalStore::Get()->ret_str;
+  raw_str.resize(0);
+
  API_BEGIN();
  CHECK_HANDLE();
  common::MemoryBufferStream fo(&raw_str);
@@ -632,6 +674,41 @@ XGB_DLL int XGBoosterGetModelRaw(BoosterHandle handle,
  API_END();
 }

+XGB_DLL int XGBoosterUnserializeFromBuffer(BoosterHandle handle,
+                                           const void *buf,
+                                           xgboost::bst_ulong len) {
+  API_BEGIN();
+  CHECK_HANDLE();
+  common::MemoryFixSizeBuffer fs((void*)buf, len);  // NOLINT(*)
+  static_cast<Learner*>(handle)->Load(&fs);
+  API_END();
+}
+
+XGB_DLL int XGBoosterLoadRabitCheckpoint(BoosterHandle handle,
+                                         int* version) {
+  API_BEGIN();
+  CHECK_HANDLE();
+  auto* bst = static_cast<Learner*>(handle);
+  *version = rabit::LoadCheckPoint(bst);
+  if (*version != 0) {
+    bst->Configure();
+  }
+  API_END();
+}
+
+XGB_DLL int XGBoosterSaveRabitCheckpoint(BoosterHandle handle) {
+  API_BEGIN();
+  CHECK_HANDLE();
+  auto* learner = static_cast<Learner*>(handle);
+  learner->Configure();
+  if (learner->AllowLazyCheckPoint()) {
+    rabit::LazyCheckPoint(learner);
+  } else {
+    rabit::CheckPoint(learner);
+  }
+  API_END();
+}
+
 inline void XGBoostDumpModelImpl(
    BoosterHandle handle,
    const FeatureMap& fmap,
@@ -758,29 +835,5 @@ XGB_DLL int XGBoosterGetAttrNames(BoosterHandle handle,
  API_END();
 }

-XGB_DLL int XGBoosterLoadRabitCheckpoint(BoosterHandle handle,
-                                         int* version) {
-  API_BEGIN();
-  CHECK_HANDLE();
-  auto* bst = static_cast<Learner*>(handle);
-  *version = rabit::LoadCheckPoint(bst);
-  if (*version != 0) {
-    bst->Configure();
-  }
-  API_END();
-}
-
-XGB_DLL int XGBoosterSaveRabitCheckpoint(BoosterHandle handle) {
-  API_BEGIN();
-  CHECK_HANDLE();
-  auto* bst = static_cast<Learner*>(handle);
-  if (bst->AllowLazyCheckPoint()) {
-    rabit::LazyCheckPoint(bst);
-  } else {
-    rabit::CheckPoint(bst);
-  }
-  API_END();
-}
-
 // force link rabit
 static DMLC_ATTRIBUTE_UNUSED int XGBOOST_LINK_RABIT_C_API_ = RabitLinkTag();
--- a/src/gbm/gblinear.cc
+++ b/src/gbm/gblinear.cc
@@ -99,6 +99,16 @@ class GBLinear : public GradientBooster {
    model_.LoadModel(model);
  }

+  void LoadConfig(Json const& in) override {
+    CHECK_EQ(get<String>(in["name"]), "gblinear");
+    fromJson(in["gblinear_train_param"], &param_);
+  }
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String{"gblinear"};
+    out["gblinear_train_param"] = toJson(param_);
+  }
+
  void DoBoost(DMatrix *p_fmat,
               HostDeviceVector<GradientPair> *in_gpair,
               ObjFunction* obj) override {
--- a/src/gbm/gblinear_model.h
+++ b/src/gbm/gblinear_model.h
@@ -112,7 +112,8 @@ class GBLinearModel : public Model {
         << "    \"weight\": [" << std::endl;
      for (unsigned i = 0; i < nfeature; ++i) {
        for (int gid = 0; gid < ngroup; ++gid) {
-          if (i != 0 || gid != 0) fo << "," << std::endl;
+          if (i != 0 || gid != 0)
+            fo << "," << std::endl;
          fo << "      " << (*this)[i][gid];
        }
      }
@@ -134,5 +135,6 @@ class GBLinearModel : public Model {
    return v;
  }
 };
+
 }  // namespace gbm
 }  // namespace xgboost
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -34,6 +34,7 @@ DMLC_REGISTRY_FILE_TAG(gbtree);

 void GBTree::Configure(const Args& cfg) {
  this->cfg_ = cfg;
+  std::string updater_seq = tparam_.updater_seq;
  tparam_.UpdateAllowUnknown(cfg);

  model_.Configure(cfg);
@@ -75,24 +76,31 @@ void GBTree::Configure(const Args& cfg) {
        "`tree_method` parameter instead.";
    // Don't drive users to silent XGBOost.
    showed_updater_warning_ = true;
-  } else {
-    this->ConfigureUpdaters();
-    LOG(DEBUG) << "Using updaters: " << tparam_.updater_seq;
  }

-  for (auto& up : updaters_) {
-    up->Configure(cfg);
+  this->ConfigureUpdaters();
+  if (updater_seq != tparam_.updater_seq) {
+    updaters_.clear();
+    this->InitUpdater(cfg);
+  } else {
+    for (auto &up : updaters_) {
+      up->Configure(cfg);
+    }
  }

  configured_ = true;
 }

-// FIXME(trivialfis): This handles updaters and predictor.  Because the choice of updaters
-// depends on whether external memory is used and how large is dataset.  We can remove the
-// dependency on DMatrix once `hist` tree method can handle external memory so that we can
-// make it default.
+// FIXME(trivialfis): This handles updaters.  Because the choice of updaters depends on
+// whether external memory is used and how large is dataset.  We can remove the dependency
+// on DMatrix once `hist` tree method can handle external memory so that we can make it
+// default.
 void GBTree::ConfigureWithKnownData(Args const& cfg, DMatrix* fmat) {
+  CHECK(this->configured_);
  std::string updater_seq = tparam_.updater_seq;
+  CHECK(tparam_.GetInitialised());
+
+  tparam_.UpdateAllowUnknown(cfg);

  this->PerformTreeMethodHeuristic(fmat);
  this->ConfigureUpdaters();
@@ -101,9 +109,8 @@ void GBTree::ConfigureWithKnownData(Args const& cfg, DMatrix* fmat) {
  if (updater_seq != tparam_.updater_seq) {
    LOG(DEBUG) << "Using updaters: " << tparam_.updater_seq;
    this->updaters_.clear();
+    this->InitUpdater(cfg);
  }
-
-  this->InitUpdater(cfg);
 }

 void GBTree::PerformTreeMethodHeuristic(DMatrix* fmat) {
@@ -141,6 +148,9 @@ void GBTree::PerformTreeMethodHeuristic(DMatrix* fmat) {
 }

 void GBTree::ConfigureUpdaters() {
+  if (specified_updater_) {
+    return;
+  }
  // `updater` parameter was manually specified
  /* Choose updaters according to tree_method parameters */
  switch (tparam_.tree_method) {
@@ -289,6 +299,46 @@ void GBTree::CommitModel(std::vector<std::vector<std::unique_ptr<RegTree>>>&& ne
  monitor_.Stop("CommitModel");
 }

+void GBTree::LoadConfig(Json const& in) {
+  CHECK_EQ(get<String>(in["name"]), "gbtree");
+  fromJson(in["gbtree_train_param"], &tparam_);
+  int32_t const n_gpus = xgboost::common::AllVisibleGPUs();
+  if (n_gpus == 0 && tparam_.predictor == PredictorType::kGPUPredictor) {
+    tparam_.UpdateAllowUnknown(Args{{"predictor", "auto"}});
+  }
+  if (n_gpus == 0 && tparam_.tree_method == TreeMethod::kGPUHist) {
+    tparam_.UpdateAllowUnknown(Args{{"tree_method", "hist"}});
+    LOG(WARNING)
+        << "Loading from a raw memory buffer on CPU only machine.  "
+           "Change tree_method to hist.";
+  }
+
+  auto const& j_updaters = get<Object const>(in["updater"]);
+  updaters_.clear();
+  for (auto const& kv : j_updaters) {
+    std::unique_ptr<TreeUpdater> up(TreeUpdater::Create(kv.first, generic_param_));
+    up->LoadConfig(kv.second);
+    updaters_.push_back(std::move(up));
+  }
+
+  specified_updater_ = get<Boolean>(in["specified_updater"]);
+}
+
+void GBTree::SaveConfig(Json* p_out) const {
+  auto& out = *p_out;
+  out["name"] = String("gbtree");
+  out["gbtree_train_param"] = toJson(tparam_);
+  out["updater"] = Object();
+
+  auto& j_updaters = out["updater"];
+  for (auto const& up : updaters_) {
+    j_updaters[up->Name()] = Object();
+    auto& j_up = j_updaters[up->Name()];
+    up->SaveConfig(&j_up);
+  }
+  out["specified_updater"] = Boolean{specified_updater_};
+}
+
 void GBTree::LoadModel(Json const& in) {
  CHECK_EQ(get<String>(in["name"]), "gbtree");
  model_.LoadModel(in["model"]);
@@ -324,7 +374,7 @@ class Dart : public GBTree {
    for (size_t i = 0; i < weight_drop_.size(); ++i) {
      j_weight_drop[i] = Number(weight_drop_[i]);
    }
-    out["weight_drop"] = Array(j_weight_drop);
+    out["weight_drop"] = Array(std::move(j_weight_drop));
  }
  void LoadModel(Json const& in) override {
    CHECK_EQ(get<String>(in["name"]), "dart");
@@ -352,6 +402,21 @@ class Dart : public GBTree {
    }
  }

+  void LoadConfig(Json const& in) override {
+    CHECK_EQ(get<String>(in["name"]), "dart");
+    auto const& gbtree = in["gbtree"];
+    GBTree::LoadConfig(gbtree);
+    fromJson(in["dart_train_param"], &dparam_);
+  }
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String("dart");
+    out["gbtree"] = Object();
+    auto& gbtree = out["gbtree"];
+    GBTree::SaveConfig(&gbtree);
+    out["dart_train_param"] = toJson(dparam_);
+  }
+
  // predict the leaf scores with dropout if ntree_limit = 0
  void PredictBatch(DMatrix* p_fmat,
                    HostDeviceVector<bst_float>* out_preds,
--- a/src/gbm/gbtree.h
+++ b/src/gbm/gbtree.h
@@ -192,6 +192,9 @@ class GBTree : public GradientBooster {
    model_.Save(fo);
  }

+  void LoadConfig(Json const& in) override;
+  void SaveConfig(Json* p_out) const override;
+
  void SaveModel(Json* p_out) const override;
  void LoadModel(Json const& in) override;

--- a/src/gbm/gbtree_model.cc
+++ b/src/gbm/gbtree_model.cc
@@ -46,7 +46,8 @@ void GBTreeModel::SaveModel(Json* p_out) const {
  for (auto const& tree : trees) {
    Json tree_json{Object()};
    tree->SaveModel(&tree_json);
-    tree_json["id"] = std::to_string(t);
+    // The field is not used in XGBoost, but might be useful for external project.
+    tree_json["id"] = Integer(t);
    trees_json.emplace_back(tree_json);
    t++;
  }
--- a/src/learner.cc
+++ b/src/learner.cc
@@ -30,6 +30,7 @@

 #include "common/common.h"
 #include "common/io.h"
+#include "common/observer.h"
 #include "common/random.h"
 #include "common/timer.h"
 #include "common/version.h"
@@ -37,27 +38,6 @@
 namespace {

 const char* kMaxDeltaStepDefaultValue = "0.7";
-
-inline bool IsFloat(const std::string& str) {
-  std::stringstream ss(str);
-  float f{};
-  return !((ss >> std::noskipws >> f).rdstate() ^ std::ios_base::eofbit);
-}
-
-inline bool IsInt(const std::string& str) {
-  std::stringstream ss(str);
-  int i{};
-  return !((ss >> std::noskipws >> i).rdstate() ^ std::ios_base::eofbit);
-}
-
-inline std::string RenderParamVal(const std::string& str) {
-  if (IsFloat(str) || IsInt(str)) {
-    return str;
-  } else {
-    return std::string("'") + str + "'";
-  }
-}
-
 }  // anonymous namespace

 namespace xgboost {
@@ -323,11 +303,77 @@ class LearnerImpl : public Learner {
    }
  }

-  void Load(dmlc::Stream* fi) override {
+  void LoadConfig(Json const& in) override {
+    CHECK(IsA<Object>(in));
+    Version::Load(in, true);
+
+    auto const& learner_parameters = get<Object>(in["learner"]);
+    fromJson(learner_parameters.at("learner_train_param"), &tparam_);
+
+    auto const& gradient_booster = learner_parameters.at("gradient_booster");
+
+    auto const& objective_fn = learner_parameters.at("objective");
+    if (!obj_) {
+      obj_.reset(ObjFunction::Create(tparam_.objective, &generic_parameters_));
+    }
+    obj_->LoadConfig(objective_fn);
+
+    tparam_.booster = get<String>(gradient_booster["name"]);
+    if (!gbm_) {
+      gbm_.reset(GradientBooster::Create(tparam_.booster,
+                                         &generic_parameters_, &learner_model_param_,
+                                         cache_));
+    }
+    gbm_->LoadConfig(gradient_booster);
+
+    auto const& j_metrics = learner_parameters.at("metrics");
+    auto n_metrics = get<Array const>(j_metrics).size();
+    metric_names_.resize(n_metrics);
+    metrics_.resize(n_metrics);
+    for (size_t i = 0; i < n_metrics; ++i) {
+      metric_names_[i]= get<String>(j_metrics[i]);
+      metrics_[i] = std::unique_ptr<Metric>(
+          Metric::Create(metric_names_.back(), &generic_parameters_));
+    }
+
+    fromJson(learner_parameters.at("generic_param"), &generic_parameters_);
+
+    this->need_configuration_ = true;
+  }
+
+  void SaveConfig(Json* p_out) const override {
+    CHECK(!this->need_configuration_) << "Call Configure before saving model.";
+    Version::Save(p_out);
+    Json& out { *p_out };
+    // parameters
+    out["learner"] = Object();
+    auto& learner_parameters = out["learner"];
+
+    learner_parameters["learner_train_param"] = toJson(tparam_);
+    learner_parameters["gradient_booster"] = Object();
+    auto& gradient_booster = learner_parameters["gradient_booster"];
+    gbm_->SaveConfig(&gradient_booster);
+
+    learner_parameters["objective"] = Object();
+    auto& objective_fn = learner_parameters["objective"];
+    obj_->SaveConfig(&objective_fn);
+
+    std::vector<Json> metrics(metrics_.size());
+    for (size_t i = 0; i < metrics_.size(); ++i) {
+      metrics[i] = String(metrics_[i]->Name());
+    }
+    learner_parameters["metrics"] = Array(metrics);
+
+    learner_parameters["generic_param"] = toJson(generic_parameters_);
+  }
+
+  // About to be deprecated by JSON format
+  void LoadModel(dmlc::Stream* fi) override {
    generic_parameters_.UpdateAllowUnknown(Args{});
    tparam_.Init(std::vector<std::pair<std::string, std::string>>{});
    // TODO(tqchen) mark deprecation of old format.
    common::PeekableInStream fp(fi);
+
    // backward compatible header check.
    std::string header;
    header.resize(4);
@@ -338,6 +384,15 @@ class LearnerImpl : public Learner {
        CHECK_EQ(fp.Read(&header[0], 4), 4U);
      }
    }
+
+    if (header[0] == '{') {
+      auto json_stream = common::FixedSizeStream(&fp);
+      std::string buffer;
+      json_stream.Take(&buffer);
+      auto model = Json::Load({buffer.c_str(), buffer.size()});
+      this->LoadModel(model);
+      return;
+    }
    // use the peekable reader.
    fi = &fp;
    // read parameter
@@ -370,43 +425,9 @@ class LearnerImpl : public Learner {
      std::vector<std::pair<std::string, std::string> > attr;
      fi->Read(&attr);
      for (auto& kv : attr) {
-        // Load `predictor`, `gpu_id` parameters from extra attributes
        const std::string prefix = "SAVED_PARAM_";
        if (kv.first.find(prefix) == 0) {
          const std::string saved_param = kv.first.substr(prefix.length());
-          bool is_gpu_predictor = saved_param == "predictor" && kv.second == "gpu_predictor";
-#ifdef XGBOOST_USE_CUDA
-          if (saved_param == "predictor" || saved_param == "gpu_id") {
-            cfg_[saved_param] = kv.second;
-            LOG(INFO)
-              << "Parameter '" << saved_param << "' has been recovered from "
-              << "the saved model. It will be set to "
-              << RenderParamVal(kv.second) << " for prediction. To "
-              << "override the predictor behavior, explicitly set '"
-              << saved_param << "' parameter as follows:\n"
-              << "  * Python package: bst.set_param('"
-              << saved_param << "', [new value])\n"
-              << "  * R package:      xgb.parameters(bst) <- list("
-              << saved_param << " = [new value])\n"
-              << "  * JVM packages:   bst.setParam(\""
-              << saved_param << "\", [new value])";
-          }
-#else
-          if (is_gpu_predictor) {
-            cfg_["predictor"] = "cpu_predictor";
-            kv.second = "cpu_predictor";
-          }
-#endif  // XGBOOST_USE_CUDA
-#if defined(XGBOOST_USE_CUDA)
-          // NO visible GPU in current environment
-          if (is_gpu_predictor && common::AllVisibleGPUs() == 0) {
-            cfg_["predictor"] = "cpu_predictor";
-            kv.second = "cpu_predictor";
-            LOG(INFO) << "Switch gpu_predictor to cpu_predictor.";
-          } else if (is_gpu_predictor) {
-            cfg_["predictor"] = "gpu_predictor";
-          }
-#endif  // defined(XGBOOST_USE_CUDA)
          if (saved_configs_.find(saved_param) != saved_configs_.end()) {
            cfg_[saved_param] = kv.second;
          }
@@ -447,26 +468,12 @@ class LearnerImpl : public Learner {
      tparam_.dsplit = DataSplitMode::kRow;
    }

-    // There's no logic for state machine for binary IO, as it has a mix of everything and
-    // half loaded model.
    this->Configure();
  }

-  // rabit save model to rabit checkpoint
-  void Save(dmlc::Stream* fo) const override {
-    if (this->need_configuration_) {
-      // Save empty model.  Calling Configure in a dummy LearnerImpl avoids violating
-      // constness.
-      LearnerImpl empty(std::move(this->cache_));
-      empty.SetParams({this->cfg_.cbegin(), this->cfg_.cend()});
-      for (auto const& kv : attributes_) {
-        empty.SetAttr(kv.first, kv.second);
-      }
-      empty.Configure();
-      empty.Save(fo);
-      return;
-    }
-
+  // Save model into binary format.  The code is about to be deprecated by more robust
+  // JSON serialization format.
+  void SaveModel(dmlc::Stream* fo) const override {
    LearnerModelParamLegacy mparam = mparam_;  // make a copy to potentially modify
    std::vector<std::pair<std::string, std::string> > extra_attr;
    // extra attributed to be added just before saving
@@ -479,14 +486,13 @@ class LearnerImpl : public Learner {
      }
    }
    {
-      std::vector<std::string> saved_params{"predictor", "gpu_id"};
+      std::vector<std::string> saved_params;
      // check if rabit_bootstrap_cache were set to non zero before adding to checkpoint
      if (cfg_.find("rabit_bootstrap_cache") != cfg_.end() &&
        (cfg_.find("rabit_bootstrap_cache"))->second != "0") {
        std::copy(saved_configs_.begin(), saved_configs_.end(),
                  std::back_inserter(saved_params));
      }
-      // Write `predictor`, `n_gpus`, `gpu_id` parameters as extra attributes
      for (const auto& key : saved_params) {
        auto it = cfg_.find(key);
        if (it != cfg_.end()) {
@@ -495,19 +501,6 @@ class LearnerImpl : public Learner {
        }
      }
    }
-#if defined(XGBOOST_USE_CUDA)
-    {
-      // Force save gpu_id.
-      if (std::none_of(extra_attr.cbegin(), extra_attr.cend(),
-                   [](std::pair<std::string, std::string> const& it) {
-                     return it.first == "SAVED_PARAM_gpu_id";
-                   })) {
-        mparam.contain_extra_attrs = 1;
-        extra_attr.emplace_back("SAVED_PARAM_gpu_id",
-                                std::to_string(generic_parameters_.gpu_id));
-      }
-    }
-#endif  // defined(XGBOOST_USE_CUDA)
    fo->Write(&mparam, sizeof(LearnerModelParamLegacy));
    fo->Write(tparam_.objective);
    fo->Write(tparam_.booster);
@@ -541,6 +534,69 @@ class LearnerImpl : public Learner {
    }
  }

+  void Save(dmlc::Stream* fo) const override {
+    if (generic_parameters_.enable_experimental_json_serialization) {
+      Json memory_snapshot{Object()};
+      memory_snapshot["Model"] = Object();
+      auto &model = memory_snapshot["Model"];
+      this->SaveModel(&model);
+      memory_snapshot["Config"] = Object();
+      auto &config = memory_snapshot["Config"];
+      this->SaveConfig(&config);
+      std::string out_str;
+      Json::Dump(memory_snapshot, &out_str);
+      fo->Write(out_str.c_str(), out_str.size());
+    } else {
+      std::string binary_buf;
+      common::MemoryBufferStream s(&binary_buf);
+      this->SaveModel(&s);
+      Json config{ Object() };
+      // Do not use std::size_t as it's not portable.
+      int64_t const json_offset = binary_buf.size();
+      this->SaveConfig(&config);
+      std::string config_str;
+      Json::Dump(config, &config_str);
+      // concatonate the model and config at final output, it's a temporary solution for
+      // continuing support for binary model format
+      fo->Write(&serialisation_header_[0], serialisation_header_.size());
+      fo->Write(&json_offset, sizeof(json_offset));
+      fo->Write(&binary_buf[0], binary_buf.size());
+      fo->Write(&config_str[0], config_str.size());
+    }
+  }
+
+  void Load(dmlc::Stream* fi) override {
+    common::PeekableInStream fp(fi);
+    char c {0};
+    fp.PeekRead(&c, 1);
+    if (c == '{') {
+      std::string buffer;
+      common::FixedSizeStream{&fp}.Take(&buffer);
+      auto memory_snapshot = Json::Load({buffer.c_str(), buffer.size()});
+      this->LoadModel(memory_snapshot["Model"]);
+      this->LoadConfig(memory_snapshot["Config"]);
+    } else {
+      std::string header;
+      header.resize(serialisation_header_.size());
+      CHECK_EQ(fp.Read(&header[0], header.size()), serialisation_header_.size());
+      CHECK_EQ(header, serialisation_header_);
+
+      int64_t json_offset {-1};
+      CHECK_EQ(fp.Read(&json_offset, sizeof(json_offset)), sizeof(json_offset));
+      CHECK_GT(json_offset, 0);
+      std::string buffer;
+      common::FixedSizeStream{&fp}.Take(&buffer);
+
+      common::MemoryFixSizeBuffer binary_buf(&buffer[0], json_offset);
+      this->LoadModel(&binary_buf);
+
+      common::MemoryFixSizeBuffer json_buf {&buffer[0] + json_offset,
+                                            buffer.size() - json_offset};
+      auto config = Json::Load({buffer.c_str() + json_offset, buffer.size() - json_offset});
+      this->LoadConfig(config);
+    }
+  }
+
  std::vector<std::string> DumpModel(const FeatureMap& fmap,
                                     bool with_stats,
                                     std::string format) const override {
@@ -551,6 +607,7 @@ class LearnerImpl : public Learner {

  void UpdateOneIter(int iter, DMatrix* train) override {
    monitor_.Start("UpdateOneIter");
+    TrainingObserver::Instance().Update(iter);
    this->Configure();
    if (generic_parameters_.seed_per_iteration || rabit::IsDistributed()) {
      common::GlobalRandom().seed(generic_parameters_.seed * kRandSeedMagic + iter);
@@ -561,9 +618,13 @@ class LearnerImpl : public Learner {
    monitor_.Start("PredictRaw");
    this->PredictRaw(train, &preds_[train]);
    monitor_.Stop("PredictRaw");
+    TrainingObserver::Instance().Observe(preds_[train], "Predictions");
+
    monitor_.Start("GetGradient");
    obj_->GetGradient(preds_[train], train->Info(), iter, &gpair_);
    monitor_.Stop("GetGradient");
+    TrainingObserver::Instance().Observe(gpair_, "Gradients");
+
    gbm_->DoBoost(train, &gpair_, obj_.get());
    monitor_.Stop("UpdateOneIter");
  }
@@ -792,6 +853,10 @@ class LearnerImpl : public Learner {
  LearnerModelParamLegacy mparam_;
  LearnerModelParam learner_model_param_;
  LearnerTrainParam tparam_;
+  // Used to identify the offset of JSON string when
+  // `enable_experimental_json_serialization' is set to false.  Will be removed once JSON
+  // takes over.
+  std::string const serialisation_header_ { u8"CONFIG-offset:" };
  // configurations
  std::map<std::string, std::string> cfg_;
  std::map<std::string, std::string> attributes_;
@@ -811,9 +876,8 @@ class LearnerImpl : public Learner {

  common::Monitor monitor_;

-  /*! \brief saved config keys used to restore failed worker */
-  std::set<std::string> saved_configs_ = {"max_depth", "tree_method", "dsplit",
-    "seed", "silent", "num_round", "gamma", "min_child_weight"};
+  /*! \brief (Deprecated) saved config keys used to restore failed worker */
+  std::set<std::string> saved_configs_ = {"num_round"};
 };

 std::string const LearnerImpl::kEvalMetric {"eval_metric"};  // NOLINT
--- a/src/tree/tree_model.cc
+++ b/src/tree/tree_model.cc
@@ -682,13 +682,13 @@ void RegTree::LoadModel(Json const& in) {
    s.leaf_child_cnt = get<Integer const>(leaf_child_counts[i]);

    auto& n = nodes_[i];
-    auto left = get<Integer const>(lefts[i]);
-    auto right = get<Integer const>(rights[i]);
-    auto parent = get<Integer const>(parents[i]);
-    auto ind = get<Integer const>(indices[i]);
-    auto cond = get<Number const>(conds[i]);
-    auto dft_left = get<Boolean const>(default_left[i]);
-    n = Node(left, right, parent, ind, cond, dft_left);
+    bst_node_t left = get<Integer const>(lefts[i]);
+    bst_node_t right = get<Integer const>(rights[i]);
+    bst_node_t parent = get<Integer const>(parents[i]);
+    bst_feature_t ind = get<Integer const>(indices[i]);
+    float cond { get<Number const>(conds[i]) };
+    bool dft_left { get<Boolean const>(default_left[i]) };
+    n = Node{left, right, parent, ind, cond, dft_left};
  }


--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -1027,8 +1027,6 @@ class GPUHistMakerSpecialised {
    param_.UpdateAllowUnknown(args);
    generic_param_ = generic_param;
    hist_maker_param_.UpdateAllowUnknown(args);
-    device_ = generic_param_->gpu_id;
-    CHECK_GE(device_, 0) << "Must have at least one device";
    dh::CheckComputeCapability();

    monitor_.Init("updater_gpu_hist");
@@ -1041,6 +1039,7 @@ class GPUHistMakerSpecialised {
  void Update(HostDeviceVector<GradientPair>* gpair, DMatrix* dmat,
              const std::vector<RegTree*>& trees) {
    monitor_.StartCuda("Update");
+
    // rescale learning rate according to size of trees
    float lr = param_.learning_rate;
    param_.learning_rate = lr / trees.size();
@@ -1064,6 +1063,8 @@ class GPUHistMakerSpecialised {
  }

  void InitDataOnce(DMatrix* dmat) {
+    device_ = generic_param_->gpu_id;
+    CHECK_GE(device_, 0) << "Must have at least one device";
    info_ = &dmat->Info();
    reducer_.Init({device_});

@@ -1162,14 +1163,24 @@ class GPUHistMakerSpecialised {
 class GPUHistMaker : public TreeUpdater {
 public:
  void Configure(const Args& args) override {
+    // Used in test to count how many configurations are performed
+    LOG(DEBUG) << "[GPU Hist]: Configure";
    hist_maker_param_.UpdateAllowUnknown(args);
-    float_maker_.reset();
-    double_maker_.reset();
+    // The passed in args can be empty, if we simply purge the old maker without
+    // preserving parameters then we can't do Update on it.
+    TrainParam param;
+    if (float_maker_) {
+      param = float_maker_->param_;
+    } else if (double_maker_) {
+      param = double_maker_->param_;
+    }
    if (hist_maker_param_.single_precision_histogram) {
      float_maker_.reset(new GPUHistMakerSpecialised<GradientPair>());
+      float_maker_->param_ = param;
      float_maker_->Configure(args, tparam_);
    } else {
      double_maker_.reset(new GPUHistMakerSpecialised<GradientPairPrecise>());
+      double_maker_->param_ = param;
      double_maker_->Configure(args, tparam_);
    }
  }