API

Archai Model

class archai.discrete_search.api.archai_model.ArchaiModel(arch: Any, archid: str, metadata: Dict[str, Any] | None = None)

Model wrapper with an architecture identifier and an optional metadata dictionary.

clear() → None

Clear architecture from memory.

Sometimes, after evaluating an ArchaiModel, there is no need to keep its architecture instantiated, which optimizes memory usage.

Model Evaluator

class archai.discrete_search.api.model_evaluator.ModelEvaluator

Abstract class for synchronous model evaluators.

Evaluators are general-use classes used to evaluate architectures in given criteria (task performance, speed, size, etc.).

Subclasses of ModelEvaluator are expected to implement ModelEvaluator.evaluate

Synchronous evaluators are computed by search algorithms sequentially. For parallel / async. execution, please refer to archai.api.AsyncModelEvaluator.

For a list of built-in evaluators, please check archai.discrete_search.evaluators.

Examples

>>> class MyValTaskAccuracy(ModelEvaluator):
>>>     def __init__(self, dataset: DatasetProvider, batch_size: int = 32):
>>>         self.dataset = dataset
>>>         self.batch_size = batch_size
>>>
>>>     @overrides
>>>     def get_name(self) -> str:
>>>         return f'MyValTaskAccuracy_on_{self.dataset.get_name()}}'
>>>
>>>     @overrides
>>>     def evaluate(self, model: ArchaiModel, budget: Optional[float] = None):
>>>         _, val_data = self.dataset.get_train_val_datasets()
>>>         val_dl = torch.utils.data.DataLoader(val_data, batch_size=self.batch_size)
>>>
>>>         with torch.no_grad():
>>>             labels = np.concatenate([y for _, y in val_dl], axis=0)
>>>             preds = np.concatenate(
>>>                 [model.arch(x).cpu().numpy() for x, _ in val_dl],
>>>                 axis=0
>>>             )
>>>
>>>         return np.mean(labels == preds)
>>>
>>> class NumberOfModules(ModelEvaluator):
>>>     @overrides
>>>     def evaluate(self, model: ArchaiModel, budget: Optional[float] = None):
>>>         return len(list(model.arch.modules()))

abstract evaluate(arch: ArchaiModel, budget: float | None = None) → float

Evaluate an ArchaiModel instance, optionally using a budget value.

Parameters:

• arch – Model to be evaluated.

• dataset – A dataset provider object.

• budget – A budget multiplier value, used by search algorithms like SuccessiveHalving to specify how much compute should be spent in this evaluation. In order to use this type of search algorithm, the implementation of evaluate() must use the passed budget value accordingly.

Returns:

Evaluation result.

class archai.discrete_search.api.model_evaluator.AsyncModelEvaluator

Abstract class for asynchronous model evaluators.

Evaluators are general-use classes used to evaluate architectures in given criteria (task performance, speed, size, etc.).

Unlike archai.api.ModelEvaluator, AsyncModelEvaluator evaluates models in asynchronous fashion, by sending evaluation jobs to a queue and fetching the results later.

Subclasses of AsyncModelEvaluator are expected to implement AsyncModelEvaluator.send(arch: ArchaiModel, budget: Optional[float]) and AsyncModelEvaluator.fetch_all().

AsyncModelEvaluator.send is a non-blocking call that schedules an evaluation job for a given (model, budget) triplet. AsyncModelEvaluator.fetch_all is a blocking call that waits and gathers the results from current evaluation jobs and cleans the job queue.

For a list of built-in evaluators, please check archai.discrete_search.evaluators.

>>> my_obj = MyAsyncObj(dataset)  # My AsyncModelEvaluator subclass
>>>
>>> # Non blocking calls
>>> my_obj.send(model_1, budget=None)
>>> my_obj.send(model_2, budget=None)
>>> my_obj.send(model_3, budget=None)
>>>
>>> # Blocking call
>>> eval_results = my_obj.fetch_all()
>>> assert len(eval_results) == 3
>>>
>>> # Job queue is reset after `fetch_call` method
>>> my_obj.send(model_4, budget=None)
>>> assert len(my_obj.fetch_all()) == 1

abstract send(arch: ArchaiModel, budget: float | None = None) → None

Send an evaluation job for a given (model, budget) triplet.

Parameters:

• arch – Model to be evaluated.

• dataset – A dataset provider object.

• budget – A budget multiplier value, used by search algorithms like SuccessiveHalving to specify how much compute should be spent in this evaluation. In order to use this type of search algorithm, the implementation of send() must use the passed budget value accordingly.

abstract fetch_all() → List[float | None]

Fetch all evaluation results from the job queue.

Returns:

List of evaluation results. Each result is a float or None if evaluation job failed.

Predictor

class archai.discrete_search.api.predictor.MeanVar(mean: ndarray, var: ndarray)

Predictive mean and variance estimates from a surrogate model (Predictor).

mean: ndarray

Alias for field number 0

var: ndarray

Alias for field number 1

class archai.discrete_search.api.predictor.Predictor

Abstract class for a predictor model.

This class can be used to predict the performance of a model given its architecture. The class enforces implementation of two methods: fit and predict.

Note

This class is inherited from EnforceOverrides and any overridden methods in the subclass should be decorated with @overrides to ensure they are properly overridden.

abstract fit(encoded_archs: ndarray, y: ndarray) → None

Fit a predictor model using an array of encoded architectures (N, #features) and a multi-dimensional array of targets (N, #targets).

Parameters:

• encoded_archs – (N, #features) numpy array.

• y – (N, #targets) numpy array.

abstract predict(encoded_archs: ndarray) → MeanVar

Predict the performance of an array of architectures encoded by by a subclass implementation of BayesOptSearchSpaceBase.encode().

Parameters:

encoded_archs – Array of encoded architectures.

Returns:

Named tuple with mean (N, #targets) and var (N, #targets) arrays.

Search Objectives

class archai.discrete_search.api.search_objectives.SearchConstraint(name, evaluator, constraint)

class archai.discrete_search.api.search_objectives.SearchObjective(name, model_evaluator, higher_is_better, compute_intensive, constraint)

class archai.discrete_search.api.search_objectives.SearchObjectives(cache_objective_evaluation: bool | None = True)

Search objectives and constraints.

property objective_names: List[str]

Return a list of all objective names.

property cheap_objective_names: List[str]

Return a list of cheap objective names.

property expensive_objective_names: List[str]

Return a list of expensive objective names.

property objectives: Dict[str, SearchObjective]

Return a dictionary of all objectives.

property cheap_objectives: Dict[str, SearchObjective]

Return a dictionary of cheap objectives.

property expensive_objectives: Dict[str, SearchObjective]

Return a dictionary of expensive objectives.

property constraints: Dict[str, SearchConstraint]

Return a dictionary of all the constraints.

add_objective(name: str, model_evaluator: ModelEvaluator | AsyncModelEvaluator, higher_is_better: bool, compute_intensive: bool | None = True, constraint: Tuple[float, float] | None = None) → None

Add an objective function to the SearchObjectives object.

Parameters:

• name – The name of the objective.

• model_evaluator – The model evaluator responsible for evaluating the objective.

• higher_is_better – Whether the objective should be maximized (True) or minimized (False).

• compute_intensive – If True, the objective is considered computationally expensive and will be estimated using surrogate models when possible.

• constraint – Objective constraint used to filter out candidate architectures. Expects (lower_bound, upper_bound) tuple. Can only be set if compute_intensive is set to False.

add_constraint(name: str, model_evaluator: ModelEvaluator | AsyncModelEvaluator, constraint: Tuple[float, float]) → None

Add a search constraint to the SearchObjectives object.

Constraints are typically evaluated multiple times by search algorithms to validate candidate architectures and should not be computationally expensive to evaluate.

Parameters:

• name – The name of the constraint.

• model_evaluator – The model evaluator responsible for evaluating the constraint.

• constraint – The valid range of the constraint. Expects a (lower_bound, upper_bound) tuple.

validate_constraints(models: List[ArchaiModel], progress_bar: bool | None = False) → Tuple[Dict[str, ndarray], ndarray]

Evaluate constraints for a list of models and returns the indices of models that satisfy all constraints.

Parameters:

• models – List of models to evaluate.

• progress_bar – Whether to show progress bar.

Returns:

Evaluation results and indices of models that satisfy all constraints.

is_model_valid(model: ArchaiModel) → bool

Check if a model satisfies all constraints.

Parameters:

model – Model to check.

Returns:

True if model is valid, False otherwise.

eval_cheap_objs(models: List[ArchaiModel], budgets: Dict[str, List] | None = None, progress_bar: bool | None = False) → Dict[str, ndarray]

Evaluate all cheap objectives for a list of models.

Parameters:

• models – List of models to evaluate.

• budgets – Budgets for each objective.

• progress_bar – Whether to show progress bar.

Returns:

Dictionary with evaluation results.

eval_expensive_objs(models: List[ArchaiModel], budgets: Dict[str, List] | None = None, progress_bar: bool | None = False) → Dict[str, ndarray]

Evaluate all expensive objective functions for a list of models.

Parameters:

• models – List of models to evaluate.

• budgets – Budgets for each objective.

• progress_bar – Whether to show progress bar. Defaults to False.

Returns:

Dictionary with evaluation results.

eval_all_objs(models: List[ArchaiModel], budgets: Dict[str, List] | None = None, progress_bar: bool | None = False) → Dict[str, ndarray]

Evaluate all objective functions for a list of models.

Parameters:

• models – List of models to evaluate.

• budgets – Budgets for each objective.

• progress_bar – Whether to show progress bar.

Returns:

Dictionary with evaluation results.

save_cache(file_path: str) → None

Save the state of the SearchObjectives object to a YAML file.

Parameters:

file_path – Path to save YAML file.

load_cache(file_path: str) → None

Load the state of the SearchObjectives object from a YAML file.

Parameters:

file_path – Path to YAML file.

lookup_cache(obj_name: str, arch_id: str, budget: int | None) → float | None

Look up the cache for a specific objective, architecture and budget.

Parameters:

• obj_name – Name of objective.

• arch_id – Architecture ID.

• budget – Budget.

Returns:

Evaluation result if found in cache, None otherwise.

Search Results

class archai.discrete_search.api.search_results.SearchResults(search_space: DiscreteSearchSpace, objectives: SearchObjectives)

Discrete search results.

This class implements search results, which consists in producing data frames and plots with information regarding the search.

property all_evaluated_objs: Dict[str, array]

Return all evaluated objectives.

add_iteration_results(models: List[ArchaiModel], evaluation_results: Dict[str, ndarray], extra_model_data: Dict[str, List] | None = None) → None

Store results of the current search iteration.

Parameters:

• models – Models evaluated in the search iteration.

• evaluation_results – Evaluation results from SearchObjectives.eval_all_objs().

• extra_model_data – Additional model information to be stored in the search state file. Must be a list of the same size as models and csv-serializable.

get_pareto_frontier(start_iteration: int | None = 0, end_iteration: int | None = None) → Dict[str, Any]

Get the pareto-frontier using the search results from iterations start_iteration to end_iteration. If end_iteration=None, uses the last iteration.

Parameters:

• start_iteration – Start search iteration.

• end_iteration – End search iteration. If None, uses the last iteration.

Returns:

Dictionary containing ‘models’, ‘evaluation_results’, ‘indices’ and

’iteration_nums’ for all pareto-frontier members.

get_search_state_df() → DataFrame

Get the search state data frame.

Returns:

Search state data frame.

save_search_state(file_path: str | Path) → None

Save the search state to a .csv file.

Parameters:

file_path – File path to save the search state.

save_pareto_frontier_models(directory: str, save_weights: bool | None = False) → None

Save the pareto-frontier models to a directory.

Parameters:

• directory – Directory to save the models.

• save_weights – If True, saves the model weights. Otherwise, only saves the architecture.

plot_2d_pareto_evolution(objective_names: Tuple[str, str], figsize: Tuple[int, int] | None = (10, 5)) → Figure

Plot the evolution of the pareto-frontier in 2D.

Parameters:

• objective_names – Names of the objectives to plot.

• figsize – Figure size.

Returns:

2D pareto-frontier evolution figure.

save_2d_pareto_evolution_plot(objective_names: Tuple[str, str], file_path: str) → None

Save the evolution of the pareto-frontier in 2D.

Parameters:

• objective_names – Names of the objectives to plot.

• file_path – Path to save the plot.

save_all_2d_pareto_evolution_plots(directory: str | Path) → None

Save all the 2D pareto-frontier evolution plots.

Parameters:

directory – Directory to save the plots.

Search Space

class archai.discrete_search.api.search_space.DiscreteSearchSpace

Abstract class for discrete search spaces.

This class serves as a base for implementing search spaces. The class enforces implementation of five methods: save_arch, load_arch, save_model_weights, load_model_weights and random_sample.

Note

This class is inherited from EnforceOverrides and any overridden methods in the subclass should be decorated with @overrides to ensure they are properly overridden.

Examples

>>> class MyDiscreteSearchSpace(DiscreteSearchSpace):
>>>     def __init__(self) -> None:
>>>         super().__init__()
>>>
>>>     @overrides
>>>     def save_arch(self, arch, file_path) -> None:
>>>         torch.save(arch, file_path)
>>>
>>>     @overrides
>>>     def load_arch(self, file_path) -> ArchaiModel:
>>>         return torch.load(file_path)
>>>
>>>     @overrides
>>>     def save_model_weights(self, model, file_path) -> None:
>>>         torch.save(model.state_dict(), file_path)
>>>
>>>     @overrides
>>>     def load_model_weights(self, model, file_path) -> None:
>>>         model.load_state_dict(torch.load(file_path))
>>>
>>>     @overrides
>>>     def random_sample(self, config) -> ArchaiModel:
>>>         return ArchaiModel(config)

abstract save_arch(model: ArchaiModel, file_path: str) → None

Save an architecture to a file without saving the weights.

Parameters:

• model – Model’s architecture to save.

• file_path – File path to save the architecture.

abstract load_arch(file_path: str) → ArchaiModel

Load from a file an architecture that was saved using SearchSpace.save_arch().

Parameters:

file_path – File path to load the architecture.

Returns:

Loaded model.

abstract save_model_weights(model: ArchaiModel, file_path: str) → None

Save the weights of a model.

Parameters:

• model – Model to save the weights.

• file_path – File path to save the weights.

abstract load_model_weights(model: ArchaiModel, file_path: str) → None

Load the weights (created with SearchSpace.save_model_weights()) into a model of the same architecture.

Parameters:

• model – Model to load the weights.

• file_path – File path to load the weights.

abstract random_sample() → ArchaiModel

Randomly sample an architecture from the search spaces.

Returns:

Sampled architecture.

class archai.discrete_search.api.search_space.EvolutionarySearchSpace

Abstract class for discrete search spaces compatible with evolutionary algorithms.

The class enforces implementation of two methods: mutate and crossover.

Note

This class is inherited from EnforceOverrides and any overridden methods in the subclass should be decorated with @overrides to ensure they are properly overridden.

abstract mutate(arch: ArchaiModel) → ArchaiModel

Mutate an architecture from the search space.

This method should not alter the base model architecture directly, only generate a new one.

Parameters:

arch – Base model.

Returns:

Mutated model.

abstract crossover(arch_list: List[ArchaiModel]) → ArchaiModel

Combine a list of architectures into a new one.

Parameters:

arch_list – List of architectures.

Returns:

Resulting model.

class archai.discrete_search.api.search_space.BayesOptSearchSpace

Abstract class for discrete search spaces compatible with Bayesian Optimization algorithms.

The class enforces implementation of a single method: encode.

Note

This class is inherited from EnforceOverrides and any overridden methods in the subclass should be decorated with @overrides to ensure they are properly overridden.

abstract encode(arch: ArchaiModel) → ndarray

Encode an architecture into a fixed-length vector representation.

Parameters:

arch – Model from the search space.

Returns:

Fixed-length vector representation of arch.

Searcher

class archai.discrete_search.api.searcher.Searcher

Abstract class for searchers.

This class serves as a base for implementing searchers, which searches for an architecture given an algorithm. The class enforces implementation of a single method: search.

Note

This class is inherited from EnforceOverrides and any overridden methods in the subclass should be decorated with @overrides to ensure they are properly overridden.

Examples

>>> class MySearcher(Searcher):
>>>     def __init__(self) -> None:
>>>         super().__init__()
>>>
>>>     @overrides
>>>     def search(self) -> SearchResults:
>>>         # Code used to search for the best architecture
>>>         return SearchResults(...)

abstract search() → SearchResults

Search for the best architecture.

Returns:

Search results.

subscribe_start_iteration(fn: Callable[[int], None])

on_start_iteration(iteration: int)