Thread safe, inplace prediction. (#5389)

Normal prediction with DMatrix is now thread safe with locks.  Added inplace prediction is lock free thread safe.

When data is on device (cupy, cudf), the returned data is also on device.

* Implementation for numpy, csr, cudf and cupy.

* Implementation for dask.

* Remove sync in simple dmatrix.
This commit is contained in:
Jiaming Yuan
2020-03-30 15:35:28 +08:00
committed by GitHub
parent 7f980e9f83
commit 6601a641d7
25 changed files with 1217 additions and 167 deletions

View File

@@ -9,6 +9,7 @@
#define XGBOOST_GBM_H_
#include <dmlc/registry.h>
#include <dmlc/any.h>
#include <xgboost/base.h>
#include <xgboost/data.h>
#include <xgboost/host_device_vector.h>
@@ -92,6 +93,22 @@ class GradientBooster : public Model, public Configurable {
PredictionCacheEntry* out_preds,
bool training,
unsigned ntree_limit = 0) = 0;
/*!
* \brief Inplace prediction.
*
* \param x A type erased data adapter.
* \param missing Missing value in the data.
* \param [in,out] out_preds The output preds.
* \param layer_begin (Optional) Begining of boosted tree layer used for prediction.
* \param layer_end (Optional) End of booster layer. 0 means do not limit trees.
*/
virtual void InplacePredict(dmlc::any const &x, float missing,
PredictionCacheEntry *out_preds,
uint32_t layer_begin = 0,
uint32_t layer_end = 0) const {
LOG(FATAL) << "Inplace predict is not supported by current booster.";
}
/*!
* \brief online prediction function, predict score for one instance at a time
* NOTE: use the batch prediction interface if possible, batch prediction is usually

View File

@@ -8,6 +8,7 @@
#ifndef XGBOOST_LEARNER_H_
#define XGBOOST_LEARNER_H_
#include <dmlc/any.h>
#include <rabit/rabit.h>
#include <xgboost/base.h>
#include <xgboost/feature_map.h>
@@ -120,6 +121,21 @@ class Learner : public Model, public Configurable, public rabit::Serializable {
bool approx_contribs = false,
bool pred_interactions = false) = 0;
/*!
* \brief Inplace prediction.
*
* \param x A type erased data adapter.
* \param type Prediction type.
* \param missing Missing value in the data.
* \param [in,out] out_preds Pointer to output prediction vector.
* \param layer_begin (Optional) Begining of boosted tree layer used for prediction.
* \param layer_end (Optional) End of booster layer. 0 means do not limit trees.
*/
virtual void InplacePredict(dmlc::any const& x, std::string const& type,
float missing,
HostDeviceVector<bst_float> **out_preds,
uint32_t layer_begin = 0, uint32_t layer_end = 0) = 0;
void LoadModel(Json const& in) override = 0;
void SaveModel(Json* out) const override = 0;

View File

@@ -16,6 +16,7 @@
#include <unordered_map>
#include <utility>
#include <vector>
#include <mutex>
// Forward declarations
namespace xgboost {
@@ -54,6 +55,7 @@ struct PredictionCacheEntry {
class PredictionContainer {
std::unordered_map<DMatrix *, PredictionCacheEntry> container_;
void ClearExpiredEntries();
std::mutex cache_lock_;
public:
PredictionContainer() = default;
@@ -133,6 +135,18 @@ class Predictor {
const gbm::GBTreeModel& model, int tree_begin,
uint32_t const ntree_limit = 0) = 0;
/**
* \brief Inplace prediction.
* \param x Type erased data adapter.
* \param model The model to predict from.
* \param missing Missing value in the data.
* \param [in,out] out_preds The output preds.
* \param tree_begin (Optional) Begining of boosted trees used for prediction.
* \param tree_end (Optional) End of booster trees. 0 means do not limit trees.
*/
virtual void InplacePredict(dmlc::any const &x, const gbm::GBTreeModel &model,
float missing, PredictionCacheEntry *out_preds,
uint32_t tree_begin = 0, uint32_t tree_end = 0) const = 0;
/**
* \brief online prediction function, predict score for one instance at a time
* NOTE: use the batch prediction interface if possible, batch prediction is