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..1c0e05d7b --- /dev/null +++ b/mmrazor/models/task_modules/predictor/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .metric_predictor import MetricPredictor +from .zero_shot_predictor import ZeroShotPredictor + +__all__ = ['MetricPredictor', 'ZeroShotPredictor'] diff --git a/mmrazor/models/task_modules/predictor/base_predictor.py b/mmrazor/models/task_modules/predictor/base_predictor.py new file mode 100644 index 000000000..f482e4bb1 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/base_predictor.py @@ -0,0 +1,18 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from abc import abstractmethod + +from mmrazor.registry import TASK_UTILS + + +class BasePredictor(): + """Base predictor.""" + + def __init__(self, handler_cfg: dict): + """init.""" + self.handler_cfg = handler_cfg + self.handler = TASK_UTILS.build(handler_cfg) + + @abstractmethod + def predict(self, model, predict_args): + """predict result.""" + pass 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..bf4cfd41c --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/base_handler.py @@ -0,0 +1,25 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from joblib import dump, load + + +class BaseHandler(): + """Base handler.""" + + def __init__(self) -> None: + pass + + def fit(self, train_data, train_label): + pass + + def predict(self, test_data): + 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..f5608a34c --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/carts_handler.py @@ -0,0 +1,69 @@ +# Copyright (c) OpenMMLab. All rights reserved. +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: + n_tree (int): number of regression trees. + """ + + def __init__(self, n_tree=1000): + self.n_tree = n_tree + self.model = None + + @staticmethod + def _make_decision_trees(train_data, train_label, n_tree): + """Construct the decision trees.""" + feature_record = [] + tree_record = [] + + for i in range(n_tree): + 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 + + def fit(self, train_data, train_label): + """Training predictor.""" + self.model = self._make_decision_trees(train_data, train_label, + self.n_tree) + + def predict(self, test_data): + """Predict the subnets' performance.""" + trees, features = self.model[0], self.model[1] + test_num, n_tree = 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(n_tree) + + 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 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..1c932a81f --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/gp_handler.py @@ -0,0 +1,112 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import numpy as np + +try: + import pydacefit + from pydacefit.dace import DACE +except ImportError: + pydacefit = None + DACE = object + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + + +def get_func(): + if pydacefit is None: + raise RuntimeError('Failed to import pydacefit. Please run ' + '"pip install pydacefit". ') + + from pydacefit.corr import (corr_cubic, corr_exp, corr_expg, corr_gauss, + corr_spherical, corr_spline) + from pydacefit.dace import regr_linear, regr_quadratic + from pydacefit.regr import regr_constant + + 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 + } + + return REGR, CORR + + +class DACE_with_smooth(DACE): + """GP model.""" + + def __init__(self, regr, corr, theta=1, thetaL=0, thetaU=100): + super(DACE_with_smooth, self).__init__(regr, corr, theta, thetaL, + thetaU) + + def fit(self, X, Y): + + 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: + from pydacefit.fit import fit + self.model = 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): + """Gaussian Process (Kriging) + + Args: + regr (str): regression kernel for GP model. + corr (str): correlation kernel for GP model. + """ + + def __init__(self, regr='linear', corr='gauss'): + REGR, CORR = get_func() + assert regr in REGR and corr in CORR, \ + NotImplementedError('Unknown GP regression or correlation !') + 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, train_label): + """Training predictor.""" + self.model.fit(train_data, train_label) + + def predict(self, test_data): + """Predict the subnets' performance.""" + 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..9d08b98bd --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/mlp_handler.py @@ -0,0 +1,8 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .base_handler import BaseHandler + + +class MLPHandler(BaseHandler): + + def __init__(self) -> None: + super().__init__() 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..78733d017 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/handler/rbf_handler.py @@ -0,0 +1,61 @@ +# Copyright (c) OpenMMLab. All rights reserved. +try: + import pySOT +except ImportError: + pySOT = None + +from mmrazor.registry import TASK_UTILS +from .base_handler import BaseHandler + + +@TASK_UTILS.register_module() +class RBFHandler(BaseHandler): + """Radial Basis Function. + + Args: + kernel (str): RBF kernel object. + tail (str): RBF polynomial tail object. + """ + + def __init__(self, kernel='tps', tail='linear'): + if pySOT is None: + raise RuntimeError('Failed to import pydacefit. Please run ' + '"pip install pySOT==0.2.3". ') + from pySOT.surrogate import (ConstantTail, CubicKernel, LinearTail, + TPSKernel) + self.kernel = kernel + self.tail = tail + self.model = None + + if kernel == 'cubic': + self.kernel = CubicKernel + elif self.kernel == 'tps': + self.kernel = TPSKernel + else: + raise NotImplementedError('unknown RBF kernel') + + if tail == 'linear': + self.tail = LinearTail + elif self.tail == 'constant': + self.tail = ConstantTail + else: + raise NotImplementedError('unknown RBF tail') + + def fit(self, train_data, train_label): + """Training predictor.""" + if train_data.shape[0] <= train_data.shape[1]: + raise ValueError('RBF only support ' + f'# of samples{train_data.shape[0]}' + f' > # of dimensions{train_data.shape[1]} !') + from pySOT.surrogate import RBFInterpolant + 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): + """Predict the subnets' performance.""" + 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..05869bc37 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/metric_predictor.py @@ -0,0 +1,163 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Dict + +import numpy as np +import scipy.stats as stats + +from mmrazor.registry import TASK_UTILS +from mmrazor.structures import export_fix_subnet +from .base_predictor import BasePredictor +from .handler import MLPHandler, RBFHandler + + +@TASK_UTILS.register_module() +class MetricPredictor(BasePredictor): + """Metric predictor. + + Args: + handler_type (str): The chosen type for different handler. + search_groups(dict) : search_groups of the specified super net. + pool_size(int) : Total nums of training samples if necessary. + handler_cfg (dict, optional): Optional parameter for predictor. + fit_cfg(dict, optional): Training parameters. Only supported for + MLP predictor. + pretrained(str, optional): Path to predictor's weights. If given, + predictor will load the specified weights and skip + training phase. + evaluation(str, optional): If not None, compute the correlations + between prediction and true label, used for evaluate the + predictor's performance. 'simple' only evaluate the final + samples, 'complex' whill evaluate samples in the + candidate_pool. Defaults to None. + retrain_samples(dict, optional): Given the training samples for + predictor. + evaluate_samples(dict, optional): Given the predicting samples for + predictor. + encoding_type(str, optional): how to encode the search space to + integer bit-string. Default is onehot. + """ + + def __init__(self, + handler_cfg: dict, + search_groups, + score_key, + train_samples: int = 2, + fit_cfg: dict = None, + pretrained=None, + evaluation=None, + evaluate_samples=None, + encoding_type='onehot', + **kwargs): + super().__init__(handler_cfg=handler_cfg) + + if isinstance(self.handler, RBFHandler): + encoding_type = 'normal' + assert encoding_type in [ + 'normal', 'onehot' + ], ('encoding_type must be `normal` or `onehot`.' + f'Got `{encoding_type}`.') + self.encoding_type = encoding_type + + self.search_groups = search_groups + self.train_samples = train_samples + self.pretrained = pretrained + self.evaluation = evaluation + assert evaluation in [None, 'simple', 'complex' + ], (f'Not support evaluation mode {evaluation}.') + + self.fit_cfg = fit_cfg + self.evaluate_samples = evaluate_samples + self.score_key_list = [score_key] + ['anticipate'] + self.initialize = False + + def predict(self, model, predict_args=dict()): + """Predict the candidate's performance.""" + metric = {} + if not self.initialize or predict_args.get('anticipate', False): + raise AssertionError( + 'Call evaluator to get metric instead of predictor.') + + if self.initialize: + candidate = export_fix_subnet(model) + input = self.preprocess(np.array([self.spec2feats(candidate)])) + score = float(np.squeeze(self.handler.predict(input))) + 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 spec2feats(self, candidate: dict) -> dict: + """Convert the candidate to N dimensions vector. + + N is different for different supernet. + """ + index = 0 + feats_dict: Dict[str, list] = dict(feats=[], onehot_feats=[]) + + for choice in candidate.values(): + assert len(self.search_groups[index]) == 1 + _candidates = self.search_groups[index][0].choices + onehot = np.zeros(len(_candidates)) + _chosen_index = _candidates.index(choice) + onehot[_chosen_index] = 1 + + feats_dict['feats'].extend([_chosen_index]) + feats_dict['onehot_feats'].extend(onehot) + index += 1 + + return feats_dict + + def feats2spec(self, feats, type='onehot'): + """Convert the N dimensions vector to original candidates. + + feats is the output comes form self.spec2feats. + """ + fix_subnet = {} + start = 0 + for key, value in self.search_groups.items(): + if type == 'onehot': + index = np.where(feats[start:start + + len(value[0].choices)] == 1)[0][0] + start += len(value) + else: + index = feats[start] + start += 1 + chosen = value[0].choices[int(index)] + fix_subnet[key] = chosen + return fix_subnet + + def load_checkpoint(self): + self.handler.load(self.pretrained) + self.initialize = True + + def save(self, path): + """Save predictor and return saved path for diff suffix;""" + return self.handler.save(path) + + def get_correlation(self, prediction, target): + """Compute the correlations between prediction and true label.""" + rmse = np.sqrt(((prediction - target)**2).mean()) + rho, _ = stats.spearmanr(prediction, target) + tau, _ = stats.kendalltau(prediction, target) + + return rmse, rho, tau + + def preprocess(self, input): + if isinstance(input[0], dict): + if self.encoding_type == 'normal': + input = np.array([x['feats'] for x in input]) + else: + input = np.array([x['onehot_feats'] for x in input]) + return input + + def fit(self, input, target): + """Training the accuracy predictor.""" + input = self.preprocess(input) + if isinstance(self.handler, MLPHandler): + self.handler.check_dimentions(input.shape[1]) + self.handler.fit(x=input, y=target, **self.fit_cfg) + else: + self.handler.fit(input, target) + + self.initialize = True diff --git a/mmrazor/models/task_modules/predictor/zero_shot_predictor.py b/mmrazor/models/task_modules/predictor/zero_shot_predictor.py new file mode 100644 index 000000000..54a310d89 --- /dev/null +++ b/mmrazor/models/task_modules/predictor/zero_shot_predictor.py @@ -0,0 +1,10 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from mmrazor.registry import TASK_UTILS +from .base_predictor import BasePredictor + + +@TASK_UTILS.register_module() +class ZeroShotPredictor(BasePredictor): + + def __init__(self): + super().__init__()