rank_metric.hpp 6 KB
Newer Older
1
2
3
4
/*!
 * Copyright (c) 2016 Microsoft Corporation. All rights reserved.
 * Licensed under the MIT License. See LICENSE file in the project root for license information.
 */
5
6
#ifndef LIGHTGBM_SRC_METRIC_RANK_METRIC_HPP_
#define LIGHTGBM_SRC_METRIC_RANK_METRIC_HPP_
Guolin Ke's avatar
Guolin Ke committed
7
8

#include <LightGBM/metric.h>
Guolin Ke's avatar
Guolin Ke committed
9
10
#include <LightGBM/utils/common.h>
#include <LightGBM/utils/log.h>
11
#include <LightGBM/utils/openmp_wrapper.h>
Guolin Ke's avatar
Guolin Ke committed
12

13
14
15
16
#include <string>
#include <sstream>
#include <vector>

Guolin Ke's avatar
Guolin Ke committed
17
18
19
namespace LightGBM {

class NDCGMetric:public Metric {
Nikita Titov's avatar
Nikita Titov committed
20
 public:
Guolin Ke's avatar
Guolin Ke committed
21
  explicit NDCGMetric(const Config& config) {
Guolin Ke's avatar
Guolin Ke committed
22
    // get eval position
Guolin Ke's avatar
Guolin Ke committed
23
24
25
26
    eval_at_ = config.eval_at;
    auto label_gain = config.label_gain;
    DCGCalculator::DefaultEvalAt(&eval_at_);
    DCGCalculator::DefaultLabelGain(&label_gain);
Guolin Ke's avatar
Guolin Ke committed
27
    // initialize DCG calculator
Guolin Ke's avatar
Guolin Ke committed
28
    DCGCalculator::Init(label_gain);
Guolin Ke's avatar
Guolin Ke committed
29
30
31
32
  }

  ~NDCGMetric() {
  }
Guolin Ke's avatar
Guolin Ke committed
33
  void Init(const Metadata& metadata, data_size_t num_data) override {
34
    for (auto k : eval_at_) {
35
      name_.emplace_back(std::string("ndcg@") + std::to_string(k));
36
    }
Guolin Ke's avatar
Guolin Ke committed
37
38
39
    num_data_ = num_data;
    // get label
    label_ = metadata.label();
40
41
    num_queries_ = metadata.num_queries();
    DCGCalculator::CheckMetadata(metadata, num_queries_);
42
    DCGCalculator::CheckLabel(label_, num_data_);
Guolin Ke's avatar
Guolin Ke committed
43
44
45
    // get query boundaries
    query_boundaries_ = metadata.query_boundaries();
    if (query_boundaries_ == nullptr) {
46
      Log::Fatal("The NDCG metric requires query information");
Guolin Ke's avatar
Guolin Ke committed
47
48
49
50
    }
    // get query weights
    query_weights_ = metadata.query_weights();
    if (query_weights_ == nullptr) {
51
      sum_query_weights_ = static_cast<double>(num_queries_);
Guolin Ke's avatar
Guolin Ke committed
52
53
54
55
56
57
    } else {
      sum_query_weights_ = 0.0f;
      for (data_size_t i = 0; i < num_queries_; ++i) {
        sum_query_weights_ += query_weights_[i];
      }
    }
58
    inverse_max_dcgs_.resize(num_queries_);
Andrew Ziem's avatar
Andrew Ziem committed
59
    // cache the inverse max DCG for all queries, used to calculate NDCG
60
    #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static)
Guolin Ke's avatar
Guolin Ke committed
61
    for (data_size_t i = 0; i < num_queries_; ++i) {
62
      inverse_max_dcgs_[i].resize(eval_at_.size(), 0.0f);
Guolin Ke's avatar
Guolin Ke committed
63
      DCGCalculator::CalMaxDCG(eval_at_, label_ + query_boundaries_[i],
64
65
                               query_boundaries_[i + 1] - query_boundaries_[i],
                               &inverse_max_dcgs_[i]);
Guolin Ke's avatar
Guolin Ke committed
66
      for (size_t j = 0; j < inverse_max_dcgs_[i].size(); ++j) {
67
68
69
        if (inverse_max_dcgs_[i][j] > 0.0f) {
          inverse_max_dcgs_[i][j] = 1.0f / inverse_max_dcgs_[i][j];
        } else {
Andrew Ziem's avatar
Andrew Ziem committed
70
          // marking negative for all negative queries.
Guolin Ke's avatar
Guolin Ke committed
71
          // if one meet this query, it's ndcg will be set as -1.
72
          inverse_max_dcgs_[i][j] = -1.0f;
Guolin Ke's avatar
Guolin Ke committed
73
74
75
76
77
        }
      }
    }
  }

Guolin Ke's avatar
Guolin Ke committed
78
  const std::vector<std::string>& GetName() const override {
79
    return name_;
80
81
  }

82
  double factor_to_bigger_better() const override {
83
    return 1.0f;
84
85
  }

Guolin Ke's avatar
Guolin Ke committed
86
  std::vector<double> Eval(const double* score, const ObjectiveFunction*) const override {
87
    int num_threads = OMP_NUM_THREADS();
88
    // some buffers for multi-threading sum up
89
    std::vector<std::vector<double>> result_buffer_;
90
    for (int i = 0; i < num_threads; ++i) {
91
92
      result_buffer_.emplace_back(eval_at_.size(), 0.0f);
    }
93
    std::vector<double> tmp_dcg(eval_at_.size(), 0.0f);
94
    if (query_weights_ == nullptr) {
95
      #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static) firstprivate(tmp_dcg)
96
97
98
99
100
101
      for (data_size_t i = 0; i < num_queries_; ++i) {
        const int tid = omp_get_thread_num();
        // if all doc in this query are all negative, let its NDCG=1
        if (inverse_max_dcgs_[i][0] <= 0.0f) {
          for (size_t j = 0; j < eval_at_.size(); ++j) {
            result_buffer_[tid][j] += 1.0f;
Guolin Ke's avatar
Guolin Ke committed
102
          }
103
104
105
        } else {
          // calculate DCG
          DCGCalculator::CalDCG(eval_at_, label_ + query_boundaries_[i],
106
107
                                score + query_boundaries_[i],
                                query_boundaries_[i + 1] - query_boundaries_[i], &tmp_dcg);
108
109
110
          // calculate NDCG
          for (size_t j = 0; j < eval_at_.size(); ++j) {
            result_buffer_[tid][j] += tmp_dcg[j] * inverse_max_dcgs_[i][j];
Guolin Ke's avatar
Guolin Ke committed
111
112
113
          }
        }
      }
114
    } else {
115
      #pragma omp parallel for num_threads(OMP_NUM_THREADS()) schedule(static) firstprivate(tmp_dcg)
116
117
118
119
120
121
122
123
124
125
      for (data_size_t i = 0; i < num_queries_; ++i) {
        const int tid = omp_get_thread_num();
        // if all doc in this query are all negative, let its NDCG=1
        if (inverse_max_dcgs_[i][0] <= 0.0f) {
          for (size_t j = 0; j < eval_at_.size(); ++j) {
            result_buffer_[tid][j] += 1.0f;
          }
        } else {
          // calculate DCG
          DCGCalculator::CalDCG(eval_at_, label_ + query_boundaries_[i],
126
127
                                score + query_boundaries_[i],
                                query_boundaries_[i + 1] - query_boundaries_[i], &tmp_dcg);
128
129
130
131
          // calculate NDCG
          for (size_t j = 0; j < eval_at_.size(); ++j) {
            result_buffer_[tid][j] += tmp_dcg[j] * inverse_max_dcgs_[i][j] * query_weights_[i];
          }
Guolin Ke's avatar
Guolin Ke committed
132
133
        }
      }
134
135
    }
    // Get final average NDCG
136
    std::vector<double> result(eval_at_.size(), 0.0f);
137
    for (size_t j = 0; j < result.size(); ++j) {
138
      for (int i = 0; i < num_threads; ++i) {
139
        result[j] += result_buffer_[i][j];
wxchan's avatar
wxchan committed
140
      }
141
      result[j] /= sum_query_weights_;
Guolin Ke's avatar
Guolin Ke committed
142
    }
143
    return result;
Guolin Ke's avatar
Guolin Ke committed
144
145
  }

Nikita Titov's avatar
Nikita Titov committed
146
 private:
Guolin Ke's avatar
Guolin Ke committed
147
148
149
  /*! \brief Number of data */
  data_size_t num_data_;
  /*! \brief Pointer of label */
150
  const label_t* label_;
Guolin Ke's avatar
Guolin Ke committed
151
  /*! \brief Name of test set */
152
  std::vector<std::string> name_;
Guolin Ke's avatar
Guolin Ke committed
153
154
155
156
157
  /*! \brief Query boundaries information */
  const data_size_t* query_boundaries_;
  /*! \brief Number of queries */
  data_size_t num_queries_;
  /*! \brief Weights of queries */
158
  const label_t* query_weights_;
Guolin Ke's avatar
Guolin Ke committed
159
  /*! \brief Sum weights of queries */
160
  double sum_query_weights_;
Guolin Ke's avatar
Guolin Ke committed
161
162
163
  /*! \brief Evaluate position of NDCG */
  std::vector<data_size_t> eval_at_;
  /*! \brief Cache the inverse max dcg for all queries */
164
  std::vector<std::vector<double>> inverse_max_dcgs_;
Guolin Ke's avatar
Guolin Ke committed
165
166
167
168
};

}  // namespace LightGBM

169
#endif   // LIGHTGBM_SRC_METRIC_RANK_METRIC_HPP_