# 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. """Tuner that uses xgboost as cost model""" from .model_based_tuner import ModelBasedTuner, ModelOptimizer from .xgboost_cost_model import XGBoostCostModel from .sa_model_optimizer import SimulatedAnnealingOptimizer class XGBTuner(ModelBasedTuner): """Tuner that uses xgboost as cost model Parameters ---------- task: Task The tuning task plan_size: int The size of a plan. After `plan_size` trials, the tuner will refit a new cost model and do planing for the next `plan_size` trials. feature_type: str, optional If is 'itervar', use features extracted from IterVar (loop variable). If is 'knob', use flatten ConfigEntity directly. If is 'curve', use sampled curve feature (relation feature). Note on choosing feature type: For single task tuning, 'itervar' and 'knob' are good. 'itervar' is more accurate but 'knob' is much faster. There are some constraints on 'itervar', if you meet problems with feature extraction when using 'itervar', you can swith to 'knob'. For cross-shape tuning (e.g. many convolutions with different shapes), 'itervar' and 'curve' has better transferability, 'knob' is faster. For cross-device or cross-operator tuning, you can use 'curve' only. loss_type: str If is 'reg', use regression loss to train cost model. The cost model predicts the normalized flops. If is 'rank', use pairwise rank loss to train cost model. The cost model predicts relative rank score. num_threads: int, optional The number of threads. optimizer: str or ModelOptimizer, optional If is 'sa', use a default simulated annealing optimizer. Otherwise it should be a ModelOptimizer object. diversity_filter_ratio: int or float, optional If is not None, the tuner will first select top-(plan_size * diversity_filter_ratio) candidates according to the cost model and then pick batch_size of them according to the diversity metric. log_interval: int, optional The verbose level. If is 0, output nothing. Otherwise, output debug information every `verbose` iterations. """ def __init__(self, task, plan_size=64, feature_type='itervar', loss_type='rank', num_threads=None, optimizer='sa', diversity_filter_ratio=None, log_interval=50): cost_model = XGBoostCostModel(task, feature_type=feature_type, loss_type=loss_type, num_threads=num_threads, log_interval=log_interval // 2) if optimizer == 'sa': optimizer = SimulatedAnnealingOptimizer(task, log_interval=log_interval) else: assert isinstance(optimizer, ModelOptimizer), "Optimizer must be " \ "a supported name string" \ "or a ModelOptimizer object." super(XGBTuner, self).__init__(task, cost_model, optimizer, plan_size, diversity_filter_ratio) def tune(self, *args, **kwargs): # pylint: disable=arguments-differ super(XGBTuner, self).tune(*args, **kwargs) # manually close pool to avoid multiprocessing issues self.cost_model._close_pool()