data.h

前往此文件文档。

 
#ifndef XGBOOST_DATA_H_
#define XGBOOST_DATA_H_
 
#include <dmlc/base.h>
#include <dmlc/io.h> // for Stream
#include <dmlc/serializer.h> // for Handler
#include <xgboost/base.h>
#include <xgboost/host_device_vector.h>
#include <xgboost/linalg.h>
#include <xgboost/span.h>
#include <xgboost/string_view.h>
 
#include <algorithm>
#include <cstdint> // for int32_t, uint8_t
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include <vector>
 
namespace xgboost {
// 提前声明 dmatrix。
class DMatrix;
struct Context;
 
enum class DataType : uint8_t {
 kFloat32 = 1,
 kDouble = 2,
 kUInt32 = 3,
 kUInt64 = 4,
 kStr = 5
};
 
enum class FeatureType : uint8_t { kNumerical = 0, kCategorical = 1 };
 
enum class DataSplitMode : int { kRow = 0, kCol = 1 };
 
// 元信息使用的容器的提前声明。
class CatContainer;
 
class MetaInfo {
 public
 static constexpr uint64_t kNumField = 13;
 
 bst_idx_t num_row_{0}; // NOLINT
 uint64_t num_col_{0}; // NOLINT
 uint64_t num_nonzero_{0}; // NOLINT
 linalg::Tensor<float, 2> labels;
 DataSplitMode data_split_mode{DataSplitMode::kRow};
 std::vector<bst_group_t> group_ptr_; // NOLINT
 HostDeviceVector<bst_float> weights_; // NOLINT
 linalg::Matrix<float> base_margin_; // NOLINT
 HostDeviceVector<bst_float> labels_lower_bound_; // NOLINT
 HostDeviceVector<bst_float> labels_upper_bound_; // NOLINT
 
 std::vector<std::string> feature_type_names;
 std::vector<std::string> feature_names;
  /*
 * \brief 每个特征的类型。在指定 feature_type_names 时自动设置。
   */
 HostDeviceVector<FeatureType> feature_types;
  /*
 * \brief 每个特征的权重，用于在使用列抽样时定义每个特征被选中的概率。
 * selected when using column sampling.
   */
 HostDeviceVector<float> feature_weights;
 
 MetaInfo();
 MetaInfo(MetaInfo&& that) = default;
 MetaInfo(MetaInfo const& that) = delete;
 MetaInfo& operator=(MetaInfo&& that) = default;
 MetaInfo& operator=(MetaInfo const& that) = delete;
 
 void Validate(DeviceOrd device) const;
 MetaInfo Slice(Context const* ctx, common::Span<bst_idx_t const> ridxs, bst_idx_t nnz) const;
 
 MetaInfo Copy() const;
 bool IsDense() const { return num_col_ * num_row_ == num_nonzero_; }
 inline bst_float GetWeight(size_t i) const {
 return weights_.Size() != 0 ? weights_.HostVector()[i] : 1.0f;
  }
 const std::vector<size_t>& LabelAbsSort(Context const* ctx) const;
 void Clear();
 void LoadBinary(dmlc::Stream* fi);
 void SaveBinary(dmlc::Stream* fo) const;
 void SetInfo(Context const& ctx, StringView key, StringView interface_str);
 
 void GetInfo(char const* key, bst_ulong* out_len, DataType dtype,
 const void** out_dptr) const;
 
 void SetFeatureInfo(const char *key, const char **info, const bst_ulong size);
 void GetFeatureInfo(const char *field, std::vector<std::string>* out_str_vecs) const;
 
 void Extend(MetaInfo const& that, bool accumulate_rows, bool check_column);
 void SynchronizeNumberOfColumns(Context const* ctx, DataSplitMode split_mode);
 
 [[nodiscard]] bool IsRowSplit() const { return data_split_mode == DataSplitMode::kRow; }
 [[nodiscard]] bool IsColumnSplit() const { return data_split_mode == DataSplitMode::kCol; }
 [[nodiscard]] bool IsRanking() const { return !group_ptr_.empty(); }
 
 [[nodiscard]] bool IsVerticalFederated() const;
 
 bool ShouldHaveLabels() const;
 bool HasCategorical() const { return has_categorical_; }
 [[nodiscard]] CatContainer const* Cats() const;
 [[nodiscard]] CatContainer* Cats();
 [[nodiscard]] std::shared_ptr<CatContainer const> CatsShared() const;
 void Cats(std::shared_ptr<CatContainer> cats);
 
 private
 void SetInfoFromHost(Context const* ctx, StringView key, Json arr);
 void SetInfoFromCUDA(Context const* ctx, StringView key, Json arr);
 
 mutable std::vector<size_t> label_order_cache_;
 bool has_categorical_{false};
 
 std::shared_ptr<CatContainer> cats_;
};
 
struct Entry {
 bst_feature_t index;
 bst_float fvalue;
 Entry() = default;
 XGBOOST_DEVICE Entry(bst_feature_t index, bst_float fvalue) : index(index), fvalue(fvalue) {}
 inline static bool CmpValue(const Entry& a, const Entry& b) {
 return a.fvalue < b.fvalue;
  }
 static bool CmpIndex(Entry const& a, Entry const& b) {
 return a.index < b.index;
  }
 inline bool operator==(const Entry& other) const {
 return (this->index == other.index && this->fvalue == other.fvalue);
  }
};
 
struct BatchParam {
 bst_bin_t max_bin{0};
 common::Span<float const> hess;
 bool regen{false};
 bool forbid_regen{false};
 double sparse_thresh{std::numeric_limits<double>::quiet_NaN()};
 bool prefetch_copy{true};
 std::int32_t n_prefetch_batches{3};
 BatchParam() = default;
 BatchParam(bst_bin_t max_bin, double sparse_thresh)
 : max_bin{max_bin}, sparse_thresh{sparse_thresh} {}
 BatchParam(bst_bin_t max_bin, common::Span<float const> hessian, bool regenerate)
 : max_bin{max_bin}, hess{hessian}, regen{regenerate} {}
 
 [[nodiscard]] bool ParamNotEqual(BatchParam const& other) const {
 // 检查非浮点参数。
 bool cond = max_bin != other.max_bin;
 // 检查稀疏阈值。
 bool l_nan = std::isnan(sparse_thresh);
 bool r_nan = std::isnan(other.sparse_thresh);
 bool st_chg = (l_nan != r_nan) || (!l_nan && !r_nan && (sparse_thresh != other.sparse_thresh));
 cond |= st_chg;
 
 return cond;
  }
 [[nodiscard]] bool Initialized() const { return max_bin != 0; }
 [[nodiscard]] BatchParam MakeCache() const {
 auto p = *this;
 // 这些参数与梯度索引最初是如何生成的无关。
 // first place.
 p.regen = false;
 p.forbid_regen = false;
 return p;
  }
};
 
struct HostSparsePageView {
 using Inst = common::Span<Entry const>;
 
 common::Span<bst_idx_t const> offset;
 common::Span<Entry const> data;
 
 [[nodiscard]] Inst operator[](std::size_t i) const {
 auto size = *(offset.data() + i + 1) - *(offset.data() + i);
 return {data.data() + *(offset.data() + i), static_cast<Inst::index_type>(size)};
  }
 
 [[nodiscard]] size_t Size() const { return offset.size() == 0 ? 0 : offset.size() - 1; }
};
 
class SparsePage {
 public
 // 每行的偏移量。
 HostDeviceVector<bst_idx_t> offset;
 HostDeviceVector<Entry> data;
 
 size_t base_rowid {0};
 
 using Inst = common::Span<Entry const>;
 
 [[nodiscard]] HostSparsePageView GetView() const {
 return {offset.ConstHostSpan(), data.ConstHostSpan()};
  }
 
 SparsePage() {
 this->Clear();
  }
 
 SparsePage(SparsePage const& that) = delete;
 SparsePage(SparsePage&& that) = default;
 SparsePage& operator=(SparsePage const& that) = delete;
 SparsePage& operator=(SparsePage&& that) = default;
 virtual ~SparsePage() = default;
 
 [[nodiscard]] size_t Size() const {
 return offset.Size() == 0 ? 0 : offset.Size() - 1;
  }
 
 [[nodiscard]] size_t MemCostBytes() const {
 return offset.Size() * sizeof(size_t) + data.Size() * sizeof(Entry);
  }
 
 inline void Clear() {
 base_rowid = 0;
 auto& offset_vec = offset.HostVector();
 offset_vec.clear();
 offset_vec.push_back(0);
 data.HostVector().clear();
  }
 
 inline void SetBaseRowId(size_t row_id) {
 base_rowid = row_id;
  }
 
 [[nodiscard]] SparsePage GetTranspose(int num_columns, int32_t n_threads) const;
 
 void SortIndices(int32_t n_threads);
 [[nodiscard]] bool IsIndicesSorted(int32_t n_threads) const;
 void Reindex(uint64_t feature_offset, int32_t n_threads);
 
 void SortRows(int32_t n_threads);
 
 template <typename AdapterBatchT>
 bst_idx_t Push(AdapterBatchT const& batch, float missing, std::int32_t nthread);
 
 void Push(const SparsePage &batch);
 void PushCSC(const SparsePage& batch);
};
 
class CSCPage: public SparsePage {
 public
 CSCPage() : SparsePage() {}
 explicit CSCPage(SparsePage page) : SparsePage(std::move(page)) {}
};
 
class ExtSparsePage {
 public
 std::shared_ptr<SparsePage const> page;
 explicit ExtSparsePage(std::shared_ptr<SparsePage const> p) : page{std::move(p)} {}
};
 
class SortedCSCPage : public SparsePage {
 public
 SortedCSCPage() : SparsePage() {}
 explicit SortedCSCPage(SparsePage page) : SparsePage(std::move(page)) {}
};
 
class EllpackPage;
class GHistIndexMatrix;
 
template<typename T>
class BatchIteratorImpl {
 public
 using iterator_category = std::forward_iterator_tag; // NOLINT
 virtual ~BatchIteratorImpl() = default;
 virtual const T& operator*() const = 0;
 virtual BatchIteratorImpl& operator++() = 0;
 [[nodiscard]] virtual bool AtEnd() const = 0;
 virtual std::shared_ptr<T const> Page() const = 0;
};
 
template<typename T>
class BatchIterator {
 public
 using iterator_category = std::forward_iterator_tag; // NOLINT
 explicit BatchIterator(BatchIteratorImpl<T>* impl) { impl_.reset(impl); }
 explicit BatchIterator(std::shared_ptr<BatchIteratorImpl<T>> impl) { impl_ = impl; }
 
 BatchIterator &operator++() {
 CHECK(impl_ != nullptr);
 ++(*impl_);
 return *this;
  }
 
 const T& operator*() const {
 CHECK(impl_ != nullptr);
 return *(*impl_);
  }
 
 [[nodiscard]] bool operator!=(const BatchIterator&) const { return !this->AtEnd(); }
 
 [[nodiscard]] bool AtEnd() const {
 CHECK(impl_ != nullptr);
 return impl_->AtEnd();
  }
 
 [[nodiscard]] std::shared_ptr<T const> Page() const {
 return impl_->Page();
  }
 
 private
 std::shared_ptr<BatchIteratorImpl<T>> impl_;
};
 
template<typename T>
class BatchSet {
 public
 explicit BatchSet(BatchIterator<T> begin_iter) : begin_iter_(std::move(begin_iter)) {}
 BatchIterator<T> begin() { return begin_iter_; } // NOLINT
 BatchIterator<T> end() { return BatchIterator<T>(nullptr); } // NOLINT
 
 private
 BatchIterator<T> begin_iter_;
};
 
struct XGBAPIThreadLocalEntry;
 
// 外部内存 DMatrix 的配置。
struct ExtMemConfig {
 // 缓存前缀，如果缓存位于主机内存中，则不使用。（on_host 为 true）
 std::string cache;
 // Ellpack 页是否存储在主机内存中。
 bool on_host;
 // GPU 实现的主机缓存/总缓存。
 float cache_host_ratio;
 // 缓存中每个 ellpack 页的最小字节数。仅用于主机内 ExtMemQdm。
 // ExtMemQdm.
 std::int64_t min_cache_page_bytes;
 // 缺失值。
 float missing;
 // CPU 线程数。
 std::int32_t n_threads{0};
 // 缓存可压缩的比例。用于测试。
 float hw_decomp_ratio{std::numeric_limits<float>::quiet_NaN()};
 // 回退到使用 nvcomp。用于测试。
 bool allow_decomp_fallback{false};
 
 ExtMemConfig() = delete;
 ExtMemConfig(std::string cache, bool on_host, float h_ratio, std::int64_t min_cache,
 float missing, std::int32_t n_threads)
 : cache{std::move(cache)},
 on_host{on_host},
 cache_host_ratio{h_ratio},
 min_cache_page_bytes{min_cache},
 missing{missing},
 n_threads{n_threads} {}
 
 ExtMemConfig& SetParamsForTest(float _hw_decomp_ratio, bool _allow_decomp_fallback) {
 this->hw_decomp_ratio = _hw_decomp_ratio;
 this->allow_decomp_fallback = _allow_decomp_fallback;
 return *this;
  }
};
 
class DMatrix {
 public
 DMatrix() = default;
 [[nodiscard]] virtual MetaInfo& Info() = 0;
 virtual void SetInfo(const char* key, std::string const& interface_str) {
 auto const& ctx = *this->Ctx();
 this->Info().SetInfo(ctx, key, StringView{interface_str});
  }
 [[nodiscard]] virtual const MetaInfo& Info() const = 0;
 
 [[nodiscard]] XGBAPIThreadLocalEntry& GetThreadLocal() const;
 [[nodiscard]] virtual Context const* Ctx() const = 0;
 
 template <typename T>
 BatchSet<T> GetBatches();
 template <typename T>
 BatchSet<T> GetBatches(Context const* ctx);
 template <typename T>
 BatchSet<T> GetBatches(Context const* ctx, const BatchParam& param);
 template <typename T>
 [[nodiscard]] bool PageExists() const;
 
 [[nodiscard]] bool SingleColBlock() const { return this->NumBatches() == 1; }
 [[nodiscard]] virtual std::int32_t NumBatches() const { return 1; }
 
 virtual ~DMatrix();
 
 [[nodiscard]] bool IsDense() const { return this->Info().IsDense(); }
 
 static DMatrix* Load(const std::string& uri, bool silent = true,
 DataSplitMode data_split_mode = DataSplitMode::kRow);
 
 template <typename AdapterT>
 static DMatrix* Create(AdapterT* adapter, float missing, int nthread,
 const std::string& cache_prefix = "",
 DataSplitMode data_split_mode = DataSplitMode::kRow);
 
 template <typename DataIterHandle, typename DMatrixHandle, typename DataIterResetCallback,
 typename XGDMatrixCallbackNext>
 static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
 DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing,
 std::int32_t nthread, bst_bin_t max_bin, std::int64_t max_quantile_blocks);
 
 template <typename DataIterHandle, typename DMatrixHandle, typename DataIterResetCallback,
 typename XGDMatrixCallbackNext>
 static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, DataIterResetCallback* reset,
 XGDMatrixCallbackNext* next, ExtMemConfig const& config);
 
 template <typename DataIterHandle, typename DMatrixHandle, typename DataIterResetCallback,
 typename XGDMatrixCallbackNext>
 static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
 DataIterResetCallback* reset, XGDMatrixCallbackNext* next,
 bst_bin_t max_bin, std::int64_t max_quantile_blocks,
 ExtMemConfig const& config);
 
 virtual DMatrix *Slice(common::Span<int32_t const> ridxs) = 0;
 
 virtual DMatrix* SliceCol(int num_slices, int slice_id) = 0;
 [[nodiscard]] CatContainer const* Cats() const { return this->CatsShared().get(); }
 [[nodiscard]] std::shared_ptr<CatContainer const> CatsShared() const {
 return this->Info().CatsShared();
  }
 
 protected
 virtual BatchSet<SparsePage> GetRowBatches() = 0;
 virtual BatchSet<CSCPage> GetColumnBatches(Context const* ctx) = 0;
 virtual BatchSet<SortedCSCPage> GetSortedColumnBatches(Context const* ctx) = 0;
 virtual BatchSet<EllpackPage> GetEllpackBatches(Context const* ctx, BatchParam const& param) = 0;
 virtual BatchSet<GHistIndexMatrix> GetGradientIndex(Context const* ctx,
 BatchParam const& param) = 0;
 virtual BatchSet<ExtSparsePage> GetExtBatches(Context const* ctx, BatchParam const& param) = 0;
 
 [[nodiscard]] virtual bool EllpackExists() const = 0;
 [[nodiscard]] virtual bool GHistIndexExists() const = 0;
 [[nodiscard]] virtual bool SparsePageExists() const = 0;
};
 
template <>
inline BatchSet<SparsePage> DMatrix::GetBatches() {
 return GetRowBatches();
}
 
template <>
inline bool DMatrix::PageExists<EllpackPage>() const {
 return this->EllpackExists();
}
 
template <>
inline bool DMatrix::PageExists<GHistIndexMatrix>() const {
 return this->GHistIndexExists();
}
 
template <>
inline bool DMatrix::PageExists<SparsePage>() const {
 return this->SparsePageExists();
}
 
template <>
inline BatchSet<SparsePage> DMatrix::GetBatches(Context const*) {
 return GetRowBatches();
}
 
template <>
inline BatchSet<CSCPage> DMatrix::GetBatches(Context const* ctx) {
 return GetColumnBatches(ctx);
}
 
template <>
inline BatchSet<SortedCSCPage> DMatrix::GetBatches(Context const* ctx) {
 return GetSortedColumnBatches(ctx);
}
 
template <>
inline BatchSet<EllpackPage> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
 return GetEllpackBatches(ctx, param);
}
 
template <>
inline BatchSet<GHistIndexMatrix> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
 return GetGradientIndex(ctx, param);
}
 
template <>
inline BatchSet<ExtSparsePage> DMatrix::GetBatches(Context const* ctx, BatchParam const& param) {
 return GetExtBatches(ctx, param);
}
} // namespace xgboost
 
DECLARE_FIELD_ENUM_CLASS(xgboost::DataSplitMode);
 
namespace dmlc {
DMLC_DECLARE_TRAITS(is_pod, xgboost::Entry, true);
 
namespace serializer {
 
template <>
struct Handler<xgboost::Entry> {
 inline static void Write(Stream* strm, const xgboost::Entry& data) {
 strm->Write(data.index);
 strm->Write(data.fvalue);
  }
 
 inline static bool Read(Stream* strm, xgboost::Entry* data) {
 return strm->Read(&data->index) && strm->Read(&data->fvalue);
  }
};
 
} // namespace serializer
} // namespace dmlc
#endif // XGBOOST_DATA_H_