device.

Author: trivialfis

src/objective/multiclass_obj.cu

@@ -14,6 +14,7 @@
 #include "../common/linalg_op.h"
 #include "../common/math.h"
 #include "../common/optional_weight.h"  // for MakeOptionalWeights
+#include "../common/stats.h"            // for Mean
 #include "../common/transform.h"
 #include "xgboost/data.h"
 #include "xgboost/json.h"
@@ -197,7 +198,7 @@ class SoftmaxMultiClassObj : public ObjFunction {
     *base_score = linalg::Zeros<float>(this->ctx_, n_classes);
 
     std::size_t n = info.labels.Size();
-
+    // Calculate probability
     auto labels = info.labels.View(ctx_->Device());
     auto weights = common::MakeOptionalWeights(this->ctx_->Device(), info.weights_);
     auto intercept = base_score->View(ctx_->Device());
@@ -210,15 +211,13 @@ class SoftmaxMultiClassObj : public ObjFunction {
     CHECK_GE(sum_weight, kRtEps);
     linalg::VecScaDiv(this->ctx_, intercept, sum_weight);
 
-    double sum_intercepts = 0.;
-    for (std::int64_t ix = 0; ix < n_classes; ix++) {
-      intercept(ix) = std::log(intercept(ix));
-      sum_intercepts += intercept(ix);
-    }
-    const double mean_intercepts = sum_intercepts / static_cast<double>(n_classes);
-    for (std::int64_t ix = 0; ix < n_classes; ix++) {
-      intercept(ix) -= mean_intercepts;
-    }
+    // Transform it back to margin
+    linalg::Vector<float> mean_intercepts;
+    CHECK_EQ(base_score->Size(), n_classes);
+    common::Mean(this->ctx_, *base_score, &mean_intercepts);
+    auto d_mean = mean_intercepts.View(this->ctx_->Device());
+    TransformKernel(this->ctx_, intercept,
+                    [=] XGBOOST_DEVICE(float v) { return log(v) - d_mean(0); });
   }
 
  private: