jcf94 commented on a change in pull request #6310: URL: https://github.com/apache/incubator-tvm/pull/6310#discussion_r476106376
########## File path: tests/python/unittest/test_auto_scheduler_evolutionary_search.py ########## @@ -0,0 +1,35 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" Test evolutionary search. """ + +import tvm +from tvm import te, auto_scheduler + +from test_auto_scheduler_common import conv2d_nchw_bn_relu_auto_scheduler_test + +def test_evo_search(): + workload_key = auto_scheduler.make_workload_key(conv2d_nchw_bn_relu_auto_scheduler_test, + (1, 56, 56, 512, 512, 3, 1, 1)) + dag = auto_scheduler.ComputeDAG(workload_key) + task = auto_scheduler.SearchTask(dag, workload_key, tvm.target.create('llvm')) + policy = auto_scheduler.SketchPolicy(task, verbose=0) + states = policy.sample_initial_population(50) + policy.evolutionary_search(states, 10) + + +if __name__ == "__main__": + test_evo_search() Review comment: Add new lines here. ########## File path: src/auto_scheduler/search_policy/sketch_policy.cc ########## @@ -363,8 +376,150 @@ Array<State> SketchPolicyNode::EvolutionarySearch(const Array<State>& init_popul Array<State> best_states; auto tic_begin = std::chrono::high_resolution_clock::now(); - // TODO(comaniac, merrymercy, jcf94): Since we haven't finished porting the cost model part - // yet, currently delete the implementation of EvolutionarySearch. To be added later. + size_t population = init_population.size(); + int num_iters = + static_cast<int>(GetIntParam(params, SketchParamKey::EvolutionarySearch::num_iters)); + double mutation_prob = static_cast<double>( + GetDoubleParam(params, SketchParamKey::EvolutionarySearch::mutation_prob)); + + // Two ping pong buffers to avoid copy. + Array<State> states_buf1{init_population}, states_buf2; + states_buf1.reserve(population); + states_buf2.reserve(population); + Array<State>* pnow = &states_buf1; + Array<State>* pnext = &states_buf2; + + // The set of explored states to avoid redendants. + std::unordered_set<std::string> explored_set; + + // The heap to maintain the so far best states. + using StateHeapItem = std::pair<State, float>; + auto cmp = [](const StateHeapItem& left, const StateHeapItem& right) { + return left.second > right.second; + }; + using StateHeap = std::priority_queue<StateHeapItem, std::vector<StateHeapItem>, decltype(cmp)>; + StateHeap heap(cmp); + auto update_heap = [&heap, &explored_set](const Array<State>& states, + const std::vector<float>& scores, const int out_size) { + float max_score = 0.0; + for (size_t i = 0; i < states.size(); ++i) { + const State& state = states[i]; + std::string state_str = state.ToStr(); + + // Skip redundant states. + if (explored_set.count(state_str) > 0) { + continue; + } + explored_set.insert(state_str); + + if (static_cast<int>(heap.size()) < out_size) { + // Directly push item if the heap is not full yet. + heap.push({state, scores[i]}); + } else if (scores[i] > heap.top().second) { + // Replace the worst state in the heap with the new state. + heap.pop(); + heap.push({state, scores[i]}); + } + max_score = (scores[i] > max_score) ? scores[i] : max_score; + } + return max_score; + }; + + // Cost model predicted scores. + std::vector<float> scores; + scores.reserve(population); + + // The function to generate prefix sum probabilities based on the given scores. + auto assign_prob = [](const std::vector<float>& scores, std::vector<double>* prefix_sum_probs) { + // Compute selection probabilities. + double sum = 0.0; + prefix_sum_probs->resize(scores.size()); + for (size_t i = 0; i < scores.size(); ++i) { + sum += std::max(scores[i], 0.0f); + (*prefix_sum_probs)[i] = sum; + } + for (size_t i = 0; i < scores.size(); ++i) { + (*prefix_sum_probs)[i] /= sum; + } + }; + + // State selection probabilities. + std::uniform_real_distribution<> uniform_dist(0.0, 1.0); + std::vector<double> state_select_probs; + state_select_probs.reserve(population); + + // Mutation rule selection probabilities. + std::vector<double> rule_select_probs; + rule_select_probs.reserve(mutation_rules.size()); + std::vector<float> rule_levels; + for (const auto& rule : mutation_rules) { + rule_levels.push_back(rule->GetLevel(search_task)); + } + assign_prob(rule_levels, &rule_select_probs); + + // Evaluate the init populations. + *pnow = search_task->compute_dag.InferBound(*pnow); + PruneInvalidState(search_task, pnow); + CHECK_GT(pnow->size(), 0) << "All initial populations are invalid"; + schedule_cost_model->Predict(search_task, *pnow, &scores); + + // Maintain the best states in the heap. + float max_score = update_heap(*pnow, scores, out_size); + + // Genetic algorithm. + for (auto iter_idx = 1; iter_idx <= num_iters; ++iter_idx) { + // Assign the selection probability to each state based on the cost model scores. + assign_prob(scores, &state_select_probs); + + // TODO(@comaniac): Perform cross over. + + // Perform mutations. + size_t fail_ct = 0; + while (pnext->size() < population && fail_ct < population * 2) { + // Select a state to be mutated. + State tmp_s = (*pnow)[RandomChoose(state_select_probs, &rand_gen)]; + if (uniform_dist(rand_gen) < mutation_prob) { + // Select a rule and mutate the state. + const auto rule = mutation_rules[RandomChoose(rule_select_probs, &rand_gen)]; Review comment: Should we have `const auto` -> `const auto&` ? ########## File path: tests/python/unittest/test_auto_scheduler_evolutionary_search.py ########## @@ -0,0 +1,35 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" Test evolutionary search. """ + +import tvm +from tvm import te, auto_scheduler + +from test_auto_scheduler_common import conv2d_nchw_bn_relu_auto_scheduler_test + +def test_evo_search(): Review comment: I suggest we add an extra test that uses SketchPolicy with XGBModel in `test_auto_scheduler_search_policy.py`, this should trigger the complete work flow including `GenerateSketches`, `SampleInitPopulation` and `EvolutionarySearch`. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
