| // Copyright 2018 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/learning_task_controller_impl.h" |
| |
| #include <memory> |
| #include <utility> |
| #include <vector> |
| |
| #include "base/bind.h" |
| #include "base/check_op.h" |
| #include "base/notreached.h" |
| #include "media/learning/impl/distribution_reporter.h" |
| #include "media/learning/impl/extra_trees_trainer.h" |
| #include "media/learning/impl/lookup_table_trainer.h" |
| |
| namespace media { |
| namespace learning { |
| |
| LearningTaskControllerImpl::LearningTaskControllerImpl( |
| const LearningTask& task, |
| std::unique_ptr<DistributionReporter> reporter, |
| SequenceBoundFeatureProvider feature_provider) |
| : task_(task), |
| training_data_(std::make_unique<TrainingData>()), |
| reporter_(std::move(reporter)), |
| helper_(std::make_unique<LearningTaskControllerHelper>( |
| task, |
| base::BindRepeating(&LearningTaskControllerImpl::AddFinishedExample, |
| AsWeakPtr()), |
| std::move(feature_provider))), |
| expected_feature_count_(task_.feature_descriptions.size()) { |
| // Note that |helper_| uses the full set of features. |
| |
| // TODO(liberato): Make this compositional. FeatureSubsetTaskController? |
| if (task_.feature_subset_size) |
| DoFeatureSubsetSelection(); |
| |
| switch (task_.model) { |
| case LearningTask::Model::kExtraTrees: |
| trainer_ = std::make_unique<ExtraTreesTrainer>(); |
| break; |
| case LearningTask::Model::kLookupTable: |
| trainer_ = std::make_unique<LookupTableTrainer>(); |
| break; |
| } |
| } |
| |
| LearningTaskControllerImpl::~LearningTaskControllerImpl() = default; |
| |
| void LearningTaskControllerImpl::BeginObservation( |
| base::UnguessableToken id, |
| const FeatureVector& features, |
| const absl::optional<TargetValue>& default_target, |
| const absl::optional<ukm::SourceId>& source_id) { |
| // TODO(liberato): Should we enforce that the right number of features are |
| // present here? Right now, we allow it to be shorter, so that features from |
| // a FeatureProvider may be omitted. Of course, they have to be at the end in |
| // that case. If we start enforcing it here, make sure that LearningHelper |
| // starts adding the placeholder features. |
| if (!trainer_) |
| return; |
| |
| // We don't support default targets, since we're the base learner and can't |
| // easily do that. However, defaults are handled by (weak) controllers |
| // handed out by LearningSessionImpl. So, we don't bother since they never |
| // get here anyway. |
| DCHECK(!default_target); |
| |
| helper_->BeginObservation(id, features, source_id); |
| } |
| |
| void LearningTaskControllerImpl::CompleteObservation( |
| base::UnguessableToken id, |
| const ObservationCompletion& completion) { |
| if (!trainer_) |
| return; |
| helper_->CompleteObservation(id, completion); |
| } |
| |
| void LearningTaskControllerImpl::CancelObservation(base::UnguessableToken id) { |
| if (!trainer_) |
| return; |
| helper_->CancelObservation(id); |
| } |
| |
| void LearningTaskControllerImpl::UpdateDefaultTarget( |
| base::UnguessableToken id, |
| const absl::optional<TargetValue>& default_target) { |
| NOTREACHED(); |
| } |
| |
| const LearningTask& LearningTaskControllerImpl::GetLearningTask() { |
| return task_; |
| } |
| |
| void LearningTaskControllerImpl::PredictDistribution( |
| const FeatureVector& features, |
| PredictionCB callback) { |
| if (model_) |
| std::move(callback).Run(model_->PredictDistribution(features)); |
| else |
| std::move(callback).Run(absl::nullopt); |
| } |
| |
| void LearningTaskControllerImpl::AddFinishedExample(LabelledExample example, |
| ukm::SourceId source_id) { |
| // Verify that we have a trainer and that we got the right number of features. |
| // We don't compare to |task_.feature_descriptions.size()| since that has been |
| // adjusted to the subset size already. We expect the original count. |
| if (!trainer_ || example.features.size() != expected_feature_count_) |
| return; |
| |
| // Now that we have the whole set of features, select the subset we want. |
| FeatureVector new_features; |
| if (task_.feature_subset_size) { |
| for (auto& iter : feature_indices_) |
| new_features.push_back(example.features[iter]); |
| example.features = std::move(new_features); |
| } // else use them all. |
| |
| // The features should now match the task. |
| DCHECK_EQ(example.features.size(), task_.feature_descriptions.size()); |
| |
| if (training_data_->size() >= task_.max_data_set_size) { |
| // Replace a random example. We don't necessarily want to replace the |
| // oldest, since we don't necessarily want to enforce an ad-hoc recency |
| // constraint here. That's a different issue. |
| (*training_data_)[rng()->Generate(training_data_->size())] = example; |
| } else { |
| training_data_->push_back(example); |
| } |
| // Either way, we have one more example that we haven't used for training yet. |
| num_untrained_examples_++; |
| |
| // Once we have a model, see if we'd get |example| correct. |
| if (model_ && reporter_) { |
| TargetHistogram predicted = model_->PredictDistribution(example.features); |
| |
| DistributionReporter::PredictionInfo info; |
| info.observed = example.target_value; |
| info.source_id = source_id; |
| info.total_training_weight = last_training_weight_; |
| info.total_training_examples = last_training_size_; |
| reporter_->GetPredictionCallback(info).Run(predicted); |
| } |
| |
| // Can't train more than one model concurrently. |
| if (training_is_in_progress_) |
| return; |
| |
| // Train every time we get enough new examples. Note that this works even if |
| // we are replacing old examples rather than adding new ones. |
| double frac = ((double)num_untrained_examples_) / training_data_->size(); |
| if (frac < task_.min_new_data_fraction) |
| return; |
| |
| num_untrained_examples_ = 0; |
| |
| // Record these for metrics. |
| last_training_weight_ = training_data_->total_weight(); |
| last_training_size_ = training_data_->size(); |
| |
| TrainedModelCB model_cb = |
| base::BindOnce(&LearningTaskControllerImpl::OnModelTrained, AsWeakPtr(), |
| training_data_->total_weight(), training_data_->size()); |
| training_is_in_progress_ = true; |
| // Note that this copies the training data, so it's okay if we add more |
| // examples to our copy before this returns. |
| // TODO(liberato): Post to a background task runner, and bind |model_cb| to |
| // the current one. Be careful about ownership if we invalidate |trainer_| |
| // on this thread. Be sure to post destruction to that sequence. |
| trainer_->Train(task_, *training_data_, std::move(model_cb)); |
| } |
| |
| void LearningTaskControllerImpl::OnModelTrained(double training_weight, |
| int training_size, |
| std::unique_ptr<Model> model) { |
| DCHECK(training_is_in_progress_); |
| training_is_in_progress_ = false; |
| model_ = std::move(model); |
| // Record these for metrics. |
| last_training_weight_ = training_weight; |
| last_training_size_ = training_size; |
| } |
| |
| void LearningTaskControllerImpl::SetTrainerForTesting( |
| std::unique_ptr<TrainingAlgorithm> trainer) { |
| trainer_ = std::move(trainer); |
| } |
| |
| void LearningTaskControllerImpl::DoFeatureSubsetSelection() { |
| // Choose a random feature, and trim the descriptions to match. |
| std::vector<size_t> features; |
| for (size_t i = 0; i < task_.feature_descriptions.size(); i++) |
| features.push_back(i); |
| |
| for (int i = 0; i < *task_.feature_subset_size; i++) { |
| // Pick an element from |i| to the end of the list, inclusive. |
| // TODO(liberato): For tests, this will happen before any rng is provided |
| // by the test; we'll use an actual rng. |
| int r = rng()->Generate(features.size() - i) + i; |
| // Swap them. |
| std::swap(features[i], features[r]); |
| } |
| |
| // Construct the feature subset from the first few elements. Also adjust the |
| // task's descriptions to match. We do this in two steps so that the |
| // descriptions are added via iterating over |feature_indices_|, so that the |
| // enumeration order is the same as when we adjust the feature values of |
| // incoming examples. In both cases, we iterate over |feature_indicies_|, |
| // which might (will) re-order them with respect to |features|. |
| for (int i = 0; i < *task_.feature_subset_size; i++) |
| feature_indices_.insert(features[i]); |
| |
| std::vector<LearningTask::ValueDescription> adjusted_descriptions; |
| for (auto& iter : feature_indices_) |
| adjusted_descriptions.push_back(task_.feature_descriptions[iter]); |
| |
| task_.feature_descriptions = adjusted_descriptions; |
| |
| if (reporter_) |
| reporter_->SetFeatureSubset(feature_indices_); |
| } |
| |
| } // namespace learning |
| } // namespace media |