Add Accelerated Failure Time loss for survival analysis task (#4763)
* [WIP] Add lower and upper bounds on the label for survival analysis * Update test MetaInfo.SaveLoadBinary to account for extra two fields * Don't clear qids_ for version 2 of MetaInfo * Add SetInfo() and GetInfo() method for lower and upper bounds * changes to aft * Add parameter class for AFT; use enum's to represent distribution and event type * Add AFT metric * changes to neg grad to grad * changes to binomial loss * changes to overflow * changes to eps * changes to code refactoring * changes to code refactoring * changes to code refactoring * Re-factor survival analysis * Remove aft namespace * Move function bodies out of AFTNormal and AFTLogistic, to reduce clutter * Move function bodies out of AFTLoss, to reduce clutter * Use smart pointer to store AFTDistribution and AFTLoss * Rename AFTNoiseDistribution enum to AFTDistributionType for clarity The enum class was not a distribution itself but a distribution type * Add AFTDistribution::Create() method for convenience * changes to extreme distribution * changes to extreme distribution * changes to extreme * changes to extreme distribution * changes to left censored * deleted cout * changes to x,mu and sd and code refactoring * changes to print * changes to hessian formula in censored and uncensored * changes to variable names and pow * changes to Logistic Pdf * changes to parameter * Expose lower and upper bound labels to R package * Use example weights; normalize log likelihood metric * changes to CHECK * changes to logistic hessian to standard formula * changes to logistic formula * Comply with coding style guideline * Revert back Rabit submodule * Revert dmlc-core submodule * Comply with coding style guideline (clang-tidy) * Fix an error in AFTLoss::Gradient() * Add missing files to amalgamation * Address @RAMitchell's comment: minimize future change in MetaInfo interface * Fix lint * Fix compilation error on 32-bit target, when size_t == bst_uint * Allocate sufficient memory to hold extra label info * Use OpenMP to speed up * Fix compilation on Windows * Address reviewer's feedback * Add unit tests for probability distributions * Make Metric subclass of Configurable * Address reviewer's feedback: Configure() AFT metric * Add a dummy test for AFT metric configuration * Complete AFT configuration test; remove debugging print * Rename AFT parameters * Clarify test comment * Add a dummy test for AFT loss for uncensored case * Fix a bug in AFT loss for uncensored labels * Complete unit test for AFT loss metric * Simplify unit tests for AFT metric * Add unit test to verify aggregate output from AFT metric * Use EXPECT_* instead of ASSERT_*, so that we run all unit tests * Use aft_loss_param when serializing AFTObj This is to be consistent with AFT metric * Add unit tests for AFT Objective * Fix OpenMP bug; clarify semantics for shared variables used in OpenMP loops * Add comments * Remove AFT prefix from probability distribution; put probability distribution in separate source file * Add comments * Define kPI and kEulerMascheroni in probability_distribution.h * Add probability_distribution.cc to amalgamation * Remove unnecessary diff * Address reviewer's feedback: define variables where they're used * Eliminate all INFs and NANs from AFT loss and gradient * Add demo * Add tutorial * Fix lint * Use 'survival:aft' to be consistent with 'survival:cox' * Move sample data to demo/data * Add visual demo with 1D toy data * Add Python tests Co-authored-by: Philip Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Avinash Barnwal
GitHub
119
src/objective/aft_obj.cc
Normal file
119
src/objective/aft_obj.cc
Normal file
@@ -0,0 +1,119 @@
|
||||
/*!
|
||||
* Copyright 2015 by Contributors
|
||||
* \file rank.cc
|
||||
* \brief Definition of aft loss.
|
||||
*/
|
||||
|
||||
#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;
|
||||
|
||||
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 default(none) \
|
||||
firstprivate(nsize, is_null_weight, aft_loss_distribution_scale) \
|
||||
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 default(none) firstprivate(ndata) 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"], ¶m_);
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user