Evaluators#

Functional#

class archai.discrete_search.evaluators.functional.EvaluationFunction(evaluation_fn: Callable)[source]#

Custom function evaluator.

This evaluator is used to wrap a custom evaluation function.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.

ONNX Model#

class archai.discrete_search.evaluators.onnx_model.AvgOnnxLatency(input_shape: Tuple[int, ...] | List[Tuple[int, ...]], num_trials: int | None = 1, input_dtype: str | None = 'torch.FloatTensor', rand_range: Tuple[float, float] | None = (0.0, 1.0), export_kwargs: Dict[str, Any] | None = None, device: str | None = 'cpu', inf_session_kwargs: Dict[str, Any] | None = None)[source]#

Evaluate the average ONNX Latency (in seconds) of an architecture.

The latency is measured by running the model on random inputs and averaging the latency over num_trials trials.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.

Progressive Training#

class archai.discrete_search.evaluators.progressive_training.ProgressiveTraining(search_space: DiscreteSearchSpace, dataset: DatasetProvider, training_fn: Callable)[source]#

Progressive training evaluator.

evaluate(arch: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.progressive_training.RayProgressiveTraining(search_space: DiscreteSearchSpace, dataset: DatasetProvider, training_fn: Callable, timeout: float | None = None, force_stop: bool | None = False, **ray_kwargs)[source]#

Progressive training evaluator using Ray.

send(arch: ArchaiModel, budget: float | None = None) → None[source]#

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.

fetch_all() → List[float | None][source]#

Fetch all evaluation results from the job queue.

Returns:

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

PyTorch Profiler#

class archai.discrete_search.evaluators.pt_profiler.TorchNumParameters(exclude_cls: List[Module] | None = None, trainable_only: bool | None = True)[source]#

Total number of parameters.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.pt_profiler.TorchFlops(forward_args: Tensor | List[Tensor] | None = None, forward_kwargs: Dict[str, Tensor] | None = None, ignore_layers: List[str] | None = None)[source]#

Total number of FLOPs.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.pt_profiler.TorchMacs(forward_args: Tensor | List[Tensor] | None = None, forward_kwargs: Dict[str, Tensor] | None = None, ignore_layers: List[str] | None = None)[source]#

Total number of MACs.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.pt_profiler.TorchLatency(forward_args: Tensor | List[Tensor] | None = None, forward_kwargs: Dict[str, Tensor] | None = None, num_warmups: int | None = 1, num_samples: int | None = 1, use_cuda: bool | None = False, use_median: bool | None = False, ignore_layers: List[str] | None = None)[source]#

Average/median latency (in seconds) of a PyTorch model using a sample input.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.pt_profiler.TorchPeakCudaMemory(forward_args: Tensor | List[Tensor] | None = None, forward_kwargs: Dict[str, Tensor] | None = None, num_warmups: int | None = 1, num_samples: int | None = 1, use_median: bool | None = False, ignore_layers: List[str] | None = None)[source]#

Measures CUDA peak memory (in bytes) of a PyTorch model using a sample input.

evaluate(model: ArchaiModel, budget: float | None = None) → float[source]#

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.evaluators.pt_profiler.TorchPeakCpuMemory(forward_args: Tensor | List[Tensor] | None = None, forward_kwargs: Dict[str, Tensor] | None = None)[source]#

Measures CPU peak memory (in bytes) of a PyTorch model using a sample input.

evaluate(model: ArchaiModel, budget: float | None = None)[source]#

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.

Ray#

class archai.discrete_search.evaluators.ray.RayParallelEvaluator(obj: ModelEvaluator, timeout: float | None = None, force_stop: bool | None = False, **ray_kwargs)[source]#

Wraps a ModelEvaluator object into an AsyncModelEvaluator with parallel execution using Ray.

RayParallelEvaluator expects a stateless objective function as input, meaning that any ModelEvaluator.evaluate(arch, …) will not alter the state of obj or arch in any way.

send(arch: ArchaiModel, budget: float | None = None) → None[source]#

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.

fetch_all() → List[float | None][source]#

Fetch all evaluation results from the job queue.

Returns:

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

Remote Azure Benchmark#

class archai.discrete_search.evaluators.remote_azure_benchmark.RemoteAzureBenchmarkEvaluator(input_shape: Tuple | List[Tuple], store: ArchaiStore, experiment_name: str, metric_key: str, overwrite: bool | None = True, max_retries: int | None = 5, retry_interval: int | None = 120, onnx_export_kwargs: Dict[str, Any] | None = None, verbose: bool = False, benchmark_only: bool = True)[source]#

Simple adapter for benchmarking architectures asynchronously on Azure.

This adapter uploads an ONNX model to a Azure Blob storage container and records the model entry on the respective Azure Table.

send(arch: ArchaiModel, budget: float | None = None) → None[source]#

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.

fetch_all() → List[float | None][source]#

Fetch all evaluation results from the job queue.

Returns:

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