metadata.cpp 24.8 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
#include <LightGBM/dataset.h>
#include <LightGBM/utils/common.h>
Guolin Ke's avatar
Guolin Ke committed
7
8

#include <string>
9
#include <vector>
Guolin Ke's avatar
Guolin Ke committed
10
11
12

namespace LightGBM {

Guolin Ke's avatar
Guolin Ke committed
13
Metadata::Metadata() {
Guolin Ke's avatar
Guolin Ke committed
14
15
16
17
  num_weights_ = 0;
  num_init_score_ = 0;
  num_data_ = 0;
  num_queries_ = 0;
18
19
20
  weight_load_from_file_ = false;
  query_load_from_file_ = false;
  init_score_load_from_file_ = false;
21
  #ifdef USE_CUDA
22
  cuda_metadata_ = nullptr;
23
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
24
25
}

26
void Metadata::Init(const char* data_filename) {
Guolin Ke's avatar
Guolin Ke committed
27
  data_filename_ = data_filename;
28
  // for lambdarank, it needs query data for partition data in distributed learning
Guolin Ke's avatar
Guolin Ke committed
29
30
  LoadQueryBoundaries();
  LoadWeights();
31
  CalculateQueryWeights();
32
  LoadInitialScore(data_filename_);
Guolin Ke's avatar
Guolin Ke committed
33
34
35
36
37
}

Metadata::~Metadata() {
}

38
void Metadata::Init(data_size_t num_data, int weight_idx, int query_idx) {
Guolin Ke's avatar
Guolin Ke committed
39
  num_data_ = num_data;
40
  label_ = std::vector<label_t>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
41
  if (weight_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
42
    if (!weights_.empty()) {
43
      Log::Info("Using weights in data file, ignoring the additional weights file");
Guolin Ke's avatar
Guolin Ke committed
44
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
45
    }
46
    weights_ = std::vector<label_t>(num_data_, 0.0f);
Guolin Ke's avatar
Guolin Ke committed
47
    num_weights_ = num_data_;
48
    weight_load_from_file_ = false;
Guolin Ke's avatar
Guolin Ke committed
49
50
  }
  if (query_idx >= 0) {
Guolin Ke's avatar
Guolin Ke committed
51
    if (!query_boundaries_.empty()) {
52
      Log::Info("Using query id in data file, ignoring the additional query file");
Guolin Ke's avatar
Guolin Ke committed
53
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
54
    }
Guolin Ke's avatar
Guolin Ke committed
55
    if (!query_weights_.empty()) { query_weights_.clear(); }
56
    queries_ = std::vector<data_size_t>(num_data_, 0);
57
    query_load_from_file_ = false;
Guolin Ke's avatar
Guolin Ke committed
58
  }
Guolin Ke's avatar
Guolin Ke committed
59
60
}

61
62
63
64
void Metadata::InitByReference(data_size_t num_data, const Metadata* reference) {
  int has_weights = reference->num_weights_ > 0;
  int has_init_scores = reference->num_init_score_ > 0;
  int has_queries = reference->num_queries_ > 0;
65
  int nclasses = reference->num_init_score_classes();
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
  Init(num_data, has_weights, has_init_scores, has_queries, nclasses);
}

void Metadata::Init(data_size_t num_data, int32_t has_weights, int32_t has_init_scores, int32_t has_queries, int32_t nclasses) {
  num_data_ = num_data;
  label_ = std::vector<label_t>(num_data_);
  if (has_weights) {
    if (!weights_.empty()) {
      Log::Fatal("Calling Init() on Metadata weights that have already been initialized");
    }
    weights_.resize(num_data_, 0.0f);
    num_weights_ = num_data_;
    weight_load_from_file_ = false;
  }
  if (has_init_scores) {
    if (!init_score_.empty()) {
      Log::Fatal("Calling Init() on Metadata initial scores that have already been initialized");
    }
    num_init_score_ = static_cast<int64_t>(num_data) * nclasses;
    init_score_.resize(num_init_score_, 0);
  }
  if (has_queries) {
    if (!query_weights_.empty()) {
      Log::Fatal("Calling Init() on Metadata queries that have already been initialized");
    }
    queries_.resize(num_data_, 0);
    query_load_from_file_ = false;
  }
}

Guolin Ke's avatar
Guolin Ke committed
96
97
98
void Metadata::Init(const Metadata& fullset, const data_size_t* used_indices, data_size_t num_used_indices) {
  num_data_ = num_used_indices;

99
  label_ = std::vector<label_t>(num_used_indices);
Guolin Ke's avatar
Guolin Ke committed
100
#pragma omp parallel for schedule(static, 512) if (num_used_indices >= 1024)
101
  for (data_size_t i = 0; i < num_used_indices; ++i) {
Guolin Ke's avatar
Guolin Ke committed
102
103
104
    label_[i] = fullset.label_[used_indices[i]];
  }

Guolin Ke's avatar
Guolin Ke committed
105
  if (!fullset.weights_.empty()) {
106
    weights_ = std::vector<label_t>(num_used_indices);
Guolin Ke's avatar
Guolin Ke committed
107
    num_weights_ = num_used_indices;
Guolin Ke's avatar
Guolin Ke committed
108
#pragma omp parallel for schedule(static, 512) if (num_used_indices >= 1024)
109
    for (data_size_t i = 0; i < num_used_indices; ++i) {
Guolin Ke's avatar
Guolin Ke committed
110
111
112
113
114
115
      weights_[i] = fullset.weights_[used_indices[i]];
    }
  } else {
    num_weights_ = 0;
  }

Guolin Ke's avatar
Guolin Ke committed
116
  if (!fullset.init_score_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
117
    int num_class = static_cast<int>(fullset.num_init_score_ / fullset.num_data_);
118
    init_score_ = std::vector<double>(static_cast<size_t>(num_used_indices) * num_class);
Guolin Ke's avatar
Guolin Ke committed
119
    num_init_score_ = static_cast<int64_t>(num_used_indices) * num_class;
120
    #pragma omp parallel for schedule(static)
121
    for (int k = 0; k < num_class; ++k) {
122
123
124
125
      const size_t offset_dest = static_cast<size_t>(k) * num_data_;
      const size_t offset_src = static_cast<size_t>(k) * fullset.num_data_;
      for (data_size_t i = 0; i < num_used_indices; ++i) {
        init_score_[offset_dest + i] = fullset.init_score_[offset_src + used_indices[i]];
126
      }
Guolin Ke's avatar
Guolin Ke committed
127
128
129
130
131
    }
  } else {
    num_init_score_ = 0;
  }

Guolin Ke's avatar
Guolin Ke committed
132
  if (!fullset.query_boundaries_.empty()) {
Guolin Ke's avatar
Guolin Ke committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
    std::vector<data_size_t> used_query;
    data_size_t data_idx = 0;
    for (data_size_t qid = 0; qid < num_queries_ && data_idx < num_used_indices; ++qid) {
      data_size_t start = fullset.query_boundaries_[qid];
      data_size_t end = fullset.query_boundaries_[qid + 1];
      data_size_t len = end - start;
      if (used_indices[data_idx] > start) {
        continue;
      } else if (used_indices[data_idx] == start) {
        if (num_used_indices >= data_idx + len && used_indices[data_idx + len - 1] == end - 1) {
          used_query.push_back(qid);
          data_idx += len;
        } else {
          Log::Fatal("Data partition error, data didn't match queries");
        }
      } else {
        Log::Fatal("Data partition error, data didn't match queries");
      }
    }
    query_boundaries_ = std::vector<data_size_t>(used_query.size() + 1);
    num_queries_ = static_cast<data_size_t>(used_query.size());
    query_boundaries_[0] = 0;
    for (data_size_t i = 0; i < num_queries_; ++i) {
      data_size_t qid = used_query[i];
      data_size_t len = fullset.query_boundaries_[qid + 1] - fullset.query_boundaries_[qid];
      query_boundaries_[i + 1] = query_boundaries_[i] + len;
    }
  } else {
    num_queries_ = 0;
  }
}

Guolin Ke's avatar
Guolin Ke committed
165
void Metadata::PartitionLabel(const std::vector<data_size_t>& used_indices) {
Guolin Ke's avatar
Guolin Ke committed
166
  if (used_indices.empty()) {
Guolin Ke's avatar
Guolin Ke committed
167
168
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
169
  auto old_label = label_;
Guolin Ke's avatar
Guolin Ke committed
170
  num_data_ = static_cast<data_size_t>(used_indices.size());
171
  label_ = std::vector<label_t>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
172
#pragma omp parallel for schedule(static, 512) if (num_data_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
173
174
175
  for (data_size_t i = 0; i < num_data_; ++i) {
    label_[i] = old_label[used_indices[i]];
  }
Guolin Ke's avatar
Guolin Ke committed
176
  old_label.clear();
Guolin Ke's avatar
Guolin Ke committed
177
178
}

179
180
181
182
183
184
185
186
187
188
void Metadata::CalculateQueryBoundaries() {
  if (!queries_.empty()) {
    // need convert query_id to boundaries
    std::vector<data_size_t> tmp_buffer;
    data_size_t last_qid = -1;
    data_size_t cur_cnt = 0;
    for (data_size_t i = 0; i < num_data_; ++i) {
      if (last_qid != queries_[i]) {
        if (cur_cnt > 0) {
          tmp_buffer.push_back(cur_cnt);
Guolin Ke's avatar
Guolin Ke committed
189
        }
190
191
        cur_cnt = 0;
        last_qid = queries_[i];
Guolin Ke's avatar
Guolin Ke committed
192
      }
193
      ++cur_cnt;
Guolin Ke's avatar
Guolin Ke committed
194
    }
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
    tmp_buffer.push_back(cur_cnt);
    query_boundaries_ = std::vector<data_size_t>(tmp_buffer.size() + 1);
    num_queries_ = static_cast<data_size_t>(tmp_buffer.size());
    query_boundaries_[0] = 0;
    for (size_t i = 0; i < tmp_buffer.size(); ++i) {
      query_boundaries_[i + 1] = query_boundaries_[i] + tmp_buffer[i];
    }
    CalculateQueryWeights();
    queries_.clear();
  }
}

void Metadata::CheckOrPartition(data_size_t num_all_data, const std::vector<data_size_t>& used_data_indices) {
  if (used_data_indices.empty()) {
     CalculateQueryBoundaries();
Guolin Ke's avatar
Guolin Ke committed
210
    // check weights
Guolin Ke's avatar
Guolin Ke committed
211
    if (!weights_.empty() && num_weights_ != num_data_) {
Guolin Ke's avatar
Guolin Ke committed
212
      weights_.clear();
Guolin Ke's avatar
Guolin Ke committed
213
      num_weights_ = 0;
Guolin Ke's avatar
Guolin Ke committed
214
      Log::Fatal("Weights size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
215
216
217
    }

    // check query boundries
Guolin Ke's avatar
Guolin Ke committed
218
    if (!query_boundaries_.empty() && query_boundaries_[num_queries_] != num_data_) {
Guolin Ke's avatar
Guolin Ke committed
219
      query_boundaries_.clear();
Guolin Ke's avatar
Guolin Ke committed
220
      num_queries_ = 0;
Guolin Ke's avatar
Guolin Ke committed
221
      Log::Fatal("Query size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
222
223
224
    }

    // contain initial score file
225
    if (!init_score_.empty() && (num_init_score_ % num_data_) != 0) {
Guolin Ke's avatar
Guolin Ke committed
226
      init_score_.clear();
Guolin Ke's avatar
Guolin Ke committed
227
      num_init_score_ = 0;
Guolin Ke's avatar
Guolin Ke committed
228
      Log::Fatal("Initial score size doesn't match data size");
Guolin Ke's avatar
Guolin Ke committed
229
230
    }
  } else {
231
    if (!queries_.empty()) {
232
      Log::Fatal("Cannot used query_id for distributed training");
233
    }
Guolin Ke's avatar
Guolin Ke committed
234
235
    data_size_t num_used_data = static_cast<data_size_t>(used_data_indices.size());
    // check weights
236
237
238
239
240
241
242
243
244
245
    if (weight_load_from_file_) {
      if (weights_.size() > 0 && num_weights_ != num_all_data) {
        weights_.clear();
        num_weights_ = 0;
        Log::Fatal("Weights size doesn't match data size");
      }
      // get local weights
      if (!weights_.empty()) {
        auto old_weights = weights_;
        num_weights_ = num_data_;
246
        weights_ = std::vector<label_t>(num_data_);
Guolin Ke's avatar
Guolin Ke committed
247
#pragma omp parallel for schedule(static, 512)
248
249
250
251
        for (int i = 0; i < static_cast<int>(used_data_indices.size()); ++i) {
          weights_[i] = old_weights[used_data_indices[i]];
        }
        old_weights.clear();
Guolin Ke's avatar
Guolin Ke committed
252
253
      }
    }
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    if (query_load_from_file_) {
      // check query boundries
      if (!query_boundaries_.empty() && query_boundaries_[num_queries_] != num_all_data) {
        query_boundaries_.clear();
        num_queries_ = 0;
        Log::Fatal("Query size doesn't match data size");
      }
      // get local query boundaries
      if (!query_boundaries_.empty()) {
        std::vector<data_size_t> used_query;
        data_size_t data_idx = 0;
        for (data_size_t qid = 0; qid < num_queries_ && data_idx < num_used_data; ++qid) {
          data_size_t start = query_boundaries_[qid];
          data_size_t end = query_boundaries_[qid + 1];
          data_size_t len = end - start;
          if (used_data_indices[data_idx] > start) {
            continue;
          } else if (used_data_indices[data_idx] == start) {
            if (num_used_data >= data_idx + len && used_data_indices[data_idx + len - 1] == end - 1) {
              used_query.push_back(qid);
              data_idx += len;
            } else {
              Log::Fatal("Data partition error, data didn't match queries");
            }
Guolin Ke's avatar
Guolin Ke committed
278
          } else {
Guolin Ke's avatar
Guolin Ke committed
279
            Log::Fatal("Data partition error, data didn't match queries");
Guolin Ke's avatar
Guolin Ke committed
280
281
          }
        }
282
283
284
285
286
287
288
289
290
291
        auto old_query_boundaries = query_boundaries_;
        query_boundaries_ = std::vector<data_size_t>(used_query.size() + 1);
        num_queries_ = static_cast<data_size_t>(used_query.size());
        query_boundaries_[0] = 0;
        for (data_size_t i = 0; i < num_queries_; ++i) {
          data_size_t qid = used_query[i];
          data_size_t len = old_query_boundaries[qid + 1] - old_query_boundaries[qid];
          query_boundaries_[i + 1] = query_boundaries_[i] + len;
        }
        old_query_boundaries.clear();
Guolin Ke's avatar
Guolin Ke committed
292
293
      }
    }
294
295
296
297
298
299
300
    if (init_score_load_from_file_) {
      // contain initial score file
      if (!init_score_.empty() && (num_init_score_ % num_all_data) != 0) {
        init_score_.clear();
        num_init_score_ = 0;
        Log::Fatal("Initial score size doesn't match data size");
      }
Guolin Ke's avatar
Guolin Ke committed
301

302
303
304
305
306
307
      // get local initial scores
      if (!init_score_.empty()) {
        auto old_scores = init_score_;
        int num_class = static_cast<int>(num_init_score_ / num_all_data);
        num_init_score_ = static_cast<int64_t>(num_data_) * num_class;
        init_score_ = std::vector<double>(num_init_score_);
Guolin Ke's avatar
Guolin Ke committed
308
#pragma omp parallel for schedule(static)
309
        for (int k = 0; k < num_class; ++k) {
310
311
          const size_t offset_dest = static_cast<size_t>(k) * num_data_;
          const size_t offset_src = static_cast<size_t>(k) * num_all_data;
312
          for (size_t i = 0; i < used_data_indices.size(); ++i) {
313
            init_score_[offset_dest + i] = old_scores[offset_src + used_data_indices[i]];
314
          }
315
        }
316
        old_scores.clear();
Guolin Ke's avatar
Guolin Ke committed
317
318
      }
    }
319
320
    // re-calculate query weight
    CalculateQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
321
  }
322
323
324
325
  if (num_queries_ > 0) {
    Log::Debug("Number of queries in %s: %i. Average number of rows per query: %f.",
      data_filename_.c_str(), static_cast<int>(num_queries_), static_cast<double>(num_data_) / num_queries_);
  }
Guolin Ke's avatar
Guolin Ke committed
326
327
}

328
329
330
331
332
333
334
335
336
337
338
void Metadata::SetInitScore(const double* init_score, data_size_t len) {
  std::lock_guard<std::mutex> lock(mutex_);
  // save to nullptr
  if (init_score == nullptr || len == 0) {
    init_score_.clear();
    num_init_score_ = 0;
    return;
  }
  if ((len % num_data_) != 0) {
    Log::Fatal("Initial score size doesn't match data size");
  }
339
  if (init_score_.empty()) { init_score_.resize(len); }
340
  num_init_score_ = len;
341

Guolin Ke's avatar
Guolin Ke committed
342
  #pragma omp parallel for schedule(static, 512) if (num_init_score_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
343
  for (int64_t i = 0; i < num_init_score_; ++i) {
344
    init_score_[i] = Common::AvoidInf(init_score[i]);
345
  }
346
  init_score_load_from_file_ = false;
347
  #ifdef USE_CUDA
348
349
350
  if (cuda_metadata_ != nullptr) {
    cuda_metadata_->SetInitScore(init_score_.data(), len);
  }
351
  #endif  // USE_CUDA
352
353
}

354
355
356
357
358
359
360
361
362
363
void Metadata::InsertInitScores(const double* init_scores, data_size_t start_index, data_size_t len, data_size_t source_size) {
  if (num_init_score_ <= 0) {
    Log::Fatal("Inserting initial score data into dataset with no initial scores");
  }
  if (start_index + len > num_data_) {
    // Note that len here is row count, not num_init_score, so we compare against num_data
    Log::Fatal("Inserted initial score data is too large for dataset");
  }
  if (init_score_.empty()) { init_score_.resize(num_init_score_); }

364
  int nclasses = num_init_score_classes();
365
366
367
368
369
370
371
372
373
374
375

  for (int32_t col = 0; col < nclasses; ++col) {
    int32_t dest_offset = num_data_ * col + start_index;
    // We need to use source_size here, because len might not equal size (due to a partially loaded dataset)
    int32_t source_offset = source_size * col;
    memcpy(init_score_.data() + dest_offset, init_scores + source_offset, sizeof(double) * len);
  }
  init_score_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

376
void Metadata::SetLabel(const label_t* label, data_size_t len) {
377
  std::lock_guard<std::mutex> lock(mutex_);
378
379
380
  if (label == nullptr) {
    Log::Fatal("label cannot be nullptr");
  }
Guolin Ke's avatar
Guolin Ke committed
381
  if (num_data_ != len) {
382
    Log::Fatal("Length of label is not same with #data");
Guolin Ke's avatar
Guolin Ke committed
383
  }
384
  if (label_.empty()) { label_.resize(num_data_); }
385

Guolin Ke's avatar
Guolin Ke committed
386
  #pragma omp parallel for schedule(static, 512) if (num_data_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
387
  for (data_size_t i = 0; i < num_data_; ++i) {
388
    label_[i] = Common::AvoidInf(label[i]);
Guolin Ke's avatar
Guolin Ke committed
389
  }
390
  #ifdef USE_CUDA
391
392
393
  if (cuda_metadata_ != nullptr) {
    cuda_metadata_->SetLabel(label_.data(), len);
  }
394
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
395
396
}

397
398
399
400
401
402
403
404
405
406
407
408
409
410
void Metadata::InsertLabels(const label_t* labels, data_size_t start_index, data_size_t len) {
  if (labels == nullptr) {
    Log::Fatal("label cannot be nullptr");
  }
  if (start_index + len > num_data_) {
    Log::Fatal("Inserted label data is too large for dataset");
  }
  if (label_.empty()) { label_.resize(num_data_); }

  memcpy(label_.data() + start_index, labels, sizeof(label_t) * len);

  // CUDA is handled after all insertions are complete
}

411
void Metadata::SetWeights(const label_t* weights, data_size_t len) {
412
  std::lock_guard<std::mutex> lock(mutex_);
413
414
415
416
417
418
  // save to nullptr
  if (weights == nullptr || len == 0) {
    weights_.clear();
    num_weights_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
419
  if (num_data_ != len) {
420
    Log::Fatal("Length of weights is not same with #data");
Guolin Ke's avatar
Guolin Ke committed
421
  }
422
  if (weights_.empty()) { weights_.resize(num_data_); }
Guolin Ke's avatar
Guolin Ke committed
423
  num_weights_ = num_data_;
424

Guolin Ke's avatar
Guolin Ke committed
425
  #pragma omp parallel for schedule(static, 512) if (num_weights_ >= 1024)
Guolin Ke's avatar
Guolin Ke committed
426
  for (data_size_t i = 0; i < num_weights_; ++i) {
427
    weights_[i] = Common::AvoidInf(weights[i]);
Guolin Ke's avatar
Guolin Ke committed
428
  }
429
  CalculateQueryWeights();
430
  weight_load_from_file_ = false;
431
  #ifdef USE_CUDA
432
433
434
  if (cuda_metadata_ != nullptr) {
    cuda_metadata_->SetWeights(weights_.data(), len);
  }
435
  #endif  // USE_CUDA
Guolin Ke's avatar
Guolin Ke committed
436
437
}

438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
void Metadata::InsertWeights(const label_t* weights, data_size_t start_index, data_size_t len) {
  if (!weights) {
    Log::Fatal("Passed null weights");
  }
  if (num_weights_ <= 0) {
    Log::Fatal("Inserting weight data into dataset with no weights");
  }
  if (start_index + len > num_weights_) {
    Log::Fatal("Inserted weight data is too large for dataset");
  }
  if (weights_.empty()) { weights_.resize(num_weights_); }

  memcpy(weights_.data() + start_index, weights, sizeof(label_t) * len);

  weight_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

Guolin Ke's avatar
Guolin Ke committed
456
void Metadata::SetQuery(const data_size_t* query, data_size_t len) {
457
  std::lock_guard<std::mutex> lock(mutex_);
458
  // save to nullptr
Guolin Ke's avatar
Guolin Ke committed
459
  if (query == nullptr || len == 0) {
460
461
462
463
    query_boundaries_.clear();
    num_queries_ = 0;
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
464
  data_size_t sum = 0;
465
  #pragma omp parallel for schedule(static) reduction(+:sum)
Guolin Ke's avatar
Guolin Ke committed
466
  for (data_size_t i = 0; i < len; ++i) {
Guolin Ke's avatar
Guolin Ke committed
467
    sum += query[i];
Guolin Ke's avatar
Guolin Ke committed
468
469
  }
  if (num_data_ != sum) {
470
    Log::Fatal("Sum of query counts is not same with #data");
Guolin Ke's avatar
Guolin Ke committed
471
472
  }
  num_queries_ = len;
473
  query_boundaries_.resize(num_queries_ + 1);
Guolin Ke's avatar
Guolin Ke committed
474
  query_boundaries_[0] = 0;
Guolin Ke's avatar
Guolin Ke committed
475
  for (data_size_t i = 0; i < num_queries_; ++i) {
Guolin Ke's avatar
Guolin Ke committed
476
    query_boundaries_[i + 1] = query_boundaries_[i] + query[i];
Guolin Ke's avatar
Guolin Ke committed
477
  }
478
  CalculateQueryWeights();
479
  query_load_from_file_ = false;
480
  #ifdef USE_CUDA
481
482
483
484
485
486
487
488
  if (cuda_metadata_ != nullptr) {
    if (query_weights_.size() > 0) {
      CHECK_EQ(query_weights_.size(), static_cast<size_t>(num_queries_));
      cuda_metadata_->SetQuery(query_boundaries_.data(), query_weights_.data(), num_queries_);
    } else {
      cuda_metadata_->SetQuery(query_boundaries_.data(), nullptr, num_queries_);
    }
  }
489
  #endif  // USE_CUDA
490
}
Guolin Ke's avatar
Guolin Ke committed
491

492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
void Metadata::InsertQueries(const data_size_t* queries, data_size_t start_index, data_size_t len) {
  if (!queries) {
    Log::Fatal("Passed null queries");
  }
  if (queries_.size() <= 0) {
    Log::Fatal("Inserting query data into dataset with no queries");
  }
  if (static_cast<size_t>(start_index + len) > queries_.size()) {
    Log::Fatal("Inserted query data is too large for dataset");
  }

  memcpy(queries_.data() + start_index, queries, sizeof(data_size_t) * len);

  query_load_from_file_ = false;
  // CUDA is handled after all insertions are complete
}

Guolin Ke's avatar
Guolin Ke committed
509
510
511
512
513
void Metadata::LoadWeights() {
  num_weights_ = 0;
  std::string weight_filename(data_filename_);
  // default weight file name
  weight_filename.append(".weight");
Guolin Ke's avatar
Guolin Ke committed
514
  TextReader<size_t> reader(weight_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
515
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
516
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
517
518
    return;
  }
519
  Log::Info("Loading weights...");
Guolin Ke's avatar
Guolin Ke committed
520
  num_weights_ = static_cast<data_size_t>(reader.Lines().size());
521
  weights_ = std::vector<label_t>(num_weights_);
522
  #pragma omp parallel for schedule(static)
Guolin Ke's avatar
Guolin Ke committed
523
  for (data_size_t i = 0; i < num_weights_; ++i) {
524
    double tmp_weight = 0.0f;
Guolin Ke's avatar
Guolin Ke committed
525
    Common::Atof(reader.Lines()[i].c_str(), &tmp_weight);
526
    weights_[i] = Common::AvoidInf(static_cast<label_t>(tmp_weight));
Guolin Ke's avatar
Guolin Ke committed
527
  }
528
  weight_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
529
530
}

531
void Metadata::LoadInitialScore(const std::string& data_filename) {
Guolin Ke's avatar
Guolin Ke committed
532
  num_init_score_ = 0;
533
534
  std::string init_score_filename(data_filename);
  init_score_filename = std::string(data_filename);
535
536
  // default init_score file name
  init_score_filename.append(".init");
Guolin Ke's avatar
Guolin Ke committed
537
  TextReader<size_t> reader(init_score_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
538
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
539
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
540
541
    return;
  }
542
543
  Log::Info("Loading initial scores...");

544
545
546
  // use first line to count number class
  int num_class = static_cast<int>(Common::Split(reader.Lines()[0].c_str(), '\t').size());
  data_size_t num_line = static_cast<data_size_t>(reader.Lines().size());
Guolin Ke's avatar
Guolin Ke committed
547
  num_init_score_ = static_cast<int64_t>(num_line) * num_class;
548

Guolin Ke's avatar
Guolin Ke committed
549
  init_score_ = std::vector<double>(num_init_score_);
550
  if (num_class == 1) {
551
    #pragma omp parallel for schedule(static)
552
    for (data_size_t i = 0; i < num_line; ++i) {
Guolin Ke's avatar
Guolin Ke committed
553
      double tmp = 0.0f;
554
      Common::Atof(reader.Lines()[i].c_str(), &tmp);
555
      init_score_[i] = Common::AvoidInf(static_cast<double>(tmp));
556
    }
557
  } else {
558
    std::vector<std::string> oneline_init_score;
559
    #pragma omp parallel for schedule(static)
560
    for (data_size_t i = 0; i < num_line; ++i) {
Guolin Ke's avatar
Guolin Ke committed
561
      double tmp = 0.0f;
562
563
      oneline_init_score = Common::Split(reader.Lines()[i].c_str(), '\t');
      if (static_cast<int>(oneline_init_score.size()) != num_class) {
564
        Log::Fatal("Invalid initial score file. Redundant or insufficient columns");
565
566
567
      }
      for (int k = 0; k < num_class; ++k) {
        Common::Atof(oneline_init_score[k].c_str(), &tmp);
568
        init_score_[static_cast<size_t>(k) * num_line + i] = Common::AvoidInf(static_cast<double>(tmp));
569
      }
570
    }
Guolin Ke's avatar
Guolin Ke committed
571
  }
572
  init_score_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
573
574
575
576
577
578
579
}

void Metadata::LoadQueryBoundaries() {
  num_queries_ = 0;
  std::string query_filename(data_filename_);
  // default query file name
  query_filename.append(".query");
Guolin Ke's avatar
Guolin Ke committed
580
  TextReader<size_t> reader(query_filename.c_str(), false);
Guolin Ke's avatar
Guolin Ke committed
581
  reader.ReadAllLines();
Guolin Ke's avatar
Guolin Ke committed
582
  if (reader.Lines().empty()) {
Guolin Ke's avatar
Guolin Ke committed
583
584
    return;
  }
585
  Log::Info("Calculating query boundaries...");
Guolin Ke's avatar
Guolin Ke committed
586
  query_boundaries_ = std::vector<data_size_t>(reader.Lines().size() + 1);
Guolin Ke's avatar
Guolin Ke committed
587
588
589
590
591
592
593
  num_queries_ = static_cast<data_size_t>(reader.Lines().size());
  query_boundaries_[0] = 0;
  for (size_t i = 0; i < reader.Lines().size(); ++i) {
    int tmp_cnt;
    Common::Atoi(reader.Lines()[i].c_str(), &tmp_cnt);
    query_boundaries_[i + 1] = query_boundaries_[i] + static_cast<data_size_t>(tmp_cnt);
  }
594
  query_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
595
596
}

597
void Metadata::CalculateQueryWeights() {
Guolin Ke's avatar
Guolin Ke committed
598
  if (weights_.size() == 0 || query_boundaries_.size() == 0) {
Guolin Ke's avatar
Guolin Ke committed
599
600
    return;
  }
Guolin Ke's avatar
Guolin Ke committed
601
  query_weights_.clear();
602
  Log::Info("Calculating query weights...");
603
  query_weights_ = std::vector<label_t>(num_queries_);
Guolin Ke's avatar
Guolin Ke committed
604
605
606
607
608
609
610
611
612
  for (data_size_t i = 0; i < num_queries_; ++i) {
    query_weights_[i] = 0.0f;
    for (data_size_t j = query_boundaries_[i]; j < query_boundaries_[i + 1]; ++j) {
      query_weights_[i] += weights_[j];
    }
    query_weights_[i] /= (query_boundaries_[i + 1] - query_boundaries_[i]);
  }
}

613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
void Metadata::InsertAt(data_size_t start_index,
  data_size_t count,
  const float* labels,
  const float* weights,
  const double* init_scores,
  const int32_t* queries) {
  if (num_data_ < count + start_index) {
    Log::Fatal("Length of metadata is too long to append #data");
  }
  InsertLabels(labels, start_index, count);
  if (weights) {
    InsertWeights(weights, start_index, count);
  }
  if (init_scores) {
    InsertInitScores(init_scores, start_index, count, count);
  }
  if (queries) {
    InsertQueries(queries, start_index, count);
  }
}

void Metadata::FinishLoad() {
  CalculateQueryBoundaries();
}

638
#ifdef USE_CUDA
639
640
641
642
void Metadata::CreateCUDAMetadata(const int gpu_device_id) {
  cuda_metadata_.reset(new CUDAMetadata(gpu_device_id));
  cuda_metadata_->Init(label_, weights_, query_boundaries_, query_weights_, init_score_);
}
643
#endif  // USE_CUDA
644

Guolin Ke's avatar
Guolin Ke committed
645
646
647
648
void Metadata::LoadFromMemory(const void* memory) {
  const char* mem_ptr = reinterpret_cast<const char*>(memory);

  num_data_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
649
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_data_));
Guolin Ke's avatar
Guolin Ke committed
650
  num_weights_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
651
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_weights_));
Guolin Ke's avatar
Guolin Ke committed
652
  num_queries_ = *(reinterpret_cast<const data_size_t*>(mem_ptr));
653
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(num_queries_));
Guolin Ke's avatar
Guolin Ke committed
654

Guolin Ke's avatar
Guolin Ke committed
655
  if (!label_.empty()) { label_.clear(); }
656
  label_ = std::vector<label_t>(num_data_);
657
  std::memcpy(label_.data(), mem_ptr, sizeof(label_t) * num_data_);
658
  mem_ptr += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
659
660

  if (num_weights_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
661
    if (!weights_.empty()) { weights_.clear(); }
662
    weights_ = std::vector<label_t>(num_weights_);
663
    std::memcpy(weights_.data(), mem_ptr, sizeof(label_t) * num_weights_);
664
    mem_ptr += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_weights_);
665
    weight_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
666
667
  }
  if (num_queries_ > 0) {
Guolin Ke's avatar
Guolin Ke committed
668
    if (!query_boundaries_.empty()) { query_boundaries_.clear(); }
Guolin Ke's avatar
Guolin Ke committed
669
    query_boundaries_ = std::vector<data_size_t>(num_queries_ + 1);
670
    std::memcpy(query_boundaries_.data(), mem_ptr, sizeof(data_size_t) * (num_queries_ + 1));
671
672
    mem_ptr += VirtualFileWriter::AlignedSize(sizeof(data_size_t) *
                                              (num_queries_ + 1));
673
    query_load_from_file_ = true;
Guolin Ke's avatar
Guolin Ke committed
674
  }
675
  CalculateQueryWeights();
Guolin Ke's avatar
Guolin Ke committed
676
677
}

678
void Metadata::SaveBinaryToFile(BinaryWriter* writer) const {
679
680
681
682
  writer->AlignedWrite(&num_data_, sizeof(num_data_));
  writer->AlignedWrite(&num_weights_, sizeof(num_weights_));
  writer->AlignedWrite(&num_queries_, sizeof(num_queries_));
  writer->AlignedWrite(label_.data(), sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
683
  if (!weights_.empty()) {
684
    writer->AlignedWrite(weights_.data(), sizeof(label_t) * num_weights_);
Guolin Ke's avatar
Guolin Ke committed
685
  }
Guolin Ke's avatar
Guolin Ke committed
686
  if (!query_boundaries_.empty()) {
687
688
    writer->AlignedWrite(query_boundaries_.data(),
                         sizeof(data_size_t) * (num_queries_ + 1));
Guolin Ke's avatar
Guolin Ke committed
689
  }
690
691
692
693
  if (num_init_score_ > 0) {
    Log::Warning("Please note that `init_score` is not saved in binary file.\n"
      "If you need it, please set it again after loading Dataset.");
  }
Guolin Ke's avatar
Guolin Ke committed
694
695
}

696
size_t Metadata::SizesInByte() const {
697
698
699
700
  size_t size = VirtualFileWriter::AlignedSize(sizeof(num_data_)) +
                VirtualFileWriter::AlignedSize(sizeof(num_weights_)) +
                VirtualFileWriter::AlignedSize(sizeof(num_queries_));
  size += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_data_);
Guolin Ke's avatar
Guolin Ke committed
701
  if (!weights_.empty()) {
702
    size += VirtualFileWriter::AlignedSize(sizeof(label_t) * num_weights_);
Guolin Ke's avatar
Guolin Ke committed
703
  }
Guolin Ke's avatar
Guolin Ke committed
704
  if (!query_boundaries_.empty()) {
705
706
    size += VirtualFileWriter::AlignedSize(sizeof(data_size_t) *
                                           (num_queries_ + 1));
Guolin Ke's avatar
Guolin Ke committed
707
708
709
710
711
712
  }
  return size;
}


}  // namespace LightGBM