Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

GPU implementation of AFT survival objective and metric (#5714)

Browse Source 
* Add interval accuracy

* De-virtualize AFT functions

* Lint

* Refactor AFT metric using GPU-CPU reducer

* Fix R build

* Fix build on Windows

* Fix copyright header

* Clang-tidy

* Fix crashing demo

* Fix typos in comment; explain GPU ID

* Remove unnecessary #include

* Add C++ test for interval accuracy

* Fix a bug in accuracy metric: use log pred

* Refactor AFT objective using GPU-CPU Transform

* Lint

* Fix lint

* Use Ninja to speed up build

* Use time, not /usr/bin/time

* Add cpu_build worker class, with concurrency = 1

* Use concurrency = 1 only for CUDA build

* concurrency = 1 for clang-tidy

* Address reviewer's feedback

* Update link to AFT paper
This commit is contained in:
Philip Hyunsu Cho
2020-07-17 01:18:13 -07:00
committed by GitHub
parent 7c2686146e
commit 71b0528a2f
20 changed files with 1050 additions and 822 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
src/objective/aft_obj.cc
View File

@@ -1,116 +1,21 @@
/*!
* Copyright 2015 by Contributors
* \file rank.cc
* \brief Definition of aft loss.
* Copyright 2019-2020 by Contributors
* \file aft_obj.cc
* \brief Definition of AFT loss for survival analysis.
* \author Avinash Barnwal, Hyunsu Cho and Toby Hocking
*/
#include <dmlc/omp.h>
#include <xgboost/logging.h>
#include <xgboost/objective.h>
#include <vector>
#include <limits>
#include <algorithm>
#include <memory>
#include <utility>
#include <cmath>
#include "xgboost/json.h"
#include "../common/math.h"
#include "../common/random.h"
#include "../common/survival_util.h"
using AFTParam = xgboost::common::AFTParam;
using AFTLoss = xgboost::common::AFTLoss;
// Dummy file to keep the CUDA conditional compile trick.
#include <dmlc/registry.h>
namespace xgboost {
namespace obj {
DMLC_REGISTRY_FILE_TAG(aft_obj);
class AFTObj : public ObjFunction {
public:
void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
param_.UpdateAllowUnknown(args);
loss_.reset(new AFTLoss(param_.aft_loss_distribution));
}
void GetGradient(const HostDeviceVector<bst_float>& preds,
const MetaInfo& info,
int iter,
HostDeviceVector<GradientPair>* out_gpair) override {
/* Boilerplate */
CHECK_EQ(preds.Size(), info.labels_lower_bound_.Size());
CHECK_EQ(preds.Size(), info.labels_upper_bound_.Size());
const auto& yhat = preds.HostVector();
const auto& y_lower = info.labels_lower_bound_.HostVector();
const auto& y_upper = info.labels_upper_bound_.HostVector();
const auto& weights = info.weights_.HostVector();
const bool is_null_weight = weights.empty();
out_gpair->Resize(yhat.size());
std::vector<GradientPair>& gpair = out_gpair->HostVector();
CHECK_LE(yhat.size(), static_cast<size_t>(std::numeric_limits<omp_ulong>::max()))
<< "yhat is too big";
const omp_ulong nsize = static_cast<omp_ulong>(yhat.size());
const float aft_loss_distribution_scale = param_.aft_loss_distribution_scale;
#pragma omp parallel for \
shared(weights, y_lower, y_upper, yhat, gpair)
for (omp_ulong i = 0; i < nsize; ++i) {
// If weights are empty, data is unweighted so we use 1.0 everywhere
const double w = is_null_weight ? 1.0 : weights[i];
const double grad = loss_->Gradient(y_lower[i], y_upper[i],
yhat[i], aft_loss_distribution_scale);
const double hess = loss_->Hessian(y_lower[i], y_upper[i],
yhat[i], aft_loss_distribution_scale);
gpair[i] = GradientPair(grad * w, hess * w);
}
}
void PredTransform(HostDeviceVector<bst_float> *io_preds) override {
// Trees give us a prediction in log scale, so exponentiate
std::vector<bst_float> &preds = io_preds->HostVector();
const long ndata = static_cast<long>(preds.size()); // NOLINT(*)
#pragma omp parallel for shared(preds)
for (long j = 0; j < ndata; ++j) { // NOLINT(*)
preds[j] = std::exp(preds[j]);
}
}
void EvalTransform(HostDeviceVector<bst_float> *io_preds) override {
// do nothing here, since the AFT metric expects untransformed prediction score
}
bst_float ProbToMargin(bst_float base_score) const override {
return std::log(base_score);
}
const char* DefaultEvalMetric() const override {
return "aft-nloglik";
}
void SaveConfig(Json* p_out) const override {
auto& out = *p_out;
out["name"] = String("survival:aft");
out["aft_loss_param"] = ToJson(param_);
}
void LoadConfig(Json const& in) override {
FromJson(in["aft_loss_param"], &param_);
loss_.reset(new AFTLoss(param_.aft_loss_distribution));
}
private:
AFTParam param_;
std::unique_ptr<AFTLoss> loss_;
};
// register the objective functions
XGBOOST_REGISTER_OBJECTIVE(AFTObj, "survival:aft")
.describe("AFT loss function")
.set_body([]() { return new AFTObj(); });
} // namespace obj
} // namespace xgboost
#ifndef XGBOOST_USE_CUDA
#include "aft_obj.cu"
#endif // XGBOOST_USE_CUDA

147
src/objective/aft_obj.cu Normal file
View File

@@ -0,0 +1,147 @@
/*!
* Copyright 2019-2020 by Contributors
* \file aft_obj.cu
* \brief Definition of AFT loss for survival analysis.
* \author Avinash Barnwal, Hyunsu Cho and Toby Hocking
*/
#include <vector>
#include <limits>
#include <memory>
#include <utility>
#include "xgboost/host_device_vector.h"
#include "xgboost/json.h"
#include "xgboost/parameter.h"
#include "xgboost/span.h"
#include "xgboost/logging.h"
#include "xgboost/objective.h"
#include "../common/transform.h"
#include "../common/survival_util.h"
using AFTParam = xgboost::common::AFTParam;
using ProbabilityDistributionType = xgboost::common::ProbabilityDistributionType;
template <typename Distribution>
using AFTLoss = xgboost::common::AFTLoss<Distribution>;
namespace xgboost {
namespace obj {
#if defined(XGBOOST_USE_CUDA)
DMLC_REGISTRY_FILE_TAG(aft_obj_gpu);
#endif // defined(XGBOOST_USE_CUDA)
class AFTObj : public ObjFunction {
public:
void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
param_.UpdateAllowUnknown(args);
}
template <typename Distribution>
void GetGradientImpl(const HostDeviceVector<bst_float> &preds,
const MetaInfo &info,
HostDeviceVector<GradientPair> *out_gpair,
size_t ndata, int device, bool is_null_weight,
float aft_loss_distribution_scale) {
common::Transform<>::Init(
[=] XGBOOST_DEVICE(size_t _idx,
common::Span<GradientPair> _out_gpair,
common::Span<const bst_float> _preds,
common::Span<const bst_float> _labels_lower_bound,
common::Span<const bst_float> _labels_upper_bound,
common::Span<const bst_float> _weights) {
const double pred = static_cast<double>(_preds[_idx]);
const double label_lower_bound = static_cast<double>(_labels_lower_bound[_idx]);
const double label_upper_bound = static_cast<double>(_labels_upper_bound[_idx]);
const float grad = static_cast<float>(
AFTLoss<Distribution>::Gradient(label_lower_bound, label_upper_bound,
pred, aft_loss_distribution_scale));
const float hess = static_cast<float>(
AFTLoss<Distribution>::Hessian(label_lower_bound, label_upper_bound,
pred, aft_loss_distribution_scale));
const bst_float w = is_null_weight ? 1.0f : _weights[_idx];
_out_gpair[_idx] = GradientPair(grad * w, hess * w);
},
common::Range{0, static_cast<int64_t>(ndata)}, device).Eval(
out_gpair, &preds, &info.labels_lower_bound_, &info.labels_upper_bound_,
&info.weights_);
}
void GetGradient(const HostDeviceVector<bst_float>& preds,
const MetaInfo& info,
int iter,
HostDeviceVector<GradientPair>* out_gpair) override {
const size_t ndata = preds.Size();
CHECK_EQ(info.labels_lower_bound_.Size(), ndata);
CHECK_EQ(info.labels_upper_bound_.Size(), ndata);
out_gpair->Resize(ndata);
const int device = tparam_->gpu_id;
const float aft_loss_distribution_scale = param_.aft_loss_distribution_scale;
const bool is_null_weight = info.weights_.Size() == 0;
if (!is_null_weight) {
CHECK_EQ(info.weights_.Size(), ndata)
<< "Number of weights should be equal to number of data points.";
}
switch (param_.aft_loss_distribution) {
case common::ProbabilityDistributionType::kNormal:
GetGradientImpl<common::NormalDistribution>(preds, info, out_gpair, ndata, device,
is_null_weight, aft_loss_distribution_scale);
break;
case common::ProbabilityDistributionType::kLogistic:
GetGradientImpl<common::LogisticDistribution>(preds, info, out_gpair, ndata, device,
is_null_weight, aft_loss_distribution_scale);
break;
case common::ProbabilityDistributionType::kExtreme:
GetGradientImpl<common::ExtremeDistribution>(preds, info, out_gpair, ndata, device,
is_null_weight, aft_loss_distribution_scale);
break;
default:
LOG(FATAL) << "Unrecognized distribution";
}
}
void PredTransform(HostDeviceVector<bst_float> *io_preds) override {
// Trees give us a prediction in log scale, so exponentiate
common::Transform<>::Init(
[] XGBOOST_DEVICE(size_t _idx, common::Span<bst_float> _preds) {
_preds[_idx] = exp(_preds[_idx]);
}, common::Range{0, static_cast<int64_t>(io_preds->Size())},
tparam_->gpu_id)
.Eval(io_preds);
}
void EvalTransform(HostDeviceVector<bst_float> *io_preds) override {
// do nothing here, since the AFT metric expects untransformed prediction score
}
bst_float ProbToMargin(bst_float base_score) const override {
return std::log(base_score);
}
const char* DefaultEvalMetric() const override {
return "aft-nloglik";
}
void SaveConfig(Json* p_out) const override {
auto& out = *p_out;
out["name"] = String("survival:aft");
out["aft_loss_param"] = ToJson(param_);
}
void LoadConfig(Json const& in) override {
FromJson(in["aft_loss_param"], &param_);
}
private:
AFTParam param_;
};
// register the objective functions
XGBOOST_REGISTER_OBJECTIVE(AFTObj, "survival:aft")
.describe("AFT loss function")
.set_body([]() { return new AFTObj(); });
} // namespace obj
} // namespace xgboost