learning_task_controller_impl_unittest.cc 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305
  1. // Copyright 2018 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "media/learning/impl/learning_task_controller_impl.h"
  5. #include <utility>
  6. #include "base/bind.h"
  7. #include "base/memory/raw_ptr.h"
  8. #include "base/test/task_environment.h"
  9. #include "base/threading/sequenced_task_runner_handle.h"
  10. #include "media/learning/impl/distribution_reporter.h"
  11. #include "testing/gtest/include/gtest/gtest.h"
  12. namespace media {
  13. namespace learning {
  14. class LearningTaskControllerImplTest : public testing::Test {
  15. public:
  16. class FakeDistributionReporter : public DistributionReporter {
  17. public:
  18. FakeDistributionReporter(const LearningTask& task)
  19. : DistributionReporter(task) {}
  20. // protected => public
  21. const absl::optional<std::set<int>>& feature_indices() const {
  22. return DistributionReporter::feature_indices();
  23. }
  24. protected:
  25. void OnPrediction(const PredictionInfo& info,
  26. TargetHistogram predicted) override {
  27. num_reported_++;
  28. TargetHistogram dist;
  29. dist += info.observed;
  30. if (dist == predicted)
  31. num_correct_++;
  32. most_recent_source_id_ = info.source_id;
  33. }
  34. public:
  35. int num_reported_ = 0;
  36. int num_correct_ = 0;
  37. ukm::SourceId most_recent_source_id_;
  38. };
  39. // Model that always predicts a constant.
  40. class FakeModel : public Model {
  41. public:
  42. FakeModel(TargetValue target) : target_(target) {}
  43. // Model
  44. TargetHistogram PredictDistribution(
  45. const FeatureVector& features) override {
  46. TargetHistogram dist;
  47. dist += target_;
  48. return dist;
  49. }
  50. private:
  51. // The value we predict.
  52. TargetValue target_;
  53. };
  54. class FakeTrainer : public TrainingAlgorithm {
  55. public:
  56. // |num_models| is where we'll record how many models we've trained.
  57. // |target_value| is the prediction that our trained model will make.
  58. FakeTrainer(int* num_models, TargetValue target_value)
  59. : num_models_(num_models), target_value_(target_value) {}
  60. ~FakeTrainer() override {}
  61. void Train(const LearningTask& task,
  62. const TrainingData& training_data,
  63. TrainedModelCB model_cb) override {
  64. task_ = task;
  65. (*num_models_)++;
  66. training_data_ = training_data;
  67. std::move(model_cb).Run(std::make_unique<FakeModel>(target_value_));
  68. }
  69. const LearningTask& task() const { return task_; }
  70. const TrainingData& training_data() const { return training_data_; }
  71. private:
  72. LearningTask task_;
  73. raw_ptr<int> num_models_ = nullptr;
  74. TargetValue target_value_;
  75. // Most recently provided training data.
  76. TrainingData training_data_;
  77. };
  78. // Increments feature 0.
  79. class FakeFeatureProvider : public FeatureProvider {
  80. public:
  81. void AddFeatures(FeatureVector features, FeatureVectorCB cb) override {
  82. features[0] = FeatureValue(features[0].value() + 1);
  83. std::move(cb).Run(features);
  84. }
  85. };
  86. LearningTaskControllerImplTest()
  87. : predicted_target_(123), not_predicted_target_(456) {
  88. // Set the name so that we can check it later.
  89. task_.name = "TestTask";
  90. // Don't require too many training examples per report.
  91. task_.max_data_set_size = 20;
  92. task_.min_new_data_fraction = 0.1;
  93. }
  94. ~LearningTaskControllerImplTest() override {
  95. // To prevent a memory leak, reset the controller. This may post
  96. // destruction of other objects, so RunUntilIdle().
  97. controller_.reset();
  98. task_environment_.RunUntilIdle();
  99. }
  100. void CreateController(SequenceBoundFeatureProvider feature_provider =
  101. SequenceBoundFeatureProvider()) {
  102. std::unique_ptr<FakeDistributionReporter> reporter =
  103. std::make_unique<FakeDistributionReporter>(task_);
  104. reporter_raw_ = reporter.get();
  105. controller_ = std::make_unique<LearningTaskControllerImpl>(
  106. task_, std::move(reporter), std::move(feature_provider));
  107. auto fake_trainer =
  108. std::make_unique<FakeTrainer>(&num_models_, predicted_target_);
  109. trainer_raw_ = fake_trainer.get();
  110. controller_->SetTrainerForTesting(std::move(fake_trainer));
  111. }
  112. void AddExample(const LabelledExample& example,
  113. absl::optional<ukm::SourceId> source_id = absl::nullopt) {
  114. base::UnguessableToken id = base::UnguessableToken::Create();
  115. controller_->BeginObservation(id, example.features, absl::nullopt,
  116. source_id);
  117. controller_->CompleteObservation(
  118. id, ObservationCompletion(example.target_value, example.weight));
  119. }
  120. void VerifyPrediction(const FeatureVector& features,
  121. absl::optional<TargetHistogram> expectation) {
  122. absl::optional<TargetHistogram> observed_prediction;
  123. controller_->PredictDistribution(
  124. features, base::BindOnce(
  125. [](absl::optional<TargetHistogram>* test_storage,
  126. const absl::optional<TargetHistogram>& predicted) {
  127. *test_storage = predicted;
  128. },
  129. &observed_prediction));
  130. task_environment_.RunUntilIdle();
  131. EXPECT_EQ(observed_prediction, expectation);
  132. }
  133. base::test::TaskEnvironment task_environment_;
  134. // Number of models that we trained.
  135. int num_models_ = 0;
  136. // Two distinct targets.
  137. const TargetValue predicted_target_;
  138. const TargetValue not_predicted_target_;
  139. raw_ptr<FakeDistributionReporter> reporter_raw_ = nullptr;
  140. raw_ptr<FakeTrainer> trainer_raw_ = nullptr;
  141. LearningTask task_;
  142. std::unique_ptr<LearningTaskControllerImpl> controller_;
  143. };
  144. TEST_F(LearningTaskControllerImplTest, AddingExamplesTrainsModelAndReports) {
  145. CreateController();
  146. LabelledExample example;
  147. // Up to the first 1/training_fraction examples should train on each example.
  148. // Make each of the examples agree on |predicted_target_|.
  149. example.target_value = predicted_target_;
  150. int count = static_cast<int>(1.0 / task_.min_new_data_fraction);
  151. for (int i = 0; i < count; i++) {
  152. AddExample(example);
  153. EXPECT_EQ(num_models_, i + 1);
  154. // All examples except the first should be reported as correct. For the
  155. // first, there's no model to test again.
  156. EXPECT_EQ(reporter_raw_->num_reported_, i);
  157. EXPECT_EQ(reporter_raw_->num_correct_, i);
  158. }
  159. // The next |count| should train every other one.
  160. for (int i = 0; i < count; i++) {
  161. AddExample(example);
  162. EXPECT_EQ(num_models_, count + (i + 1) / 2);
  163. }
  164. // The next |count| should be the same, since we've reached the max training
  165. // set size.
  166. for (int i = 0; i < count; i++) {
  167. AddExample(example);
  168. EXPECT_EQ(num_models_, count + count / 2 + (i + 1) / 2);
  169. }
  170. // We should have reported results for each except the first. All of them
  171. // should be correct, since there's only one target so far.
  172. EXPECT_EQ(reporter_raw_->num_reported_, count * 3 - 1);
  173. EXPECT_EQ(reporter_raw_->num_correct_, count * 3 - 1);
  174. // Adding a value that doesn't match should report one more attempt, with an
  175. // incorrect prediction.
  176. example.target_value = not_predicted_target_;
  177. AddExample(example);
  178. EXPECT_EQ(reporter_raw_->num_reported_, count * 3);
  179. EXPECT_EQ(reporter_raw_->num_correct_, count * 3 - 1); // Unchanged.
  180. }
  181. TEST_F(LearningTaskControllerImplTest, FeatureProviderIsUsed) {
  182. // If a FeatureProvider factory is provided, make sure that it's used to
  183. // adjust new examples.
  184. task_.feature_descriptions.push_back({"AddedByFeatureProvider"});
  185. SequenceBoundFeatureProvider feature_provider =
  186. base::SequenceBound<FakeFeatureProvider>(
  187. base::SequencedTaskRunnerHandle::Get());
  188. CreateController(std::move(feature_provider));
  189. LabelledExample example;
  190. example.features.push_back(FeatureValue(123));
  191. example.weight = 321u;
  192. AddExample(example);
  193. task_environment_.RunUntilIdle();
  194. EXPECT_EQ(trainer_raw_->training_data()[0].features[0], FeatureValue(124));
  195. EXPECT_EQ(trainer_raw_->training_data()[0].weight, example.weight);
  196. }
  197. TEST_F(LearningTaskControllerImplTest, FeatureSubsetsWork) {
  198. const char* feature_names[] = {
  199. "feature0", "feature1", "feature2", "feature3", "feature4", "feature5",
  200. "feature6", "feature7", "feature8", "feature9", "feature10", "feature11",
  201. };
  202. const int num_features = sizeof(feature_names) / sizeof(feature_names[0]);
  203. for (int i = 0; i < num_features; i++)
  204. task_.feature_descriptions.push_back({feature_names[i]});
  205. const size_t subset_size = 4;
  206. task_.feature_subset_size = subset_size;
  207. CreateController();
  208. // Verify that the reporter is given a subset of the features.
  209. auto subset = *reporter_raw_->feature_indices();
  210. EXPECT_EQ(subset.size(), subset_size);
  211. // Train a model. Each feature will have a unique value.
  212. LabelledExample example;
  213. for (int i = 0; i < num_features; i++)
  214. example.features.push_back(FeatureValue(i));
  215. AddExample(example);
  216. // Verify that all feature names in |subset| are present in the task.
  217. FeatureVector expected_features;
  218. expected_features.resize(subset_size);
  219. EXPECT_EQ(trainer_raw_->task().feature_descriptions.size(), subset_size);
  220. for (auto& iter : subset) {
  221. bool found = false;
  222. for (size_t i = 0; i < subset_size; i++) {
  223. if (trainer_raw_->task().feature_descriptions[i].name ==
  224. feature_names[iter]) {
  225. // Also build a vector with the features in the expected order.
  226. expected_features[i] = example.features[iter];
  227. found = true;
  228. break;
  229. }
  230. }
  231. EXPECT_TRUE(found);
  232. }
  233. // Verify that the training data has the adjusted features.
  234. EXPECT_EQ(trainer_raw_->training_data().size(), 1u);
  235. EXPECT_EQ(trainer_raw_->training_data()[0].features, expected_features);
  236. }
  237. TEST_F(LearningTaskControllerImplTest, PredictDistribution) {
  238. CreateController();
  239. // Predictions should be absl::nullopt until we have a model.
  240. LabelledExample example;
  241. VerifyPrediction(example.features, absl::nullopt);
  242. AddExample(example);
  243. TargetHistogram expected_histogram;
  244. expected_histogram += predicted_target_;
  245. VerifyPrediction(example.features, expected_histogram);
  246. }
  247. TEST_F(LearningTaskControllerImplTest,
  248. SourceIdIsProvidedToDistributionReporter) {
  249. CreateController();
  250. LabelledExample example;
  251. ukm::SourceId source_id(123);
  252. // Add two examples, so that the second causes a prediction to be reported.
  253. AddExample(example, source_id);
  254. AddExample(example, source_id);
  255. EXPECT_EQ(reporter_raw_->most_recent_source_id_, source_id);
  256. }
  257. } // namespace learning
  258. } // namespace media