diff --git a/configs/nas/mmcls/spos/spos_shufflenet_search_predictor_8xb128_in1k.py b/configs/nas/mmcls/spos/spos_shufflenet_search_predictor_8xb128_in1k.py new file mode 100644 index 000000000..fb08d3494 --- /dev/null +++ b/configs/nas/mmcls/spos/spos_shufflenet_search_predictor_8xb128_in1k.py @@ -0,0 +1,21 @@ +_base_ = ['./spos_shufflenet_supernet_8xb128_in1k.py'] + +model = dict(norm_training=True) + +train_cfg = dict( + _delete_=True, + type='mmrazor.EvolutionSearchLoop', + dataloader=_base_.val_dataloader, + evaluator=_base_.val_evaluator, + max_epochs=20, + num_candidates=50, + top_k=10, + num_mutation=25, + num_crossover=25, + mutate_prob=0.1, + constraints_range=dict(flops=(0., 360.)), + predictor_cfg=dict( + type='mmrazor.MetricPredictor', + train_samples=20, + handler_cfg=dict(type='mmrazor.GaussProcessHandler')), +) diff --git a/mmrazor/engine/hooks/estimate_resources_hook.py b/mmrazor/engine/hooks/estimate_resources_hook.py index dc27f2906..e8c4d8446 100644 --- a/mmrazor/engine/hooks/estimate_resources_hook.py +++ b/mmrazor/engine/hooks/estimate_resources_hook.py @@ -7,7 +7,7 @@ from mmengine.registry import HOOKS from mmengine.structures import BaseDataElement -from mmrazor.models.task_modules import ResourceEstimator +from mmrazor.registry import TASK_UTILS DATA_BATCH = Optional[Sequence[dict]] @@ -23,7 +23,7 @@ class EstimateResourcesHook(Hook): by_epoch (bool): Saving checkpoints by epoch or by iteration. Default to True. estimator_cfg (Dict[str, Any]): Used for building a resource estimator. - Default to dict(). + Default to None. Example: >>> add the `EstimatorResourcesHook` in custom_hooks as follows: @@ -41,11 +41,14 @@ class EstimateResourcesHook(Hook): def __init__(self, interval: int = -1, by_epoch: bool = True, - estimator_cfg: Dict[str, Any] = dict(), + estimator_cfg: Dict[str, Any] = None, **kwargs) -> None: self.interval = interval self.by_epoch = by_epoch - self.estimator = ResourceEstimator(**estimator_cfg) + estimator_cfg = dict() if estimator_cfg is None else estimator_cfg + if 'type' not in estimator_cfg: + estimator_cfg['type'] = 'mmrazor.ResourceEstimator' + self.estimator = TASK_UTILS.build(estimator_cfg) def after_val_epoch(self, runner, diff --git a/mmrazor/engine/runner/evolution_search_loop.py b/mmrazor/engine/runner/evolution_search_loop.py index d85c0ed30..fc907f3aa 100644 --- a/mmrazor/engine/runner/evolution_search_loop.py +++ b/mmrazor/engine/runner/evolution_search_loop.py @@ -1,11 +1,11 @@ # Copyright (c) OpenMMLab. All rights reserved. -import copy import os import os.path as osp import random import warnings from typing import Any, Dict, List, Optional, Tuple, Union +import numpy as np import torch from mmengine import fileio from mmengine.dist import broadcast_object_list @@ -14,7 +14,6 @@ from mmengine.utils import is_list_of from torch.utils.data import DataLoader -from mmrazor.models.task_modules import ResourceEstimator from mmrazor.registry import LOOPS, TASK_UTILS from mmrazor.structures import Candidates, export_fix_subnet from mmrazor.utils import SupportRandomSubnet @@ -45,8 +44,10 @@ class EvolutionSearchLoop(EpochBasedTrainLoop): crossover_prob (float): The probability of crossover. Defaults to 0.5. constraints_range (Dict[str, Any]): Constraints to be used for screening candidates. Defaults to dict(flops=(0, 330)). - resource_estimator_cfg (dict, Optional): Used for building a - resource estimator. Defaults to None. + estimator_cfg (dict, Optional): Used for building a resource estimator. + Defaults to None. + predictor_cfg (dict, Optional): Used for building a metric predictor. + Defaults to None. score_key (str): Specify one metric in evaluation results to score candidates. Defaults to 'accuracy_top-1'. init_candidates (str, optional): The candidates file path, which is @@ -68,7 +69,8 @@ def __init__(self, mutate_prob: float = 0.1, crossover_prob: float = 0.5, constraints_range: Dict[str, Any] = dict(flops=(0., 330.)), - resource_estimator_cfg: Optional[Dict] = None, + estimator_cfg: Optional[Dict] = None, + predictor_cfg: Optional[Dict] = None, score_key: str = 'accuracy/top1', init_candidates: Optional[str] = None) -> None: super().__init__(runner, dataloader, max_epochs) @@ -109,56 +111,28 @@ def __init__(self, else: self.model = runner.model - # Build resource estimator. - resource_estimator_cfg = dict( - ) if resource_estimator_cfg is None else resource_estimator_cfg - self.estimator = self.build_resource_estimator(resource_estimator_cfg) - - def build_resource_estimator( - self, resource_estimator: Union[ResourceEstimator, - Dict]) -> ResourceEstimator: - """Build resource estimator for search loop. - - Examples of ``resource_estimator``: - - # `ResourceEstimator` will be used - resource_estimator = dict() - - # custom resource_estimator - resource_estimator = dict(type='mmrazor.ResourceEstimator') - - Args: - resource_estimator (ResourceEstimator or dict): A - resource_estimator or a dict to build resource estimator. - If ``resource_estimator`` is a resource estimator object, - just returns itself. - - Returns: - :obj:`ResourceEstimator`: Resource estimator object build from - ``resource_estimator``. - """ - if isinstance(resource_estimator, ResourceEstimator): - return resource_estimator - elif not isinstance(resource_estimator, dict): - raise TypeError( - 'resource estimator should be a ResourceEstimator object or' - f'dict, but got {resource_estimator}') - - resource_estimator_cfg = copy.deepcopy( - resource_estimator) # type: ignore - - if 'type' in resource_estimator_cfg: - estimator = TASK_UTILS.build(resource_estimator_cfg) - else: - estimator = ResourceEstimator( - **resource_estimator_cfg) # type: ignore - - return estimator # type: ignore + # initialize estimator + estimator_cfg = dict() if estimator_cfg is None else estimator_cfg + if 'type' not in estimator_cfg: + estimator_cfg['type'] = 'mmrazor.ResourceEstimator' + self.estimator = TASK_UTILS.build(estimator_cfg) + + # initialize predictor + self.use_predictor = False + self.predictor_cfg = predictor_cfg + if self.predictor_cfg is not None: + self.predictor_cfg['score_key'] = self.score_key + self.predictor_cfg['search_groups'] = \ + self.model.mutator.search_groups + self.predictor = TASK_UTILS.build(self.predictor_cfg) def run(self) -> None: """Launch searching.""" self.runner.call_hook('before_train') + if self.predictor_cfg is not None: + self._init_predictor() + if self.resume_from: self._resume() @@ -174,7 +148,7 @@ def run_epoch(self) -> None: """Iterate one epoch. Steps: - 1. Sample some new candidates from the supernet.Then Append them + 1. Sample some new candidates from the supernet. Then Append them to the candidates, Thus make its number equal to the specified number. 2. Validate these candidates(step 1) and update their scores. @@ -240,8 +214,8 @@ def update_candidates_scores(self) -> None: top-k candicates.""" for i, candidate in enumerate(self.candidates.subnets): self.model.set_subnet(candidate) - metrics = self._val_candidate() - score = metrics[self.score_key] \ + metrics = self._val_candidate(use_predictor=self.use_predictor) + score = round(metrics[self.score_key], 2) \ if len(metrics) != 0 else 0. self.candidates.set_resource(i, score, 'score') self.runner.logger.info( @@ -250,7 +224,7 @@ def update_candidates_scores(self) -> None: f'Flops: {self.candidates.resources("flops")[i]} ' f'Params: {self.candidates.resources("params")[i]} ' f'Latency: {self.candidates.resources("latency")[i]} ' - f'Score: {self.candidates.scores} ') + f'Score: {self.candidates.scores[i]} ') def gen_mutation_candidates(self): """Generate specified number of mutation candicates.""" @@ -340,13 +314,23 @@ def _save_best_fix_subnet(self): f'{save_name} saved in {self.runner.work_dir}.') @torch.no_grad() - def _val_candidate(self) -> Dict: - """Run validation.""" - self.runner.model.eval() - for data_batch in self.dataloader: - outputs = self.runner.model.val_step(data_batch) - self.evaluator.process(outputs, data_batch) - metrics = self.evaluator.evaluate(len(self.dataloader.dataset)) + def _val_candidate(self, use_predictor: bool = False) -> Dict: + """Run validation. + + Args: + use_predictor (bool): Whether to use predictor to get metrics. + Defaults to False. + """ + if use_predictor: + assert self.predictor is not None + metrics = self.predictor.predict(self.model) + else: + self.runner.model.eval() + for data_batch in self.dataloader: + outputs = self.runner.model.val_step(data_batch) + self.evaluator.process( + data_samples=outputs, data_batch=data_batch) + metrics = self.evaluator.evaluate(len(self.dataloader.dataset)) return metrics def _save_searcher_ckpt(self) -> None: @@ -391,3 +375,43 @@ def _check_constraints( constraints_range=self.constraints_range) return is_pass, results + + def _init_predictor(self): + """Initialize predictor, training is required.""" + if self.predictor.handler_ckpt: + self.predictor.load_checkpoint() + self.runner.logger.info( + f'Loaded Checkpoints from {self.predictor.handler_ckpt}') + else: + self.runner.logger.info('No predictor checkpoints found. ' + 'Start pre-training the predictor.') + if isinstance(self.predictor.train_samples, str): + self.runner.logger.info('Find specified samples in ' + f'{self.predictor.train_samples}') + train_samples = fileio.load(self.predictor.train_samples) + self.candidates = train_samples['subnets'] + else: + self.runner.logger.info( + 'Without specified samples. Start random sampling.') + temp_num_candidates = self.num_candidates + self.num_candidates = self.predictor.train_samples + + assert self.use_predictor is False, ( + 'Real evaluation is required when initializing predictor.') + self.sample_candidates() + self.update_candidates_scores() + self.num_candidates = temp_num_candidates + + inputs = [] + for candidate in self.candidates.subnets: + inputs.append(self.predictor.model2vector(candidate)) + inputs = np.array(inputs) + labels = np.array(self.candidates.scores) + self.predictor.fit(inputs, labels) + if self.runner.rank == 0: + predictor_dir = self.predictor.save_checkpoint( + osp.join(self.runner.work_dir, 'predictor')) + self.runner.logger.info( + f'Predictor pre-trained, saved in {predictor_dir}.') + self.use_predictor = True + self.candidates = Candidates() diff --git a/mmrazor/engine/runner/subnet_sampler_loop.py b/mmrazor/engine/runner/subnet_sampler_loop.py index 273561568..56c4f893c 100644 --- a/mmrazor/engine/runner/subnet_sampler_loop.py +++ b/mmrazor/engine/runner/subnet_sampler_loop.py @@ -1,5 +1,4 @@ # Copyright (c) OpenMMLab. All rights reserved. -import copy import math import os import random @@ -13,7 +12,6 @@ from mmengine.utils import is_list_of from torch.utils.data import DataLoader -from mmrazor.models.task_modules import ResourceEstimator from mmrazor.registry import LOOPS, TASK_UTILS from mmrazor.structures import Candidates from mmrazor.utils import SupportRandomSubnet @@ -102,8 +100,8 @@ class GreedySamplerTrainLoop(BaseSamplerTrainLoop): candidates. Defaults to 'accuracy_top-1'. constraints_range (Dict[str, Any]): Constraints to be used for screening candidates. Defaults to dict(flops=(0, 330)). - resource_estimator_cfg (dict, Optional): Used for building a - resource estimator. Defaults to None. + estimator_cfg (dict, Optional): Used for building a resource estimator. + Defaults to None. num_candidates (int): The number of the candidates consist of samples from supernet and itself. Defaults to 1000. num_samples (int): The number of sample in each sampling subnet. @@ -138,7 +136,7 @@ def __init__(self, val_interval: int = 1000, score_key: str = 'accuracy/top1', constraints_range: Dict[str, Any] = dict(flops=(0, 330)), - resource_estimator_cfg: Optional[Dict] = None, + estimator_cfg: Optional[Dict] = None, num_candidates: int = 1000, num_samples: int = 10, top_k: int = 5, @@ -176,51 +174,11 @@ def __init__(self, self.candidates = Candidates() self.top_k_candidates = Candidates() - # Build resource estimator. - resource_estimator_cfg = dict( - ) if resource_estimator_cfg is None else resource_estimator_cfg - self.estimator = self.build_resource_estimator(resource_estimator_cfg) - - def build_resource_estimator( - self, resource_estimator: Union[ResourceEstimator, - Dict]) -> ResourceEstimator: - """Build resource estimator for search loop. - - Examples of ``resource_estimator``: - - # `ResourceEstimator` will be used - resource_estimator = dict() - - # custom resource_estimator - resource_estimator = dict(type='mmrazor.ResourceEstimator') - - Args: - resource_estimator (ResourceEstimator or dict): - A resource_estimator or a dict to build resource estimator. - If ``resource_estimator`` is a resource estimator object, - just returns itself. - - Returns: - :obj:`ResourceEstimator`: Resource estimator object build from - ``resource_estimator``. - """ - if isinstance(resource_estimator, ResourceEstimator): - return resource_estimator - elif not isinstance(resource_estimator, dict): - raise TypeError( - 'resource estimator should be a ResourceEstimator object or' - f'dict, but got {resource_estimator}') - - resource_estimator_cfg = copy.deepcopy( - resource_estimator) # type: ignore - - if 'type' in resource_estimator_cfg: - estimator = TASK_UTILS.build(resource_estimator_cfg) - else: - estimator = ResourceEstimator( - **resource_estimator_cfg) # type: ignore - - return estimator # type: ignore + # initialize estimator + estimator_cfg = dict() if estimator_cfg is None else estimator_cfg + if 'type' not in estimator_cfg: + estimator_cfg['type'] = 'mmrazor.ResourceEstimator' + self.estimator = TASK_UTILS.build(estimator_cfg) def run(self) -> None: """Launch training.""" diff --git a/mmrazor/models/algorithms/nas/dsnas.py b/mmrazor/models/algorithms/nas/dsnas.py index 5434ce0ac..a763c75ea 100644 --- a/mmrazor/models/algorithms/nas/dsnas.py +++ b/mmrazor/models/algorithms/nas/dsnas.py @@ -68,8 +68,10 @@ def __init__(self, **kwargs): super().__init__(architecture, data_preprocessor, **kwargs) - if estimator_cfg is None: - estimator_cfg = dict(type='mmrazor.ResourceEstimator') + # initialize estimator + estimator_cfg = dict() if estimator_cfg is None else estimator_cfg + if 'type' not in estimator_cfg: + estimator_cfg['type'] = 'mmrazor.ResourceEstimator' self.estimator = TASK_UTILS.build(estimator_cfg) if fix_subnet: # Avoid circular import diff --git a/mmrazor/models/task_modules/__init__.py b/mmrazor/models/task_modules/__init__.py index 56cb681e7..b86bebbb9 100644 --- a/mmrazor/models/task_modules/__init__.py +++ b/mmrazor/models/task_modules/__init__.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. from .delivery import * # noqa: F401,F403 from .estimators import ResourceEstimator +from .predictor import * # noqa: F401,F403 from .recorder import * # noqa: F401,F403 from .tracer import * # noqa: F401,F403 diff --git a/mmrazor/models/task_modules/predictor/__init__.py b/mmrazor/models/task_modules/predictor/__init__.py new file mode 100644 index 000000000..6caa0dbf8 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .metric_predictor import MetricPredictor + +__all__ = ['MetricPredictor'] diff --git a/mmrazor/models/task_modules/predictor/handler/__init__.py b/mmrazor/models/task_modules/predictor/handler/__init__.py new file mode 100644 index 000000000..96337921d --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .carts_handler import CartsHandler +from .gp_handler import GaussProcessHandler +from .mlp_handler import MLPHandler +from .rbf_handler import RBFHandler + +__all__ = ['CartsHandler', 'GaussProcessHandler', 'MLPHandler', 'RBFHandler'] diff --git a/mmrazor/models/task_modules/predictor/handler/base_handler.py b/mmrazor/models/task_modules/predictor/handler/base_handler.py new file mode 100644 index 000000000..40246ac3d --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/base_handler.py @@ -0,0 +1,32 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from joblib import dump, load + + +class BaseHandler: + """Base class for a handler. + + Note: + The handler works through a specific machine leanring algorithm, + and is designed for predicting the evaluation metric of a model. + """ + + def __init__(self) -> None: + pass + + def fit(self, train_data, train_label): + """Training the model of handler.""" + pass + + def predict(self, test_data): + """Predicting the metric using the model of handler.""" + pass + + def load(self, path): + """Load pretrained weights for the handler.""" + self.model = load(path) + + def save(self, path): + """Save the handler and return saved path for diff suffix.""" + path += f'_{self.__class__.__name__}.joblib'.lower() + dump(self.model, path) + return path diff --git a/mmrazor/models/task_modules/predictor/handler/carts_handler.py b/mmrazor/models/task_modules/predictor/handler/carts_handler.py new file mode 100644 index 000000000..318503b6f --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/carts_handler.py @@ -0,0 +1,92 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import List + +import numpy as np +from sklearn.tree import DecisionTreeRegressor + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + + +@TASK_UTILS.register_module() +class CartsHandler(BaseHandler): + """Classification and Regression Tree. + + Args: + num_trees (int): number of regression trees. + """ + + def __init__(self, num_trees=1000): + self.num_trees = num_trees + + def fit(self, train_data: np.array, train_label: np.array) -> None: + """Define the model of handler. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + """ + self.model = self._make_decision_trees(train_data, train_label, + self.num_trees) + + def predict(self, test_data: np.array) -> np.array: + """Predict the evaluation metric of the model. + + Args: + test_data (numpy.array): input data for testing. + + Returns: + numpy.array: predicted metric. + """ + trees, features = self.model[0], self.model[1] + test_num, num_trees = len(test_data), len(trees) + + predict_labels = np.zeros((test_num, 1)) + for i in range(test_num): + this_test_data = test_data[i, :] + predict_this_list = np.zeros(num_trees) + + for j, (tree, feature) in enumerate(zip(trees, features)): + predict_this_list[j] = tree.predict([this_test_data[feature] + ])[0] + + predict_this_list = np.sort(predict_this_list) + predict_this_list = predict_this_list[::-1] + this_predict = np.mean(predict_this_list) + predict_labels[i, 0] = this_predict + + return predict_labels + + @staticmethod + def _make_decision_trees(train_data: np.array, train_label: np.array, + num_trees: int) -> List[list]: + """Construct the decision trees. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + num_trees (int): num of decision trees. + + Returns: + List[list]: List of built models. + """ + feature_record = [] + tree_record = [] + + for _ in range(num_trees): + sample_idx = np.arange(train_data.shape[0]) + np.random.shuffle(sample_idx) + train_data = train_data[sample_idx, :] + train_label = train_label[sample_idx] + + feature_idx = np.arange(train_data.shape[1]) + np.random.shuffle(feature_idx) + n_feature = np.random.randint(1, train_data.shape[1] + 1) + selected_feature_ids = feature_idx[0:n_feature] + feature_record.append(selected_feature_ids) + + dt = DecisionTreeRegressor() + dt.fit(train_data[:, selected_feature_ids], train_label) + tree_record.append(dt) + + return [tree_record, feature_record] diff --git a/mmrazor/models/task_modules/predictor/handler/gp_handler.py b/mmrazor/models/task_modules/predictor/handler/gp_handler.py new file mode 100644 index 000000000..4935abfc1 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/gp_handler.py @@ -0,0 +1,115 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import numpy as np +from pydacefit.corr import (corr_cubic, corr_exp, corr_expg, corr_gauss, + corr_spherical, corr_spline) +from pydacefit.dace import DACE, regr_linear, regr_quadratic +from pydacefit.fit import fit as pydace_fit +from pydacefit.regr import regr_constant + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + +REGR = { + 'linear': regr_linear, + 'constant': regr_constant, + 'quadratic': regr_quadratic +} + +CORR = { + 'gauss': corr_gauss, + 'cubic': corr_cubic, + 'exp': corr_exp, + 'expg': corr_expg, + 'spline': corr_spline, + 'spherical': corr_spherical +} + + +class DACE_with_smooth(DACE): + """GP model.""" + + def __init__(self, + regr, + corr, + theta: float = 1.0, + thetaL: float = 0.0, + thetaU: float = 100.0): + super(DACE_with_smooth, self).__init__(regr, corr, theta, thetaL, + thetaU) + + def fit(self, X, Y): + """Build the model.""" + if len(Y.shape) == 1: + Y = Y[:, None] + + if X.shape[0] != Y.shape[0]: + raise Exception('X and Y must have the same number of rows.') + + mX, sX = np.mean(X, axis=0), np.std(X, axis=0, ddof=1) + 1e-6 + mY, sY = np.mean(Y, axis=0), np.std(Y, axis=0, ddof=1) + 1e-6 + + nX = (X - mX) / sX + nY = (Y - mY) / sY + + if self.tl is not None and self.tu is not None: + self.model = {'nX': nX, 'nY': nY} + self.boxmin() + self.model = self.itpar['best'] + else: + self.model = pydace_fit(nX, nY, self.regr, self.kernel, self.theta) + + self.model = { + **self.model, 'mX': mX, + 'sX': sX, + 'mY': mY, + 'sY': sY, + 'nX': nX, + 'nY': nY + } + self.model['sigma2'] = np.square(sY) @ self.model['_sigma2'] + + +@TASK_UTILS.register_module() +class GaussProcessHandler(BaseHandler): + """GaussProcess handler of the metric predictor. It uses Gaussian Process + (Kriging) to predict the metric of a trained model. + + Args: + regr (str): regression kernel for GP model. Defaults to 'linear'. + corr (str): correlation kernel for GP model. Defaults to 'gauss'. + """ + + def __init__(self, regr: str = 'linear', corr: str = 'gauss'): + assert regr in REGR, \ + ValueError(f'`regr` should be in `REGR`. Got `{regr}`.') + assert corr in CORR, \ + ValueError(f'`corr` should be in `CORR`. Got `{corr}`.') + self.regr = REGR[regr] + self.corr = CORR[corr] + + self.model = DACE_with_smooth( + regr=self.regr, + corr=self.corr, + theta=1.0, + thetaL=0.00001, + thetaU=100) + + def fit(self, train_data: np.array, train_label: np.array) -> None: + """Training the model of handler. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + """ + self.model.fit(train_data, train_label) + + def predict(self, test_data: np.array) -> np.array: + """Predict the evaluation metric of the model. + + Args: + test_data (numpy.array): input data for testing. + + Returns: + numpy.array: predicted metric. + """ + return self.model.predict(test_data) diff --git a/mmrazor/models/task_modules/predictor/handler/mlp_handler.py b/mmrazor/models/task_modules/predictor/handler/mlp_handler.py new file mode 100644 index 000000000..65e1ccb80 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/mlp_handler.py @@ -0,0 +1,192 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import copy +from typing import Dict + +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +from mmcv.cnn.bricks import build_activation_layer +from mmdet.models.losses import SmoothL1Loss +from mmengine.model import BaseModule +from mmengine.optim.scheduler import CosineAnnealingLR + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + + +class MLP(BaseModule): + """MLP implemented with nn.Linear. + + Input: Tensor with shape [B, C, H, W]. + Output: Tensor with shape [B, C, H, W]. + + Args: + in_features (int): Dimension of input features. + hidden_features (int): Dimension of hidden features. + out_features (int): Dimension of output features. + act_cfg (dict): The config dict for activation between pointwise + convolution. Defaults to ``dict(type='ReLU')``. + drop (float): Dropout rate. Defaults to 0.0. + """ + + def __init__(self, + in_features: int = 78, + hidden_features: int = 300, + out_features: int = 1, + num_hidden_layers: int = 2, + act_cfg: Dict = dict(type='ReLU'), + drop: float = 0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = build_activation_layer(act_cfg) + + hidden_layers = [] + for _ in range(num_hidden_layers): + hidden_layers.append(nn.Linear(hidden_features, hidden_features)) + hidden_layers.append(build_activation_layer(act_cfg)) + self.hidden_layers = nn.Sequential(*hidden_layers) + + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + self.init_weights() + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.hidden_layers(x) + x = self.drop(x) + x = self.fc2(x) + return x + + +@TASK_UTILS.register_module() +class MLPHandler(BaseHandler): + """MLP handler of the metric predictor. It uses MLP network to predict the + metric of a trained model. + + Args: + epochs (int, optional): num of epochs for MLP network training. + Defaults to 100. + data_split_ratio (float, optional): split ratio of train/valid of + input data. Defaults to 0.8. + model_cfg (dict, optional): configs for MLP network. Defaults to None. + device (str, optional): device for MLP Handler. Defaults to 'cuda'. + """ + + def __init__(self, + epochs: int = 100, + data_split_ratio: float = 0.8, + model_cfg: Dict = None, + device: str = 'cpu'): + self.epochs = epochs + self.data_split_ratio = data_split_ratio + + self.model_cfg = model_cfg if model_cfg is not None else dict() + self.model = MLP(**self.model_cfg) + + self.device = device + + def fit(self, train_data: np.array, train_label: np.array) -> None: + """Training the model of handler. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + """ + if train_data.shape[1] != self.model.fc1.in_features: + self.model.fc1 = nn.Linear(train_data.shape[1], + self.model.fc1.out_features) + self.model = self.train_mlp(train_data, train_label) + + def predict(self, test_data: np.array) -> np.array: + """Predict the evaluation metric of the model. + + Args: + test_data (numpy.array): input data for testing. + + Returns: + numpy.array: predicted metric. + """ + if test_data.ndim < 2: + data = torch.zeros(1, test_data.shape[0]) + data[0, :] = torch.from_numpy(test_data).float() + else: + data = torch.from_numpy(test_data).float() + + self.model = self.model.to(device=self.device) + self.model.eval() + with torch.no_grad(): + data = data.to(device=self.device) + pred = self.model(data) + + return pred.cpu().detach().numpy() + + def load(self, path: str) -> None: + """Load predictor's pretrained weights.""" + self.model.load_state_dict( + torch.load(path, map_location='cpu')['state_dict']) + + def save(self, path: str) -> str: + """Save predictor and return saved path for diff suffix.""" + path = path + '_mlp.pth' + torch.save({'state_dict': self.model.state_dict(), 'meta': {}}, path) + return path + + def train_mlp(self, train_data: np.array, + train_label: np.array) -> nn.Module: + """Train MLP network. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + + Returns: + nn.Module: the well-trained MLP network. + """ + num_samples = train_data.shape[0] + target = torch.zeros(num_samples, 1) + perm = torch.randperm(target.size(0)) + train_index = perm[:int(num_samples * self.data_split_ratio)] + valid_index = perm[int(num_samples * self.data_split_ratio):] + + inputs = torch.from_numpy(train_data).float() + target[:, 0] = torch.from_numpy(train_label).float() + + self.model = self.model.to(device=self.device) + self.optimizer = optim.Adam(self.model.parameters(), lr=8e-4) + self.criterion = SmoothL1Loss() + + self.scheduler = CosineAnnealingLR( + self.optimizer, T_max=self.epochs, eta_min=0, by_epoch=True) + + best_loss = 1e33 + for _ in range(self.epochs): + train_inputs = inputs[train_index].to(self.device) + train_labels = target[train_index].to(self.device) + + self.model.train() + self.optimizer.zero_grad() + pred = self.model(train_inputs) + loss = self.criterion(pred, train_labels) + loss.backward() + self.optimizer.step() + + self.model.eval() + with torch.no_grad(): + valid_inputs = inputs[valid_index].to(self.device) + valid_labels = target[valid_index].to(self.device) + + pred = self.model(valid_inputs) + valid_loss = self.criterion(pred, valid_labels).item() + + self.scheduler.step() + + if valid_loss < best_loss: + best_loss = valid_loss + best_net = copy.deepcopy(self.model) + + return best_net.to(device='cpu') diff --git a/mmrazor/models/task_modules/predictor/handler/rbf_handler.py b/mmrazor/models/task_modules/predictor/handler/rbf_handler.py new file mode 100644 index 000000000..a1b620fb8 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/rbf_handler.py @@ -0,0 +1,60 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import numpy as np +from pySOT.surrogate import (ConstantTail, CubicKernel, Kernel, LinearTail, + RBFInterpolant, Tail, TPSKernel) + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + + +@TASK_UTILS.register_module() +class RBFHandler(BaseHandler): + """RBF handler of the metric predictor. It uses `Radial Basis Function` to + predict the metric of a trained model. + + Args: + kernel (str): RBF kernel object. Defaults to 'tps'. + tail (str): RBF polynomial tail object. Defaults to 'linear'. + """ + kernel_mapping = {'cubic': CubicKernel, 'tps': TPSKernel} + tail_mapping = {'linear': LinearTail, 'constant': ConstantTail} + + def __init__(self, kernel: str = 'tps', tail: str = 'linear'): + assert kernel in self.kernel_mapping.keys(), ( + f'Got unknown RBF kernel `{kernel}`.') + self.kernel: Kernel = self.kernel_mapping[kernel] + + assert tail in self.tail_mapping.keys(), ( + f'Got unknown RBF tail `{tail}`.') + self.tail: Tail = self.tail_mapping[tail] + + def fit(self, train_data: np.array, train_label: np.array) -> None: + """Training the model of handler. + + Args: + train_data (numpy.array): input data for training. + train_label (numpy.array): input label for training. + """ + if train_data.shape[0] <= train_data.shape[1]: + raise ValueError('In RBF, dim 0 of data (got ' + f'{train_data.shape[0]}) should be larger than ' + f'dim 1 of data (got {train_data.shape[1]}).') + + self.model = RBFInterpolant( + dim=train_data.shape[1], + kernel=self.kernel(), + tail=self.tail(train_data.shape[1])) + + for i in range(len(train_data)): + self.model.add_points(train_data[i, :], train_label[i]) + + def predict(self, test_data: np.array) -> np.array: + """Predict the evaluation metric of the model. + + Args: + test_data (numpy.array): input data for testing. + + Returns: + numpy.array: predicted metric. + """ + return self.model.predict(test_data) diff --git a/mmrazor/models/task_modules/predictor/metric_predictor.py b/mmrazor/models/task_modules/predictor/metric_predictor.py new file mode 100644 index 000000000..8660badd9 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/metric_predictor.py @@ -0,0 +1,198 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Dict, List, Union + +import numpy as np +import scipy.stats as stats + +from mmrazor.registry import TASK_UTILS +from mmrazor.structures import export_fix_subnet +from mmrazor.utils.typing import DumpChosen +from .handler import RBFHandler + + +@TASK_UTILS.register_module() +class MetricPredictor: + """A predictor for predicting evaluation metrics in different tasks. + + Args: + handler_cfg (dict): Config to build a predict handler. + search_groups (dict) : The search_groups of the specified supernet. + train_samples (int): Num of training samples for the handler. + Defaults to 2. + handler_ckpt (str, optional): Path to handler's checkpoint. If given, + predictor will load weights directly instead of handler training. + encoding_type (str, optional): Type of how to encode the search space + to integer bit-string. Defaults to `onehot`. + score_key (str): Specify one metric in evaluation results to score + models. Defaults to 'accuracy_top-1'. + """ + + def __init__(self, + handler_cfg: Dict, + search_groups: Dict, + train_samples: int = 2, + handler_ckpt: str = None, + encoding_type: str = 'onehot', + score_key: str = 'accuracy_top-1', + **kwargs): + self.handler_cfg = handler_cfg + self.handler = TASK_UTILS.build(handler_cfg) + + assert encoding_type in [ + 'normal', 'onehot' + ], ('encoding_type must be `normal` or `onehot`.' + f'Got `{encoding_type}`.') + if isinstance(self.handler, RBFHandler): + encoding_type = 'normal' + self.encoding_type = encoding_type + + self.search_groups = search_groups + self.train_samples = train_samples + self.handler_ckpt = handler_ckpt + + self.score_key_list = [score_key] + ['anticipate'] + self.initialize = False + + def predict(self, model) -> Dict[str, float]: + """Predict the evaluation metric of input model using the handler. + + Args: + model: input model. + + Returns: + Dict[str, float]: evaluation metric of the model. + """ + metric: Dict[str, float] = {} + assert self.initialize is True, ( + 'Before predicting, evaluator is required to be executed first, ' + 'cause the model of handler in predictor needs to be initialized.') + + if self.initialize: + model = export_fix_subnet(model) + data = self.preprocess(np.array([self.model2vector(model)])) + score = float(np.squeeze(self.handler.predict(data))) + if metric.get(self.score_key_list[0], None): + metric.update({self.score_key_list[1]: score}) + else: + metric.update({self.score_key_list[0]: score}) + return metric + + def model2vector( + self, model: Dict[str, Union[str, DumpChosen]]) -> Dict[str, list]: + """Convert the input model to N-dims vector. + + Args: + model (Dict[str, Union[str, DumpChosen]]): input model. + + Returns: + Dict[str, list]: converted vector. + """ + index = 0 + vector_dict: Dict[str, list] = \ + dict(normal_vector=[], onehot_vector=[]) + + for key, choice in model.items(): + if isinstance(choice, DumpChosen): + assert choice.meta is not None, ( + f'`DumpChosen.meta` of current {key} should not be None ' + 'when converting the search space.') + onehot = np.zeros( + len(choice.meta['all_choices']), dtype=np.int) + _chosen_index = choice.meta['all_choices'].index(choice.chosen) + else: + assert len(self.search_groups[index]) == 1 + choices = self.search_groups[index][0].choices + onehot = np.zeros(len(choices), dtype=np.int) + _chosen_index = choices.index(choice) + onehot[_chosen_index] = 1 + + vector_dict['normal_vector'].extend([_chosen_index]) + vector_dict['onehot_vector'].extend(onehot) + index += 1 + + return vector_dict + + def vector2model(self, vector: np.array) -> Dict[str, str]: + """Convert the N-dims vector to original model. + + Args: + vector (numpy.array): input vector which represents the model. + + Returns: + Dict[str, str]: converted model. + """ + start = 0 + model = {} + for key, value in self.search_groups.items(): + if self.encoding_type == 'onehot': + index = np.where(vector[start:start + + len(value[0].choices)] == 1)[0][0] + start += len(value) + else: + index = vector[start] + start += 1 + chosen = value[0].choices[int(index)] + model[key] = chosen + + return model + + @staticmethod + def get_correlation(prediction: np.array, + label: np.array) -> List[np.array]: + """Compute the correlations between prediction and ground-truth label. + + Args: + prediction (numpy.array): predict vector. + label (numpy.array): ground-truth label. + + Returns: + List[numpy.array]: coefficients of correlations between predicton + and ground-truth label. + """ + rmse = np.sqrt(((prediction - label)**2).mean()) + rho, _ = stats.spearmanr(prediction, label) + tau, _ = stats.kendalltau(prediction, label) + return [rmse, rho, tau] + + def preprocess(self, data: List[Dict[str, list]]) -> np.array: + """Preprocess the data, convert it into np.array format. + + Args: + data (List[Dict[str, list]]): input data for training. + + Returns: + numpy.array: input data in numpy.array format. + """ + if self.encoding_type == 'normal': + data = np.array([x['normal_vector'] for x in data]) + else: + data = np.array([x['onehot_vector'] for x in data]) + return data + + def fit(self, data: List[Dict[str, list]], label: np.array) -> None: + """Training the handler using the structure information of a model. The + weights of handler will be fixed after that. + + Args: + data (List[Dict[str, list]]): input data for training. + label (numpy.array): input label for training. + """ + data = self.preprocess(data) + self.handler.fit(data, label) + self.initialize = True + + def load_checkpoint(self) -> None: + """Load checkpoint for handler.""" + self.handler.load(self.handler_ckpt) + self.initialize = True + + def save_checkpoint(self, path: str) -> str: + """Save checkpoint of handler and return saved path for diff suffix. + + Args: + path (str): save path for the handler. + + Returns: + (str): specific checkpoint path of the current handler. + """ + return self.handler.save(path) diff --git a/requirements/optional.txt b/requirements/optional.txt index 32f7d6fd0..ffbd59d79 100644 --- a/requirements/optional.txt +++ b/requirements/optional.txt @@ -1,3 +1,5 @@ albumentations>=0.3.2 +pydacefit +pySOT==0.2.3 scipy # timm diff --git a/tests/test_models/test_task_modules/test_predictors/test_metric_predictor.py b/tests/test_models/test_task_modules/test_predictors/test_metric_predictor.py new file mode 100644 index 000000000..d9293cbf2 --- /dev/null +++ b/tests/test_models/test_task_modules/test_predictors/test_metric_predictor.py @@ -0,0 +1,196 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import tempfile +from unittest import TestCase + +import numpy as np +import torch.nn as nn +from mmengine.model import BaseModel + +from mmrazor.models import OneShotMutableOP +from mmrazor.registry import TASK_UTILS + +convs = nn.ModuleDict({ + 'conv1': nn.Conv2d(3, 8, 1), + 'conv2': nn.Conv2d(3, 8, 1), + 'conv3': nn.Conv2d(3, 8, 1), +}) +MutableOP = OneShotMutableOP(convs) + + +class ToyModel(BaseModel): + + def __init__(self, data_preprocessor=None): + super().__init__(data_preprocessor=data_preprocessor, init_cfg=None) + self.mutable = MutableOP + self.bn = nn.BatchNorm2d(8) + + def forward(self, batch_inputs, data_samples=None, mode='tensor'): + if mode == 'loss': + out = self.bn(self.mutable(batch_inputs)) + return dict(loss=out) + elif mode == 'predict': + out = self.bn(self.mutable(batch_inputs)) + 1 + return out + elif mode == 'tensor': + out = self.bn(self.mutable(batch_inputs)) + 2 + return out + + +class TestMetricPredictorWithGP(TestCase): + + def setUp(self) -> None: + self.temp_dir = tempfile.mkdtemp() + self.search_groups = {0: [MutableOP], 1: [MutableOP]} + self.candidates = [{0: 'conv1'}, {0: 'conv2'}, {0: 'conv3'}] + predictor_cfg = dict( + type='MetricPredictor', + handler_cfg=dict(type='GaussProcessHandler'), + search_groups=self.search_groups, + train_samples=4, + ) + self.predictor = TASK_UTILS.build(predictor_cfg) + self.model = ToyModel() + + def generate_data(self): + inputs = [] + for candidate in self.candidates: + inputs.append(self.predictor.model2vector(candidate)) + inputs = np.array(inputs) + labels = np.random.rand(3) + return inputs, labels + + def test_init_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.assertFalse(self.predictor.initialize) + self.predictor.fit(inputs, labels) + self.assertTrue(self.predictor.initialize) + + def test_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.predictor.fit(inputs, labels) + + metrics = self.predictor.predict(self.model) + self.assertIsInstance(metrics, dict) + self.assertGreater(metrics['accuracy_top-1'], 0.0) + + +class TestMetricPredictorWithCart(TestCase): + + def setUp(self) -> None: + self.temp_dir = tempfile.mkdtemp() + self.search_groups = {0: [MutableOP], 1: [MutableOP]} + self.candidates = [{0: 'conv1'}, {0: 'conv2'}, {0: 'conv3'}] + predictor_cfg = dict( + type='MetricPredictor', + handler_cfg=dict(type='CartsHandler'), + search_groups=self.search_groups, + train_samples=4, + ) + self.predictor = TASK_UTILS.build(predictor_cfg) + self.model = ToyModel() + + def generate_data(self): + inputs = [] + for candidate in self.candidates: + inputs.append(self.predictor.model2vector(candidate)) + inputs = np.array(inputs) + labels = np.random.rand(3) + return inputs, labels + + def test_init_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.assertFalse(self.predictor.initialize) + self.predictor.fit(inputs, labels) + self.assertTrue(self.predictor.initialize) + + def test_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.predictor.fit(inputs, labels) + + metrics = self.predictor.predict(self.model) + self.assertIsInstance(metrics, dict) + self.assertGreater(metrics['accuracy_top-1'], 0.0) + + +class TestMetricPredictorWithRBF(TestCase): + + def setUp(self) -> None: + self.temp_dir = tempfile.mkdtemp() + self.search_groups = {0: [MutableOP], 1: [MutableOP]} + self.candidates = [{0: 'conv1'}, {0: 'conv2'}, {0: 'conv3'}] + predictor_cfg = dict( + type='MetricPredictor', + handler_cfg=dict(type='RBFHandler'), + search_groups=self.search_groups, + train_samples=4, + ) + self.predictor = TASK_UTILS.build(predictor_cfg) + self.model = ToyModel() + + def generate_data(self): + inputs = [] + for candidate in self.candidates: + inputs.append(self.predictor.model2vector(candidate)) + inputs = np.array(inputs) + labels = np.random.rand(3) + return inputs, labels + + def test_init_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.assertFalse(self.predictor.initialize) + self.predictor.fit(inputs, labels) + self.assertTrue(self.predictor.initialize) + + def test_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.predictor.fit(inputs, labels) + + metrics = self.predictor.predict(self.model) + self.assertIsInstance(metrics, dict) + self.assertGreater(metrics['accuracy_top-1'], 0.0) + + +class TestMetricPredictorWithMLP(TestCase): + + def setUp(self) -> None: + self.temp_dir = tempfile.mkdtemp() + self.search_groups = {0: [MutableOP], 1: [MutableOP]} + self.candidates = [{0: 'conv1'}, {0: 'conv2'}, {0: 'conv3'}] + predictor_cfg = dict( + type='MetricPredictor', + handler_cfg=dict(type='MLPHandler'), + search_groups=self.search_groups, + train_samples=4, + ) + self.predictor = TASK_UTILS.build(predictor_cfg) + self.model = ToyModel() + + def generate_data(self): + inputs = [] + for candidate in self.candidates: + inputs.append(self.predictor.model2vector(candidate)) + inputs = np.array(inputs) + labels = np.random.rand(3) + return inputs, labels + + def test_init_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.assertFalse(self.predictor.initialize) + self.predictor.fit(inputs, labels) + self.assertTrue(self.predictor.initialize) + + def test_predictor(self): + self.model.mutable.current_choice = 'conv1' + inputs, labels = self.generate_data() + self.predictor.fit(inputs, labels) + + metrics = self.predictor.predict(self.model) + self.assertIsInstance(metrics, dict) + self.assertGreater(metrics['accuracy_top-1'], 0.0) diff --git a/tests/test_runners/test_evolution_search_loop.py b/tests/test_runners/test_evolution_search_loop.py index 6d8814a7b..a92731545 100644 --- a/tests/test_runners/test_evolution_search_loop.py +++ b/tests/test_runners/test_evolution_search_loop.py @@ -13,6 +13,7 @@ from torch.utils.data import DataLoader, Dataset from mmrazor.engine import EvolutionSearchLoop +from mmrazor.models import OneShotMutableOP from mmrazor.registry import LOOPS from mmrazor.structures import Candidates @@ -209,3 +210,168 @@ def test_run_loop(self, mock_flops, mock_export_fix_subnet): self.runner.rank = 0 loop.run() self.assertEqual(loop._max_epochs, 1) + + +class TestEvolutionSearchLoopWithPredictor(TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + convs = nn.ModuleDict({ + 'conv1': nn.Conv2d(3, 8, 1), + 'conv2': nn.Conv2d(3, 8, 1), + 'conv3': nn.Conv2d(3, 8, 1), + }) + MutableOP = OneShotMutableOP(convs) + self.search_groups = {0: [MutableOP], 1: [MutableOP]} + train_cfg = dict( + type='EvolutionSearchLoop', + max_epochs=4, + max_keep_ckpts=3, + resume_from=None, + num_candidates=4, + top_k=2, + num_mutation=2, + num_crossover=2, + mutate_prob=0.1, + constraints_range=dict(flops=(0, 330)), + score_key='bbox_mAP', + predictor_cfg=dict( + type='MetricPredictor', + handler_cfg=dict(type='GaussProcessHandler'), + search_groups=self.search_groups, + train_samples=4, + )) + self.train_cfg = Config(train_cfg) + self.runner = MagicMock(spec=ToyRunner) + self.dataloader = DataLoader(ToyDataset(), collate_fn=collate_fn) + self.evaluator = MagicMock() + + def tearDown(self): + shutil.rmtree(self.temp_dir) + + def test_init(self): + # test_init: dataloader and evaluator are instances + loop_cfg = copy.deepcopy(self.train_cfg) + loop_cfg.runner = self.runner + loop_cfg.dataloader = self.dataloader + loop_cfg.evaluator = self.evaluator + loop = LOOPS.build(loop_cfg) + self.assertIsInstance(loop, EvolutionSearchLoop) + + # test init_candidates is not None + fake_subnet = {'1': 'choice1', '2': 'choice2'} + fake_candidates = Candidates(fake_subnet) + init_candidates_path = os.path.join(self.temp_dir, 'candidates.yaml') + fileio.dump(fake_candidates, init_candidates_path) + loop_cfg.init_candidates = init_candidates_path + loop = LOOPS.build(loop_cfg) + self.assertIsInstance(loop, EvolutionSearchLoop) + self.assertEqual(loop.candidates, fake_candidates) + + @patch('mmrazor.engine.runner.utils.check.load_fix_subnet') + @patch('mmrazor.engine.runner.utils.check.export_fix_subnet') + @patch('mmrazor.models.task_modules.estimators.resource_estimator.' + 'get_model_flops_params') + def test_run_epoch(self, flops_params, mock_export_fix_subnet, + load_status): + # test_run_epoch: distributed == False + loop_cfg = copy.deepcopy(self.train_cfg) + loop_cfg.runner = self.runner + loop_cfg.dataloader = self.dataloader + loop_cfg.evaluator = self.evaluator + loop = LOOPS.build(loop_cfg) + self.runner.rank = 0 + self.runner.distributed = False + self.runner.work_dir = self.temp_dir + fake_subnet = {'1': 'choice1', '2': 'choice2'} + loop.model.sample_subnet = MagicMock(return_value=fake_subnet) + load_status.return_value = True + flops_params.return_value = 0, 0 + loop.run_epoch() + self.assertEqual(len(loop.candidates), 4) + self.assertEqual(len(loop.top_k_candidates), 2) + self.assertEqual(loop._epoch, 1) + + # test_run_epoch: distributed == True + loop = LOOPS.build(loop_cfg) + self.runner.rank = 0 + self.runner.distributed = True + self.runner.work_dir = self.temp_dir + fake_subnet = {'1': 'choice1', '2': 'choice2'} + self.runner.model.sample_subnet = MagicMock(return_value=fake_subnet) + loop.run_epoch() + self.assertEqual(len(loop.candidates), 4) + self.assertEqual(len(loop.top_k_candidates), 2) + self.assertEqual(loop._epoch, 1) + + # test_check_constraints + loop_cfg.constraints_range = dict(params=(0, 100)) + loop = LOOPS.build(loop_cfg) + self.runner.rank = 0 + self.runner.distributed = True + self.runner.work_dir = self.temp_dir + fake_subnet = {'1': 'choice1', '2': 'choice2'} + loop.model.sample_subnet = MagicMock(return_value=fake_subnet) + flops_params.return_value = (50., 1) + mock_export_fix_subnet.return_value = fake_subnet + loop.run_epoch() + self.assertEqual(len(loop.candidates), 4) + self.assertEqual(len(loop.top_k_candidates), 2) + self.assertEqual(loop._epoch, 1) + + @patch('mmrazor.engine.runner.utils.check.export_fix_subnet') + @patch('mmrazor.models.task_modules.predictor.metric_predictor.' + 'MetricPredictor.model2vector') + @patch('mmrazor.models.task_modules.estimators.resource_estimator.' + 'get_model_flops_params') + def test_run_loop(self, mock_flops, mock_model2vector, + mock_export_fix_subnet): + # test a new search: resume == None + loop_cfg = copy.deepcopy(self.train_cfg) + loop_cfg.runner = self.runner + loop_cfg.dataloader = self.dataloader + loop_cfg.evaluator = self.evaluator + loop = LOOPS.build(loop_cfg) + self.runner.rank = 0 + loop._epoch = 1 + + fake_subnet = {'1': 'choice1', '2': 'choice2'} + loop.model.sample_subnet = MagicMock(return_value=fake_subnet) + + self.runner.work_dir = self.temp_dir + loop.update_candidate_pool = MagicMock() + loop.val_candidate_pool = MagicMock() + + mutation_candidates = Candidates([fake_subnet] * loop.num_mutation) + for i in range(loop.num_mutation): + mutation_candidates.set_resource(i, 0.1 + 0.1 * i, 'flops') + mutation_candidates.set_resource(i, 99 + i, 'score') + crossover_candidates = Candidates([fake_subnet] * loop.num_crossover) + for i in range(loop.num_crossover): + crossover_candidates.set_resource(i, 0.1 + 0.1 * i, 'flops') + crossover_candidates.set_resource(i, 99 + i, 'score') + loop.gen_mutation_candidates = \ + MagicMock(return_value=mutation_candidates) + loop.gen_crossover_candidates = \ + MagicMock(return_value=crossover_candidates) + loop.candidates = Candidates([fake_subnet] * 4) + + mock_flops.return_value = (0.5, 101) + mock_export_fix_subnet.return_value = fake_subnet + mock_model2vector.return_value = dict( + normal_vector=[0, 1], onehot_vector=[0, 1, 0, 1]) + + loop.run() + assert os.path.exists( + os.path.join(self.temp_dir, 'best_fix_subnet.yaml')) + self.assertEqual(loop._epoch, loop._max_epochs) + assert os.path.exists( + os.path.join(self.temp_dir, + f'search_epoch_{loop._max_epochs-1}.pkl')) + # test resuming search + loop_cfg.resume_from = os.path.join( + self.temp_dir, f'search_epoch_{loop._max_epochs-1}.pkl') + loop = LOOPS.build(loop_cfg) + self.runner.rank = 0 + loop.run() + self.assertEqual(loop._max_epochs, 1)