blob: f858e68756e9faaa96e8c452a5b5460dce765f4b (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
|
// Copyright 2019 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#include "media/learning/impl/lookup_table_trainer.h"
#include <map>
namespace media {
namespace learning {
class LookupTable : public Model {
public:
LookupTable(const TrainingData& training_data) {
for (auto& example : training_data)
buckets_[example.features] += example;
}
// Model
TargetHistogram PredictDistribution(const FeatureVector& instance) override {
auto iter = buckets_.find(instance);
if (iter == buckets_.end())
return TargetHistogram();
return iter->second;
}
private:
std::map<FeatureVector, TargetHistogram> buckets_;
};
LookupTableTrainer::LookupTableTrainer() = default;
LookupTableTrainer::~LookupTableTrainer() = default;
void LookupTableTrainer::Train(const LearningTask& task,
const TrainingData& training_data,
TrainedModelCB model_cb) {
std::unique_ptr<LookupTable> lookup_table =
std::make_unique<LookupTable>(training_data);
// TODO(liberato): post?
std::move(model_cb).Run(std::move(lookup_table));
}
} // namespace learning
} // namespace media
|
